调试或跳过顺序 pytorch 模型的最佳方法

Posted 2023-03-28

【中文标题】调试或跳过顺序 pytorch 模型的最佳方法

【英文标题】：Best way to debug or step over a sequential pytorch model

【发布时间】：2021-07-11 21:36:23

【问题描述】：

我曾经使用 nn.Module 编写 PyTorch 模型，其中包括 __init__ 和转发，以便我可以跨过我的模型以检查变量维度如何沿网络变化。 不过我后来意识到你也可以用nn.Sequential 来做，只需要__init__，你不需要编写如下的转发函数：

但是，问题是当我尝试跨过这个网络时，再检查变量并不容易。它只是跳到另一个地方然后返回。

有谁知道在这种情况下如何跨步？

P.S：我正在使用 PyCharm。

【参考方案1】：

您可以像下面这样遍历模型的子代并打印尺寸以进行调试。这类似于前写，但您编写一个单独的函数而不是创建一个 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])