深度学习中模型结构可视化的两个方法
本文主要说两个方式
- torchsummary库
- 关于一个可视化的函数，原文链接是：可视化（输出pdf）
由于 我们需要展示，这里我使用CycleGAN的利用残差块的生成器
代码如下：
class ResidualBlock(nn.Module):
def __init__(self, in_features):
super(ResidualBlock, self).__init__()
conv_block = [ nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features) ]
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x):
return x + self.conv_block(x)
class GeneratorResNet(nn.Module):
def __init__(self, in_channels, out_channels, res_blocks ):
super(GeneratorResNet, self).__init__()
#in_channels = args.input_nc
#out_channels = args.output_nc
#res_blocks = args.n_residual_blocks
# Initial convolution block
model = [ nn.ReflectionPad2d(3),
nn.Conv2d(in_channels, 64, 7),
nn.InstanceNorm2d(64),
nn.ReLU(inplace=True) ]
# Downsampling
in_features = 64
out_features = in_features*2
for _ in range(2):
model += [ nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True) ]
in_features = out_features
out_features = in_features*2
# Residual blocks
for _ in range(res_blocks):
model += [ResidualBlock(in_features)]
# Upsampling
out_features = in_features//2
for _ in range(2):
model += [ nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True) ]
in_features = out_features
out_features = in_features//2
# Output layer
model += [ nn.ReflectionPad2d(3),
nn.Conv2d(64, out_channels, 7),
nn.Tanh() ]
self.model = nn.Sequential(*model)
def forward(self, x):
return self.model(x)
1. torchsummary
1.1 使用torchsummary需要先安装，（有虚拟环境打开虚拟环境）
pip install torchsummary
1.2 使用
我们只需在上面的代码基础上使用
from torchsummmary import summary
#1.初始一个上面模型的对象
model = GeneratorResNet(3,3,9)
# 2.使用summary
#参数说明：
#参数1：model对象
#参数2：模型的一个输入shape，例如你是用MNIST，则shape=(1,28,28)
print(summary(model,(3,256,256))) #这里我们设置shape为（3,256,256）
注意：出现错误时的解决方法
方法1：
#①
print(summary(model,(3,256,256)).cuda())
#②
#或者如果你已经设置了device,则使用下面的代码
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(summary(model,(3,256,256)).to(device))
方法2：
print(summary(model,(3,256,256),device="cpu"))
2. 关于一个文件visualize.py
具体使用就是你创建一个py文件，把下面的代码复制进去，在其他 py文件使用时，导入这个visualize.py文件即可
2.1 visualize.py内容
def make_dot(var, params=None):
""" Produces Graphviz representation of PyTorch autograd graph
Blue nodes are the Variables that require grad, orange are Tensors
saved for backward in torch.autograd.Function
Args:
var: output Variable
params: dict of (name, Variable) to add names to node that
require grad (TODO: make optional)
"""
if params is not None:
assert isinstance(params.values()[0], Variable)
param_map = {id(v): k for k, v in params.items()}
node_attr = dict(style='filled',
shape='box',
align='left',
fontsize='12',
ranksep='0.1',
height='0.2')
dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"))
seen = set()
def size_to_str(size):
return '(' + (', ').join(['%d' % v for v in size]) + ')'
def add_nodes(var):
if var not in seen:
if torch.is_tensor(var):
dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange')
elif hasattr(var, 'variable'):
u = var.variable
name = param_map[id(u)] if params is not None else ''
node_name = '%s\n %s' % (name, size_to_str(u.size()))
dot.node(str(id(var)), node_name, fillcolor='lightblue')
else:
dot.node(str(id(var)), str(type(var).__name__))
seen.add(var)
if hasattr(var, 'next_functions'):
for u in var.next_functions:
if u[0] is not None:
dot.edge(str(id(u[0])), str(id(var)))
add_nodes(u[0])
if hasattr(var, 'saved_tensors'):
for t in var.saved_tensors:
dot.edge(str(id(t)), str(id(var)))
add_nodes(t)
add_nodes(var.grad_fn)
return dot
2.2 具体使用
我们依然将上面的生成器模型作为例子
#1. 初始化一个model对象（后面的3,3,9可以不用管，这是初始化类需要的参数）
model = GeneratorResNet(3,3,9)
#2.这个x就是一个默认传入，其与上面的summary的第二个参数差不多，但是这个需要设置一个batchsize，我这里写成1。
x = torch.rand(1,3, 256, 256)
#3.通过model得到一个输出y
y = model(x)
#4.调用visualize.py文件的函数make_dot
g = make_dot(y)
g.view()
之后你的项目目录下就会出现一个pdf，打开就是（例如下面的样子），对于了解模型的架构非常清晰
3.总结
- 方法1，我们可以在控制台看到每一层的输出shape，以及参数数量等。
- 方法2，对于我们从github上得到的工程，可以快速的理解模型的架构。
