PyTorch 模型输入形状

PyTorch model input shape

我加载了一个自定义 PyTorch 模型,我想找出它的输入形状。像这样:

model.input_shape

是否可以获取此信息?


更新: print()summary() 没有显示这个模型的输入形状,所以它们不是我要找的。

print(model)

会给你一个模型的总结,在这里你可以看到每一层的形状。

您也可以使用 pytorch-summary 包。

如果你的网络有一个 FC 作为第一层,你可以很容易地计算出它的输入形状。你提到你在前面有一个卷积层。由于也存在全连接层,网络将仅针对一种特定的输入大小生成输出。我建议通过使用各种形状来解决这个问题,即喂养具有某种形状的玩具批次,然后检查 FC 层之前的 Conv 层的输出。

因为这取决于第一个 FC 层之前的网络架构(conv 层数、内核等),所以我无法为您提供正确输入的确切公式。如前所述,您必须通过尝试各种输入形状以及在第一个 FC 之前生成的网络输出来解决这个问题。有(几乎)总有一种方法可以用代码解决问题,但我现在想不出别的方法。

PyTorch 的灵活性

PyTorch 模型是非常灵活的对象,以至于它们不强制或通常期望数据的固定输入形状。

如果您有某些图层,则可能会有限制,例如:

  • 展平后跟宽度为 N 的完全连接层将强制原始输入的尺寸 (M1 x M2 x ... Mn) 的乘积等于 N
  • N 个输入通道的二维卷积会将数据强制为 3 维,第一维的大小为 N

但如您所见,这些都没有强制执行数据的 形状。

We might not realize it right now, but in more complex models, getting the size of the first linear layer right is sometimes a source of frustration. We’ve heard stories of famous practitioners putting in arbitrary numbers and then relying on error messages from PyTorch to backtrack the correct sizes for their linear layers. Lame, eh? Nah, it’s all legit!

  • 使用 PyTorch 进行深度学习

调查

简单情况:第一层全连接

如果您的模型的第一层是全连接层,那么 print(model) 中的第一层将详细说明单个样本的预期维度。

模棱两可的案例:CNN

但是,如果它是一个卷积层,因为它们是动态的,并且会在输入允许的情况下以 long/wide 步幅移动,所以没有简单的方法可以从模型本身检索此信息。 1 这种灵活性意味着对于许多体系结构 多个兼容的输入大小 2 都将被网络接受。

这是 PyTorch 的一个特性

人工检查

您需要做的是研究网络架构,一旦您找到一个可解释层(如果存在,例如完全连接),就可以“逆向”其维度,确定前面的层(例如池化和卷积)有 compressed/modified 它。

例子

例如在以下来自 Deep Learning with PyTorch (8.5.1) 的模型中:

class NetWidth(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(16 * 8 * 8, 32)
        self.fc2 = nn.Linear(32, 2)
    
    def forward(self, x):
        out = F.max_pool2d(torch.tanh(self.conv1(x)), 2)
        out = F.max_pool2d(torch.tanh(self.conv2(out)), 2)
        out = out.view(-1, 16 * 8 * 8)
        out = torch.tanh(self.fc1(out))
        out = self.fc2(out)
        return out

我们看到模型接受输入 2.d。具有 3 个通道和:

的图像
  • Conv2d -> 将其发送到具有32个通道的相同大小的图像
  • max_pool2d(,2) -> 将每个维度的图像大小减半
  • Conv2d -> 将其发送到具有16个通道的相同大小的图像
  • max_pool2d(,2) -> 将每个维度的图像大小减半
  • view -> 重塑图像
  • Linear -> 获取大小为 16 * 8 * 8 的张量并发送到大小为 32
  • ...

所以逆向计算,我们有:

  • 形状为16 * 8 * 8
  • 的张量
  • 未重塑形状(通道 x 高 x 宽)
  • un-max_pooled 在 2d 中,因子为 2,因此高度和宽度未减半
  • 未从 16 个通道卷积到 32 个
    假设:产品中的16很可能是指通道数,view看到的图像是形状(通道,8,8),并且目前是(频道,16,16)2
  • un-max_pooled 在 2d 中,因子为 2,因此高度和宽度再次减半(通道,32,32)
  • 未从 32 个通道卷积到 3 个

所以假设 kernel_size 和填充足以使卷积本身保持图像尺寸,输入图像很可能是形状 (3,32,32),即 RGB 32x32 像素正方形图像。


备注:

  1. 即使是外包也需要你提供输入形状才能显示每一层输出的形状

  2. 然而,它可以是任何 2 个乘积等于 8*8 的数字,例如(64,1)、(32,2)、(16,4) 等,但是由于代码写为 8*8,所以作者很可能使用了实际尺寸。

您可以从模型参数中的第一个张量获取输入形状。

例如创建一些模型:

class CustomNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1568, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 20)

    def forward(self, x):
        out = self.fc1(x)
        out = F.relu(out)
        out = self.fc2(out)
        out = F.relu(out)
        out = self.fc3(out)
        return out

model = CustomNet()

所以 model.parameters() 方法 returns 是 torch.Tensor class 模块参数的迭代器。查看文档 https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.parameters

第一个参数是输入张量。

first_parameter = next(model.parameters())
input_shape = first_parameter.size()