将pytorch框架网络转换为onnx格式，并使用netron可视化HRNet的网络结构

文章目录

目录

前言

一、环境配置

1.配置运行环境

2.安装netron

二、pytorch到onnx

1.定位主干网络

2.onnx格式转换 

附：HRNet emantic-segmentation运行的调试

---

前言

       torch.onnx的官方文档及绝大多数博客只给出了转换简单网络的过程，这里我们以转换HRNet网络为例，给出转换过程及主干网络的可视化结果；netron对于原生pytorch网络支持不够好，转换为onnx格式后可以用netron方便的可视化复杂网络结构。

       文章主要聚焦于用netron将pytorch框架的复杂网络结构可视化，未对onnx格式做详细测试。

一、环境配置

Ubuntu 20.04

pytorch 1.10

1.配置运行环境

HRNet代码来自于@Bubbliiiing 在github发布的图像分割模型github地址https://github.com/bubbliiiing/hrnet-pytorch

克隆代码（略）

安装相关依赖

pip install -r requirements.txt

 对于requirements.txt中的包，安装最新版本即可。

安装onnx

pip install onnx

2.安装netron

github地址https://github.com/lutzroeder/netron

任意安装方式均可，这里下载其appimage文件，下载后执行

chmod a+x Netron-5.7.0.AppImage
./netron-5.7.0.AppImage

 即可运行

二、pytorch到onnx

1.定位主干网络

       使用torch.onnx函数进行格式转换时需要在程序中定位到主干网络，大多数博客引用网络较为简单，而对于本地的复杂网络如HRNet，由多个子网络子程序组成。该HRNet分割网络主干相对路径是nets.hrnet.HRnet：

class HRnet(nn.Module):
    def __init__(self, num_classes = 21, backbone = 'hrnetv2_w18', pretrained = False):
        super(HRnet, self).__init__()
        self.backbone       = HRnet_Backbone(backbone = backbone, pretrained = pretrained)

        last_inp_channels   = np.int(np.sum(self.backbone.model.pre_stage_channels))

        self.last_layer = nn.Sequential(
            nn.Conv2d(in_channels=last_inp_channels, out_channels=last_inp_channels, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(last_inp_channels, momentum=BN_MOMENTUM),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=last_inp_channels, out_channels=num_classes, kernel_size=1, stride=1, padding=0)
        )

    def forward(self, inputs):
        H, W = inputs.size(2), inputs.size(3)
        x = self.backbone(inputs)
        
        # Upsampling
        x0_h, x0_w = x[0].size(2), x[0].size(3)
        x1 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
        x2 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
        x3 = F.interpolate(x[3], size=(x0_h, x0_w), mode='bilinear', align_corners=True)

        x = torch.cat([x[0], x1, x2, x3], 1)

        x = self.last_layer(x)
        x = F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True)
        return x

2.onnx格式转换 

在程序主目录新建.py文件并运行，程序如下：

from nets.hrnet import HRnet
import torch.onnx

batch_size = 1 #可以取其他值
torch_model = HRnet()

# torch_model.load_state_dict(torch.load('model_data/hrnetv2_w32_weights_voc.pth'), strict=False)
torch_model.eval()

x = torch.randn(batch_size, 3, 480, 480, requires_grad=True) #模拟输入数据尺寸
torch_out = torch_model(x)

torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "test.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True, # store the trained parameter weights inside the model file
                  opset_version = 11,
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})

其中

# torch_model.load_state_dict(torch.load('model_data/hrnetv2_w32_weights_voc.pth'), strict=False)

解除注释后程序运行可能报错，猜测是参数与网络结构不完全匹配导致，有解决方法的大佬请指点一下，感谢。 

执行上述程序后可以发现在程序主目录出现了我们需要的onnx格式文件：

打开netron并导入test.onnx，可以看到HRNet的网络结构已经显示出来啦，可以帮助我们方便的学习代码及文章，如下图所示（网络结构复杂，只截取部分）：

 作为对比，下图展示netron直接读取HRNet的训练权重文件出现的效果：

 可以看到其对pytorch的支持不算完善。

附：HRNet emantic-segmentation运行的调试

1.报错：

No module named ‘torchvision.models.utils’,无法加载load_state_dict_from_url...

由于torchvision的版本更新，对应的函数转移到了新版本的torch里面

将backbone.py中

from torchvision.models.utils import load_state_dict_from_url

改为

from torch.hub import load_state_dict_from_url

2.使用预测程序对视频进行分割测试时报错：

If you are on Ubuntu or Debian, install libgtk2.0-dev and pkg-config, then re-run cmake or config...

解决办法：

解决办法https://blog.csdn.net/tudou2013goodluck/article/details/108402055以上两个问题解决后就可以进行分割预测了

3.文章的转换方法仍然存在瑕疵，转换格式时加上.pth文件就会报错，希望知道原因的大佬可以指出解决方法，谢谢。

---