PyTorch 学习笔记(二): 可视化与模型参数计算

PyTorch 学习笔记(二):可视化与模型参数计算

1. 可视化

from models import Darknet
from torchviz import make_dot, make_dot_from_trace
import torch
from tensorboardX import SummaryWriter

# torchviz 可视化
model = torch.nn.Sequential()
model.add_module('W0', torch.nn.Linear(8, 16))
model.add_module('tanh', torch.nn.Tanh())
model.add_module('W1', torch.nn.Linear(16, 1))
x = torch.randn(1,8)
g = make_dot(model(x), params=dict(model.named_parameters()))
g.view()

# hiddenlayer 可视化
# pip install hiddenlayer
import torchvision
# 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
    # 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
g = hl.build_graph(model, torch.zeros([1, 3, 224, 224]), transforms=transforms)
g.save('1.pdf')

# tensorboardx 可视化
 writer = SummaryWriter(logdir="./logs/", comment="TestView")
 with writer:
     writer.add_graph(model, input_to_model=torch.rand(1, 8))
# 命令行窗口输入 tensorboard --logdir logs
# 浏览器输入以下网址
# http://localhost:6006

2. 计算模型参数

# 计算模型参数个数
def get_parameter_number(net):
    total_num = sum(p.numel() for p in net.parameters())
    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}
print(get_parameter_number(model))
# 结果: {'Trainable': 161, 'Total': 161}

参考文献
1.可视化参考
2.计算参数参考