pytorch模型可视化

1. 使用dot

1.1 安装graphviz和torchviz

sudo apt-get install graphviz

sudo pip install torchviz

1.2 使用torchviz

import torch
from torch import nn
from torchviz import make_dot, make_dot_from_trace

# Visualize gradients of simple MLP
# The method below is for building directed graphs of PyTorch operations, built during forward propagation and showing which operations will be called on backward. It omits subgraphs which do not require gradients.

model=nn.Sequential()
model.add_module("W0", nn.Linear(8, 16))
model.add_module("tanh", nn.Tanh())
model.add_module("W1", nn.Linear(16, 1))

x = Variable(torch.randn(1, 8))
y = model(x)

make_dot(y.mean(), params=dict(model.named_parameters()))  # 直接在ipython notebook中显示

dot=make_dot(y.mean(), params=dict(model.named_parameters()))
dot.render("model.pdf")  #保存为pdf

见github教程

2. 使用tensorwatch

安装pytorch = 1.2, tensorwatch = 0.8.7

import tensorwatch as tw
import torchvision.models

alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224]).save("alextnet.png")

教程： tensorwatch

3. 使用hiddenlayer

pip install torch===1.4.0 torchvision===0.5.0 -f https://download.pytorch.org/whl/torch_stable.html

pip install hiddenlayer

import torch
import torchvision.models
import hiddenlayer as hl

# Resnet101
model = torchvision.models.resnet101()

# Rather than using the default transforms, build custom ones to group
# nodes of residual and bottleneck blocks.
transforms = [
    # Fold Conv, BN, RELU layers into one
    hl.transforms.Fold("Conv > BatchNorm > Relu", "ConvBnRelu"),
    # Fold Conv, BN layers together
    hl.transforms.Fold("Conv > BatchNorm", "ConvBn"),
    # Fold bottleneck blocks
    hl.transforms.Fold("""
        ((ConvBnRelu > ConvBnRelu > ConvBn) | ConvBn) > Add > Relu
        """, "BottleneckBlock", "Bottleneck Block"),
    # Fold residual blocks
    hl.transforms.Fold("""ConvBnRelu > ConvBnRelu > ConvBn > Add > Relu""",
                       "ResBlock", "Residual Block"),
    # Fold repeated blocks
    hl.transforms.FoldDuplicates(),
]

# Display graph using the transforms above
dot = hl.build_graph(model, torch.zeros([1, 3, 224, 224]), transforms=transforms)
dot.attr("graph", rankdir="TD")
dot.render("resnet101")  # save as resnet101.pdf

教程: hiddenlayer