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.計算參數參考