示例

import torch
from torch import nn
from torch.nn import functional as F 
import netron

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 10, (3, 3), stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(10)
        self.conv2 = nn.Conv2d(10, 32, (3, 3), stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        return x

if __name__=="__main__":
    net = Net()
    x = torch.randn(1, 3, 512, 512)
    onnx = torch.onnx.export(net, x, 'test.onnx')
    netron.start(onnx)

运行上面代码后会在浏览器弹出netron查看页面，打开test.onnx文件即可查看网络. 如图