深度学习根据代码可视化模型结构图的方法

news2025/1/10 16:20:40

方法1. Netron

Netron 是一个支持多种深度学习模型格式的可视化工具，可以将 PyTorch 模型转换为 ONNX 格式，然后使用 Netron 进行可视化。

安装 Netron：

pip install netron

使用示例：

import torch.onnx

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

# 导出为 ONNX 格式
torch.onnx.export(model, x, "ema_model.onnx")

# 启动 Netron
import netron
netron.start('ema_model.onnx')

在浏览器中打开 http://localhost:8080/ 即可查看模型的结构图。

方法2. tensorboard

TensorBoard 是 TensorFlow 的可视化工具，也可以与 PyTorch 集成，适用于查看模型的计算图和训练过程。

安装 TensorBoard：

pip install tensorboard

使用示例：

import torch
from torch.utils.tensorboard import SummaryWriter

# 定义模型
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)

# 创建一个SummaryWriter对象
writer = SummaryWriter('runs/ema_example')

# 创建模型和输入数据
model = EMA(channels=64)
x = torch.randn(1, 64, 128, 128)

# 添加计算图到TensorBoard
writer.add_graph(model, x)
writer.close()

然后使用以下命令启动 TensorBoard 并查看模型的计算图：

tensorboard --logdir=runs/ema_example

在浏览器中打开 http://localhost:6006/ 即可查看模型的结构图。

方法3. Torchviz

Torchviz 是一个可以将 PyTorch 模型的计算图转换为可视化结构图的工具。

需要先安装 graphviz

sudo apt-get update
sudo apt-get install graphviz

安装 Torchviz：

pip install torchviz

使用示例：

import torch
from torchviz import make_dot

# 定义模型
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)

# 生成计算图
y = model(x)
dot = make_dot(y, params=dict(model.named_parameters()))

# 保存计算图为文件
dot.format = 'png'
dot.render('ema_model')