阶段三：CIPA 双流多模态模型 C++ TensorRT 边缘部署总结

本阶段的核心目标是将基于 PyTorch 训练的算法模型，转化为脱离 Python 环境、纯 C++ 驱动且针对底层 GPU（如 RTX 3090 / Jetson Nano）极致优化的工业级推理引擎。

整个部署流程分为三大步骤：静态图转换 (PyTorch -> ONNX)、硬件级引擎锻造 (ONNX -> TensorRT)、C++ 视觉推理流水线构建。

步骤一：静态图转换与动态解包 (PyTorch -> ONNX)

具体工作：编写export_onnx.py脚本，脱离庞大的训练框架，通过伪造配置文件（MockConfig）和构建双流虚拟输入（Dummy Inputs），将.pth动态图导出为跨平台的静态网络图.onnx。

解决的核心痛点（史诗级排错）：

1. 单双通道维度错位陷阱：训练时由于使用了BCEWithLogitsLoss，模型实际上是单通道输出（num_classes=1），但导出配置默认写成了 2。这导致 PyTorch 在加载权重时悄悄丢弃了最后一层的预训练权重，导致模型生成充满“棋盘网格伪影”的乱码图。解决：强制锁定num_classes=1并开启strict=True。

2. 多模态.pth字典解包失败：保存的权重文件不仅包含网络参数，还打包了 optimizer、epoch 等训练状态，且不包含标准的'state_dict'键值。解决：编写智能解析代码，自动遍历字典提取真实权重，并剥离了潜在的多卡训练module.前缀。

3. Axial Attention 相对位置编码崩塌：导出时默认使用了 $512 \times 512$ 的分辨率，但模型实际上是在 $640 \times 640$ 尺寸下训练的，导致底层轴向注意力机制的尺寸不匹配（[64, 319] vs [64, 255]）。解决：将输入张量全面拉升至 $640 \times 640$ 模式。

关键代码 (export_onnx.py核心片段)：

Python

class MockConfig: def __init__(self): self.image_height = 640 # 顺应真实训练分辨率 self.image_width = 640 # 顺应真实训练分辨率 self.patch_size = 4 self.decoder = 'MambaDecoder' self.num_classes = 1 # 严格遵照 BCEWithLogitsLoss 的 1 通道 def export_cipa_onnx(): # 1. 智能解析权重文件 checkpoint = torch.load("save_model/2026-01-03_20_24_20_Exp5_SOTA/best_model.pth", map_location="cuda") raw_weights = checkpoint['model'] if 'model' in checkpoint else checkpoint clean_weights = {} for k, v in raw_weights.items(): new_key = k[7:] if k.startswith('module.') else k clean_weights[new_key] = v model.load_state_dict(clean_weights, strict=True) # 严禁随机初始化 model.eval() # 2. 构造 640 分辨率双流输入 dummy_ct = torch.randn(1, 3, 640, 640, device="cuda") dummy_pet = torch.randn(1, 3, 640, 640, device="cuda") # 3. 导出 torch.onnx.export( model, (dummy_ct, dummy_pet), "cipa_mamba_dual_stream.onnx", export_params=True, opset_version=13, do_constant_folding=True, input_names=['input_ct', 'input_pet'], output_names=['output_last', 'output_0', 'output_1', 'output_2'] )

执行指令：python export_onnx.py

步骤二：硬件级引擎锻造 (ONNX -> TensorRT Engine)

具体工作：使用 NVIDIA 官方的trtexec工具，对生成的 ONNX 模型进行算子融合（Layer Fusion）、内核自动调优（Kernel Auto-Tuning），生成直接与显卡物理架构绑定的二进制 Engine 文件。

解决的核心痛点：

• Mamba 架构的 FP16 精度溢出：最初尝试使用--fp16混合精度编译时，由于 Mamba 底层状态空间方程和 LayerNorm 在 16 位半精度下发生数值溢出（Overflow），导致推理结果大面积崩坏（输出极端的网格噪点）。解决：放弃 FP16 的激进量化，采用原生的 FP32 精度，在保证医疗图像极高分割精度的前提下，依然跑出了优异的吞吐量。

