Pytorch中查看CNN网络结构和各层输出参数

在使用Pytorch搭建多层CNN网络结构时，需要调整各层输入参数，一般输入参数可以通过计算获得，但是会比价麻烦，特别是遇到填充和修改步长之后。其实没那么复杂，只需要在定义完模型之后，模拟输入数据，就可以获得各层的输出。

1、查看整体结构

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(64, 128, 5, 1, 2), 
    nn.BatchNorm2d(128),
    nn.ReLU(), 
    nn.MaxPool2d(2)
)

print(model)

output：

Sequential(
  (0): Conv2d(64, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU()
  (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

2、查看各层size

x = torch.randn([10, 64, 25, 25])

for i in range(len(model)):
    x = model[i](x)
    print(x.size())

output：

torch.Size([10, 128, 25, 25])
torch.Size([10, 128, 25, 25])
torch.Size([10, 128, 25, 25])
torch.Size([10, 128, 12, 12])