您可以像下面这样遍历模型的子代并打印尺寸以进行调试。这类似于前写,但您编写一个单独的函数而不是创建一个 nn.Module 类。
import torch
from torch import nn
model = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Conv2d(20,64,5),
nn.ReLU()
)
def print_sizes(model, input_tensor):
output = input_tensor
for m in model.children():
output = m(output)
print(m, output.shape)
return output
input_tensor = torch.rand(100, 1, 28, 28)
print_sizes(model, input_tensor)
# output:
# Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1)) torch.Size([100, 20, 24, 24])
# ReLU() torch.Size([100, 20, 24, 24])
# Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1)) torch.Size([100, 64, 20, 20])
# ReLU() torch.Size([100, 64, 20, 20])
# you can also nest the Sequential models like this. In this case inner Sequential will be considered as module itself.
model1 = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Sequential(
nn.Conv2d(20,64,5),
nn.ReLU()
)
)
print_sizes(model1, input_tensor)
# output:
# Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1)) torch.Size([100, 20, 24, 24])
# ReLU() torch.Size([100, 20, 24, 24])
# Sequential(
# (0): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
# (1): ReLU()
# ) torch.Size([100, 64, 20, 20])