Netron可视化网络结构

netron可视化在线地址：https://netron.app/
netron Windows安装包下载

方式一：代码部分（没有细节信息）：

if __name__ == '__main__':
    net = ResidualAttentionModel_448input()
    #其中ResidualAttentionModel_448input()为自己搭建的网络结构
    net.cuda()

    print('net: {}'.format(net))

    path = "../checkpoints/"
    model = path+'best.pth'
    torch.save(net, model)
    model = torch.load(model)  # 加载

第一步：使用以上代码生成best.pt文件。
第二步：采用netron软件加载生成的best.pth文件。
结果如下图所示，可以发现 没有网络结构的细节信息。

因此，采用方式二来保存网络结构的细节信息。


方式二：（保存细节信息）代码部分：

if __name__ == '__main__':
    net = ResidualAttentionModel_448input()
    #其中ResidualAttentionModel_448input()为自己搭建的网络结构
    net.cuda()

    print('net: {}'.format(net))

    path = "../checkpoints/"
    model = path+'best.onnx'
    # torch.save(net, model)
    # model = torch.load(model)  # 加载

    import torch
    import torch.onnx
    from torch.autograd import Variable
    import onnx
    from onnx import shape_inference

    x = Variable(torch.randn(1, 3, 448, 448)).cuda()
    torch_out = torch.onnx.export(net, x ,model,export_params=True,opset_version=11)

    onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)

第一步：采用以上代码生成best.onnx文件
第二步：下载netron软件，打开保存好的模型（best.onnx），即可看到已经加载了网路结构的细节信息。

分别介绍各部分代码作用：

1、ResidualAttentionModel_448input()为自己搭建的网络结构

net = ResidualAttentionModel_448input()
net.cuda()

2、将pytorch生成的模型转换为onnx格式

import torch
import torch.onnx
from torch.autograd import Variable

x = Variable(torch.randn(1, 3, 448, 448)).cuda()
torch_out = torch.onnx.export(net, x ,model,export_params=True,opset_version=11)

3、保存模型结构的细节信息

import onnx
from onnx import shape_inference

onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)

以下是踩坑过程记录：

踩坑一： RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

解决方式：添加 net.cuda()
net = ResidualAttentionModel_448input()
net.cuda()

原因：是cpu和cdua 使用的不一致，这里model和可视化都使用的是cuda()

踩坑二： UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9. This operator might cause results to not match the expected results by PyTorch.
ONNX’s Upsample/Resize operator did not match Pytorch’s Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch’s behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator.
“” + str(_export_onnx_opset_version) + ". "

解决方式：添加 opset_version=11（pytorch版本必须为1.13以上）
torch_out = torch.onnx.export(net, x ,model,export_params=True,opset_version=11)
net.cuda()