当前位置：首页 > news >正文

pytorch使用tensorboardX面板自动生成模型结构图和各类可视化图像

news 2025/6/24 16:06:36

总结：

在原本代码中额外添加如下几行即可实现查看模型结构：

    from tensorboardX import SummaryWriter  # 用于进行可视化# 1. 来用tensorflow进行可视化with SummaryWriter("./log", comment="sample_model_visualization") as sw:  sw.add_graph(modelviz, sampledata)

操作步骤如下

安装完torch之后，再安装tensorboardX

pip install tensorboardX -i https://pypi.tuna.tsinghua.edu.cn/simple

运行下面代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from tensorboardX import SummaryWriter  # 用于进行可视化 class modelViz(nn.Module):def __init__(self):super(modelViz, self).__init__()self.conv1 = nn.Conv2d(3, 16, 3, 1, padding=1)self.bn1 = nn.BatchNorm2d(16)self.conv2 = nn.Conv2d(16, 64, 3, 1, padding=1)self.bn2 = nn.BatchNorm2d(64)self.conv3 = nn.Conv2d(64, 10, 3, 1, padding=1)self.bn3 = nn.BatchNorm2d(10)def forward(self, x):x = self.bn1(self.conv1(x))x = F.relu(x)x = self.bn2(self.conv2(x))x = F.relu(x)x = self.bn3(self.conv3(x))x = F.relu(x)return xif __name__  == "__main__":# 首先来搭建一个模型modelviz = modelViz()# 创建输入sampledata = torch.rand(1, 3, 4, 4)# 看看输出结果对不对out = modelviz(sampledata)print(out)  # 测试有输出，网络没有问题# 1. 来用tensorflow进行可视化with SummaryWriter("./log", comment="sample_model_visualization") as sw:sw.add_graph(modelviz, sampledata)# # 2. 保存成pt文件后进行可视化# torch.save(modelviz, "./log/modelviz.pt")

运行代码后会在"./log"路径下生成一个tfevents文件，在终端中进入代码的主目录下执行命令：

tensorboard --logdir=./

然后会输出

(base) jie@dell:~/桌面/fno_task$ tensorboard --logdir=./
TensorFlow installation not found - running with reduced feature set.NOTE: Using experimental fast data loading logic. To disable, pass"--load_fast=false" and report issues on GitHub. More details:https://github.com/tensorflow/tensorboard/issues/4784Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.16.2 at http://localhost:6006/ (Press CTRL+C to quit)

http://localhost:6006/

然后按照提示打开浏览器，输入上面这个网址就可以看到我们搭建的网络结构了，如下图所示，可以双击打开每一个节点查看其内容。也可以查看详细的结构以及每一层的输入输出shape。通过双击模型的组件实现展示网络细节和收起细节。