PyTorch使用技巧2：netron可视化模型结构

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx

import netron


class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(64, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 32, 1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64)
        )

        self.conv1 = nn.Conv2d(3, 64, 3, padding=1, bias=False)
        self.output = nn.Sequential(
            nn.Conv2d(64, 1, 3, padding=1, bias=True),
            nn.Sigmoid()
        )

    def forward(self, x):

        x = self.conv1(x)

        identity = x
        x = F.relu(self.block1(x) + identity)
        x = self.output(x)
        # print("x:", x.shape) # 在不知道torch.rand(1, 3, 416, 416)的情况下可以使用这个命令来查找
        return x


print(model())


d = torch.rand(1, 3, 416, 416)
# print(len(d[0][1]))
# print(d[0])
m = model()
o = m(d)



onnx_path = "onnx_model_name.onnx"
torch.onnx.export(m, d, onnx_path)
netron.start(onnx_path)

参考资料：
https://blog.csdn.net/qq_30263737/article/details/114268646
https://blog.csdn.net/weixin_43863869/article/details/121915379
https://blog.csdn.net/weixin_43183872/article/details/108329776