方法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 = factorassert channels // self.groups > 0self.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,wx_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, hwx21 = 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, hwweights = (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 = factorassert channels // self.groups > 0self.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,wx_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, hwx21 = 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, hwweights = (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')
执行上面这个代码之后,然后会在当前目录下生成一个 png 文件和另一个文件