对Pytorch网络进行可视化

最新推荐文章于 2023-07-21 14:05:06 发布

yhwang-hub

最新推荐文章于 2023-07-21 14:05:06 发布

阅读量427

点赞数

CC 4.0 BY-SA版权

分类专栏： TensorRT-部署-加速 TVM

本文链接：https://blog.csdn.net/qq_33287871/article/details/115921897

TensorRT-部署-加速同时被 2 个专栏收录

52 篇文章 ¥19.90 ¥99.00

订阅专栏

超级会员免费看

TVM

41 篇文章 ¥19.90 ¥99.00

订阅专栏

超级会员免费看

本文介绍了三种在Pytorch中可视化神经网络的方法：1)利用tensorboard进行可视化，需在Ubuntu系统安装graphviz，生成的网络结构可通过tensorboard查看；2)使用netron工具，可以直接在线或下载使用进行模型结构可视化；3)借助torchviz生成pdf文件进行网络结构展示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

对pytorch模型进行可视化主要包括以下几种方法:
1.使用tensorboard
2.使用netron
3.使用torchviz

1.使用tensorboard进行可视化

# -*-coding:utf-8 -*-
# --------------------
# author: cjs
# time: 20200910
# usage: 进行pytorch模型的可视化
# packages： pytorch, tensorflow, tensorboard, tensorboardX
# --------------------

import torch
import torch.nn as nn
import torch.nn.functional as F
from tensorboardX import SummaryWriter  # 用于进行可视化
from torchviz import make_dot


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)