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

最新推荐文章于 2025-03-07 20:30:52 发布

Wils0nEdwards

最新推荐文章于 2025-03-07 20:30:52 发布

阅读量1.1k

点赞数 9

文章标签：深度学习人工智能

本文链接：https://blog.csdn.net/coldasice342/article/details/140526283

版权

方法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.