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}