关键指令：

Bash

./TensorRT-8.6.1.6/bin/trtexec --onnx=cipa_mamba_dual_stream.onnx --saveEngine=cipa_mamba_fp32.engine --workspace=4096

步骤三：C++ 视觉推理流水线构建 (C++ Infer)

具体工作：从零开始编写main.cpp，利用 CUDA Runtime API 和 TensorRT C++ API，手动接管了 GPU 显存分配、异步数据拷贝（H2D/D2H）、以及多分支前向推理计算。同时引入 OpenCV 处理图像的输入与输出。

解决的核心痛点：

1. 跨语言像素级对齐（对齐 Python 的 Dataset）：C++ 和 OpenCV 默认读取图像的维度和逻辑与 Python 存在巨大鸿沟。解决：在 C++ 中强制读取单通道灰度图，手动展开为CHW格式的三通道张量，并完美复刻了 Python 端独特的归一化公式(x / 255.0) * 3.2 - 1.6。

2. 动态输入端口寻址：避免了因为硬编码索引导致 CT 和 PET 图像“张冠李戴”送错模态分支的问题。解决：采用getBindingIndex依据算子名称动态获取物理内存地址。

3. 极速单通道后处理：摒弃了复杂的 Softmax 操作，利用数学等价性将 Python 的Sigmoid(x) >= 0.5直接简化为底层的logit >= 0.0f。

关键代码 (main.cpp核心预处理与推理片段)：

C++

// 1. 跨语言对齐：完美复刻 Python 预处理逻辑 void preprocessImage(const std::string& image_path, std::vector<float>& buffer) { cv::Mat img = cv::imread(image_path, cv::IMREAD_GRAYSCALE); cv::resize(img, img, cv::Size(640, 640)); // 严格匹配 640 架构 int index = 0; for (int c = 0; c < 3; ++c) { // 模拟 Python 的 repeat(1, 3, 1, 1) for (int i = 0; i < 640; ++i) { for (int j = 0; j < 640; ++j) { float val = img.at<uchar>(i, j) / 255.0f; buffer[index++] = val * 3.2f - 1.6f; // 复刻独特归一化公式 } } } } int main() { // ... (加载引擎、创建 ExecutionContext 略) ... // 2. 动态获取物理显存通道，防止模态错乱 int ct_idx = engine->getBindingIndex("input_ct"); int pet_idx = engine->getBindingIndex("input_pet"); int out_idx = engine->getBindingIndex("output_last"); // 3. 异步数据流搬运 (Host to Device) cudaStream_t stream; cudaStreamCreate(&stream); cudaMemcpyAsync(gpu_buffers[ct_idx], cpu_input_ct.data(), size, cudaMemcpyHostToDevice, stream); cudaMemcpyAsync(gpu_buffers[pet_idx], cpu_input_pet.data(), size, cudaMemcpyHostToDevice, stream); // 4. 核心推理点火 context->enqueueV2(gpu_buffers.data(), stream, nullptr); // 5. 结果取回 (Device to Host) 与 O(1) 极速后处理 cudaMemcpyAsync(cpu_output_last.data(), gpu_buffers[out_idx], size, cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream); cv::Mat mask(640, 640, CV_8UC1); int area = 640 * 640; for (int i = 0; i < area; ++i) { // Sigmoid(x) >= 0.5 数学等价于 x >= 0.0f mask.data[i] = (cpu_output_last[i] >= 0.0f) ? 255 : 0; } cv::imwrite("final_segmentation_mask.png", mask); }

编译与执行指令：```bash

cd cpp_infer/build

cmake .. && make

./cipa_infer

运行图片截取

python export_onnx.py

行trtexec结束时，终端打印的=== Performance summary ===那一段（包含Throughput: 41.7412 qps、Latency、GPU Compute Time等数据）。

运行./cipa_infer时，终端打印的✅ 完美对齐图像、🚀 数据装载完毕、⚡ 启动张量计算、🎉 绝杀！完美分割图已生成这一完整流程。

成果展示

测试照片：

CT:

PET:

MASK:

测试结果生成MASK: