BisenetV1/2網絡以及模型推理轉換

1 BiSenetV1

相關開源地址，eg1，

??為了加速模型推理，很多模型裁剪圖像尺寸、剪枝卷積網絡通道、卷積步驟數等，這樣損失了空間信息以及相應的感受野，BiSenet這個網絡選擇融合SP和CP兩個通道來實現加速

class BiSeNetV1(nn.Module):def __init__(self, n_classes, aux_mode='train', *args, **kwargs):super(BiSeNetV1, self).__init__()self.cp = ContextPath()self.sp = SpatialPath()self.ffm = FeatureFusionModule(256, 256)self.conv_out = BiSeNetOutput(256, 256, n_classes, up_factor=8)self.aux_mode = aux_modeif self.aux_mode == 'train':self.conv_out16 = BiSeNetOutput(128, 64, n_classes, up_factor=8)self.conv_out32 = BiSeNetOutput(128, 64, n_classes, up_factor=16)self.init_weight()def forward(self, x):H, W = x.size()[2:]feat_cp8, feat_cp16 = self.cp(x)feat_sp = self.sp(x)feat_fuse = self.ffm(feat_sp, feat_cp8)feat_out = self.conv_out(feat_fuse)if self.aux_mode == 'train':feat_out16 = self.conv_out16(feat_cp8)feat_out32 = self.conv_out32(feat_cp16)return feat_out, feat_out16, feat_out32elif self.aux_mode == 'eval':return feat_out,elif self.aux_mode == 'pred':#  feat_out = feat_out.argmax(dim=1)feat_out = CustomArgMax.apply(feat_out, 1)return feat_outelse:raise NotImplementedErrordef init_weight(self):for ly in self.children():if isinstance(ly, nn.Conv2d):nn.init.kaiming_normal_(ly.weight, a=1)if not ly.bias is None: nn.init.constant_(ly.bias, 0)def get_params(self):wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []for name, child in self.named_children():child_wd_params, child_nowd_params = child.get_params()if isinstance(child, (FeatureFusionModule, BiSeNetOutput)):lr_mul_wd_params += child_wd_paramslr_mul_nowd_params += child_nowd_paramselse:wd_params += child_wd_paramsnowd_params += child_nowd_paramsreturn wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params

1.1 Contex Path

??上下文通道是的backbone基于Resnet18去實現的

# 上下文path
class ContextPath(nn.Module):def __init__(self, *args, **kwargs):super(ContextPath, self).__init__()self.resnet = Resnet18()        # 調用殘差網絡  resnet.pyself.arm16 = AttentionRefinementModule(256, 128)self.arm32 = AttentionRefinementModule(512, 128)self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)self.up32 = nn.Upsample(scale_factor=2.)self.up16 = nn.Upsample(scale_factor=2.)self.init_weight()def forward(self, x):# 下采樣的特征是通過殘差網絡Resnet18得到的feat8, feat16, feat32 = self.resnet(x)avg = torch.mean(feat32, dim=(2, 3), keepdim=True)avg = self.conv_avg(avg)feat32_arm = self.arm32(feat32)feat32_sum = feat32_arm + avgfeat32_up = self.up32(feat32_sum)   # 原來下采樣，現在上采樣恢復feat32_up = self.conv_head32(feat32_up)feat16_arm = self.arm16(feat16)feat16_sum = feat16_arm + feat32_up         # 兩次采樣相加feat16_up = self.up16(feat16_sum)feat16_up = self.conv_head16(feat16_up)return feat16_up, feat32_up # x8, x16    feat16_up就是最終的輸出

1.2 Spatial Path

空間通道比較簡單

class SpatialPath(nn.Module):def __init__(self, *args, **kwargs):# super() 函數用于調用父類nn.Module的方法，這里調用了調用了 nn.Module 的構造函數# super() 函數的第一個參數是當前類（SpatialPath）# super() 第二個參數是當前類的實例（self）super(SpatialPath, self).__init__()     self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)           # 卷積核大小為7x7，步長為2，填充為3self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)self.init_weight()      # 初始化權重def forward(self, x):       # x 是輸入數據，通常是一個四維張量，形狀為 (batch_size, channels, height, width)# 定義前向傳播過程----卷積計算feat = self.conv1(x)feat = self.conv2(feat)feat = self.conv3(feat)feat = self.conv_out(feat)return feat

1.3 ARM

1.4 FFM

1.5 backbone

這里使用的是resnet網絡，可視化殘差網絡第二層

2 模型推理代碼流程分析

2.1 加載模型

net = model_factory[cfg.model_type](cfg.n_cats, aux_mode='eval')
net.load_state_dict(torch.load(args.weight_path, map_location='cpu'), strict=False)
net.eval()     # 設置為推理模式
net.cuda()     # 將模型加載到GPU上

2.2 模型推理

① 轉換張量

對于一張圖像

# prepare data
to_tensor = T.ToTensor(mean=(0.3257, 0.3690, 0.3223), # city, rgbstd=(0.2112, 0.2148, 0.2115),
)
im = cv2.imread(args.img_path)[:, :, ::-1]   # [:, :, ::-1] 將 BGR 格式轉換為 RGB（因 OpenCV 默認讀取為 BGR，而PyTorch需要 RGB）
im = to_tensor(dict(im=im, lb=None))['im'].unsqueeze(0).cuda()      # 按均值 mean=(0.3257, 0.3690, 0.3223) 和標準差 std=(0.2112, 0.2148, 0.2115) 對圖像進行標準化（逐通道操作）# PyTorch 模型通常接受批處理輸入（即使只有一張圖像），格式為 [B, C, H, W]。如果沒有 unsqueeze(0)，模型會因維度不匹配而報錯。假設輸入圖像經過 ToTensor 轉換后的形狀是 [C, H, W]（通道、高度、寬度）：原始 Tensor：(3, 224, 224) （RGB 圖像，3 通道，224x224 分辨率）
unsqueeze(0) 后：(1, 3, 224, 224) （添加了 batch 維度，表示 batch_size=1）

對于視頻流，那么這里的B就會變大，比如一次處理列表里面的8幀，就是一個批次8幀

② 輸入尺寸調整

im = F.interpolate(im, size=new_size, align_corners=False, mode='bilinear')

• 作用：將輸入圖像 im 的尺寸調整為 new_size（之前計算的32的倍數，如 512x704）。
• 參數說明：
  • size=new_size：目標尺寸（高度和寬度）。
  • mode='bilinear'：使用雙線性插值（平衡速度和質量）。
  • align_corners=False：避免邊緣像素對齊問題（推薦設置）。
• 為什么需要調整輸入尺寸？
  確保輸入尺寸符合模型的架構要求（如全卷積網絡需要尺寸能被下采樣因子整除）。

---

③ 模型推理

out = net(im)[0]

• 作用：將調整后的圖像輸入模型 net，獲取輸出。

• [0] 的含義：

  如果模型返回多個輸出（如主輸出 + 輔助頭輸出），[0]表示只取主輸出。例如：

  • 輸出形狀可能是 [B, C, H, W]（Batch, 類別數, 高, 寬）。
  • 若 B=1（單張圖像），out 的形狀為 [C, H, W]。

---

④ 輸出尺寸還原

out = F.interpolate(out, size=org_size, align_corners=False, mode='bilinear')

• 作用：將模型輸出的特征圖 out 插值回原始圖像尺寸 org_size（如 480x640）。
• 為什么需要還原尺寸？
  • 模型輸出的分辨率通常比輸入低（因下采樣），需還原到原始尺寸以便后續處理（如可視化或評估）。
  • 例如：
    • 輸入尺寸：512x704 → 模型輸出：16x22（下采樣32倍）→ 插值回 480x640。

---

⑤ 類別預測

out = out.argmax(dim=1)		# 張量模式，每一個像素值都是一種類別，用相應的數字表示！

• 作用：對特征圖 out 沿 類別維度（dim=1） 取 argmax，得到每個像素的預測類別。
• 輸入輸出形狀：
  • 輸入 out：[C, H, W]（C是類別數，如分割任務有20類）。
  • 輸出 out：[H, W]，每個像素值為 0 到 C-1 的整數（表示類別ID）。
• 示例：
  如果 C=3（3類：背景、貓、狗），out 的每個像素值會是 0、1 或 2。

---

⑥ 保存繪制

先轉換張量

out = out.squeeze().detach().cpu().numpy()      # 將張量轉換為NumPy數組
pred = palette[out]# 生成256*3的數組，每個數處于0-255之間
palette = np.random.randint(0, 256, (256, 3), dtype=np.uint8)		# 一個預定義的調色板（顏色映射表），形狀為 [num_classes, 3]，等價下面
palette = np.array([		[0, 0, 0],      # 類別0：黑色（背景）[255, 0, 0],    # 類別1：紅色（貓）[0, 255, 0],    # 類別2：綠色（狗）# ...其他類別顏色
], dtype=np.uint8)

⑦ 關于視頻流

[視頻文件]|v
生產者進程 ──(in_q)──> 主進程 ──(out_q)──> 消費者進程|                      |                    || (幀: [1,3,1080,1920])| (結果: [1,1080,1920])|v                      v                    v預處理                [GPU批量推理]           保存為MP4|                      |v                      v
BGR→RGB              合并8幀→[8,3,768,768]|                      |v                      v
Tensor化             縮放→推理→還原分辨率

3 pth轉onnx

3.1 為什么要轉換

pth這種權重加載時候需要定義具體的網絡模型，而onnx則不需要。而且pth每一次推理耗時會更久，但onnx則每一次減少一半以上

所以這里把pth權重轉換為onnx

(segment) D:\wpj\Bisnetv1>python tools/demo.py --config configs/bisenetv1_ipm0121.py --weight-path res/IPM0121/model_epoch_21.pth --img-path ./datasets/ipm_all_annoted_0121/images/training/0112.jpg
[張量] 耗時: 0.0000 秒
[圖像讀取與預處理] 耗時: 0.0036 秒
[圖像轉化為tensor] 耗時: 0.0180 秒
[推理：] 耗時: 0.7578 秒					# 耗時在0.75-0.8s
[預測：] 耗時: 0.0001 秒

3.2 轉換代碼

① 整個onnx模型導出還是需要借助BiSenetV1網絡模型，配置文件也是需要的

② 設置基本的輸入輸出，以及算子版本

torch.set_grad_enabled(False)parse = argparse.ArgumentParser()
parse.add_argument('--config', dest='config', type=str,default='configs/bisenetv2.py',)
parse.add_argument('--weight-path', dest='weight_pth', type=str,default='model_final.pth')
parse.add_argument('--outpath', dest='out_pth', type=str,default='model.onnx')
parse.add_argument('--aux-mode', dest='aux_mode', type=str,						# 注意，preddefault='pred')
args = parse.parse_args()
cfg = set_cfg_from_file(args.config)											# 加載配置文件
if cfg.use_sync_bn: cfg.use_sync_bn = False						net = model_factory[cfg.model_type](cfg.n_cats, aux_mode=args.aux_mode)			# 定義Bisenet網絡模型，用了aux_mode，會指定pred對應的argmax函數，所以說這里有可能報錯，如果你的模型不指定這個，估計就不會調用這個；所以今天遇到的問題，最簡單解決方式，應該是把aux_mode設置為eval模式，直接返回輸出！！！！
net.load_state_dict(torch.load(args.weight_pth, map_location='cpu',				weights_only=True), strict=False)				# 加載權重	
net.eval()dummy_input = torch.randn(1, 3, *cfg.cropsize)							#  dummy_input = torch.randn(1, 3, 1024, 1024)
input_names = ['input_image']											# 輸入和輸出，以及dynamic_axes設置形式一般這樣
output_names = ['preds']
dynamic_axes = {'input_image': {0: 'batch'}, 'preds': {0: 'batch'}}		# 指定模型的哪些輸入/輸出維度是動態的torch.onnx.export(net, dummy_input, 						# 輸入args.out_pth,						# 輸出路徑input_names=input_names, 			# 輸入名output_names=output_names,			# 輸出名verbose=False, opset_version=16,                   #  算子集版本原來是18，不兼容這個pytorch版本，改成16---這里pytorch-1.13# operator_export_type=torch.onnx.OperatorExportTypes.ONNX,dynamic_axes=dynamic_axes)

③ 一直報錯是關于BiSeNetV1網絡中forward函數中關于CustomArgMax的調用，這個CustomArgMax是一個自定義的argmax（在對一系列類別預測中，選擇概率最大的哪一個），在train和pred模式下是沒有調用，故此沒有問題；但是在轉換模型，即為pred模式下，發生了錯誤！因為它在symbolic返回中填的錯誤算子，系統找不到！

④ 定義Bisenet網絡模型，用了aux_mode，會指定pred對應的argmax函數，所以說這里有可能報錯，如果你的模型不指定這個，估計就不會調用這個；所以今天遇到的問題，最簡單解決方式，應該是把aux_mode設置為eval模式，直接返回輸出！經過實驗，確實如此！！因為他寫的CustomArgMax存在問題，eval是后面修改過的函數！！

class BiSeNetV1(nn.Module):def __init__(self, n_classes, aux_mode='train', *args, **kwargs):super(BiSeNetV1, self).__init__()......def forward(self, x):H, W = x.size()[2:]feat_cp8, feat_cp16 = self.cp(x)feat_sp = self.sp(x)feat_fuse = self.ffm(feat_sp, feat_cp8)feat_out = self.conv_out(feat_fuse)if self.aux_mode == 'train':feat_out16 = self.conv_out16(feat_cp8)feat_out32 = self.conv_out32(feat_cp16)return feat_out, feat_out16, feat_out32elif self.aux_mode == 'eval':return feat_out,elif self.aux_mode == 'pred':#  feat_out = feat_out.argmax(dim=1)feat_out = CustomArgMax.apply(feat_out, 1)return feat_outelse:raise NotImplementedErrorclass CustomArgMax(torch.autograd.Function):@staticmethoddef forward(ctx, feat_out, dim):if torch.onnx.is_in_onnx_export():# ONNX導出時：返回原始logits（保持梯度流）ctx.mark_dirty(feat_out)  # 避免警告return feat_out# 訓練/推理時：正常執行argmaxreturn feat_out.argmax(dim=dim).int()				# 實際上，pred應該也可以使用argmax，在推理時候就不可以使用了！！@staticmethoddef symbolic(g, feat_out, dim: int):# 導出為Identity節點（避免引入ArgMax）return g.op("Identity", feat_out)			# 原始代碼return g.op('CustomArgMax', feat_out, dim_i=dim)，這里執行不了         

3.3 轉換測試

python ./tools/xx.py --config configs/xx.py --weight-path ./res/x/x.pth

可視化模型，忘記放代碼了，后面補一下

python ./tools/x.py		# 可視化onnx權重，打印輸入輸出形狀# Dtype: 1   Dtype 是 Data Type 的縮寫，代表張量的數據類型。  eg, 1	torch.float32    2	torch.float64
===== Inputs =====
Name: input_image, Shape: ['batch', 3, 1024, 1024], Dtype: 1===== Outputs =====
Name: preds, Shape: ['batch', 2, 1024, 1024], Dtype: 1		# 2是類別數

推理測試 ，推理測試代碼待補充

python ./tools/x.py --config configs/xxx.py --weight-path ./model.onnx --img-path ./datasets/xxx/x.jpg		# 推理一張

python ./tools/xx.py		# 推理一個文件內的所有IPM圖像

4 onnx轉MNN

4.1 編譯MNN庫

下載MNN庫，為了轉換

git clone https://github.com/alibaba/MNN.git
cd MNN
mkdir build && cd build
cmake .. -DMNN_BUILD_CONVERTER=ON -DMNN_OPENCL=ON -DMNN_VULKAN=ON  # opencl是為了GPU推理
make -j$(nproc)			# 工具位于 ./build/MNNConvert
sudo make install

./MNNConvert -f ONNX --modelFile xxx.onnx --MNNModel xxx.mnn ./MNNConvert --framework=ONNX --model=your_model.onnx --MNNModel=output.mnn --optimizeLevel=3 感覺沒啥用 --bizCode=your_code 版權-f ONNX	指定輸入模型格式為 ONNX
--modelFile	輸入模型的路徑（這里是 model.onnx）	
--MNNModel	輸出模型的路徑（轉換后的 MNN 模型保存為 model.mnn）	--fp16	啟用 FP16 量化（減小模型體積，提升推理速度，可能損失少量精度）
--weightQuantBits 8	INT8 權重量化（進一步壓縮模型，需校準數據）
--optimizeLevel 2	優化級別（0不優化，1基礎優化，2激進優化，可能改變計算圖結構）
--inputConfigFile	指定動態輸入的形狀（通過 JSON 文件配置，解決 None 維度問題）
--forTraining	保留訓練相關算子（默認轉為推理模型）
--saveStaticModel	固定輸入形狀（移除動態維度，提升推理性能）

4.2 模型推理測試

??關于c++調用MNN模型的代碼還是有很多坑的（python代碼比較簡單），官方文檔中例子也給的很奇怪，代碼后續給出來

5 BisenetV2測試

?? BisenetV2優勢在于用更加輕量模塊來替換V1中原有Backbone即Resnet-18網絡。論文中V2的速度應該是156 FPS / 65 FPS ≈ 2.4于V1。但實際測試完全達不到論文結果，原始論文代碼未開源，論文中關于網絡細節不像V1清楚，搜索發現網上復現達不到論文效果。

板卡測試—8核 –ubunt20.04------mnn模型

V2（13 MB）相較于 V1（50 MB）參數量下降，但計算貌似沒有太大區別----因此不同模型的參數量不能直接作為推理速度標準，只可參考

單線程下V2（682 ms）快于V1網絡（840 ms）158 ms

2線程下V2（435 ms）快于V1網絡（481 ms）46 ms

4線程下V2（356 ms）快于V1網絡（370 ms）14 ms

8線程下V2（363 ms）慢于V1網絡（351 ms）11 ms

6 tee工具

??在測試過程中發現，想要保存終端打印的信息，又不想自己寫保存代碼，還得編譯cpp，找到tee這個工具，可以在linux終端命令行執行時，捕捉終端打印記錄并保存，還是很好用的。

python test.py  2>&1 | tee output.txt