AutoDL中怎么在本地查看使用netron可视化得到的模型结构图

首先需要安装netron包

pip install netron

然后需要自定义autodl服务

在 Autodl 实例这里的自定义服务里实名认证后开通自定义服务，Autodl 目前只开放了 6006 端口给自定义服务，

在 autodl 实例中指定6006端口启动netron服务

执行这个py代码后，

# 启动 Netron
import netron
# 启动 Netron 并指定端口为6006
import torch.onnx
import os

# 启动 Netron 并指定端口为6006
os.system('netron --port 6006 ema_model.onnx')

可以看到服务启动在6006端口了。

这里的 ema_model.onnx 是通过另外一个代码利用torch.onnx.export导出保存的 onnx 文件

import torch.onnx

import torch
from torchviz import make_dot
import torch.nn as nn
# 定义模型
class EMA(nn.Module):
    def __init__(self, channels, c2=None, factor=32):
        super(EMA, self).__init__()
        self.groups = factor
        assert channels // self.groups > 0
        self.softmax = nn.Softmax(-1)
        self.agp = nn.AdaptiveAvgPool2d((1, 1))
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)
        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)
        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        b, c, h, w = x.size()
        group_x = x.reshape(b * self.groups, -1, h, w)  # b*g,c//g,h,w
        x_h = self.pool_h(group_x)
        x_w = self.pool_w(group_x).permute(0, 1, 3, 2)
        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))
        x_h, x_w = torch.split(hw, [h, w], dim=2)
        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())
        x2 = self.conv3x3(group_x)
        x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        x12 = x2.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
        x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        x22 = x1.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
        weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)
        return (group_x * weights.sigmoid()).reshape(b, c, h, w)


# 定义模型
model = EMA(channels=64)
x = torch.randn(1, 64, 128, 128)

# 导出为 ONNX 格式，然后在 netron.app上把这个onnx文件拖进去就可以看到了
torch.onnx.export(model, x, "ema_model.onnx")

然后在本地电脑上，打开cmd窗口，输入

ssh -CNgv -L 8097:127.0.0.1:8097 root@region-1.autodl.com -p 37881

需要把其中的8097修改成6006，root@region-1.autodl.com 和 37881 替换成自己模型实例的地址和端口，然后会提示输入密码

输入正确的密码之后，在本地浏览器输入 127.0.0.1:8097 就能成功看到 netro 可视化得到的模型结构图了，这里我查看的是EMA模块的结构图