如何在 Pytorch 中自动计算卷积层的形状？ [重复]

Question

我想自动计算任何Pytorch卷积层的输出维度。 当创建更大的 CNN 时，最后一个卷积的输出的计算。图层变得非常乏味。 有没有一种方法可以在不使用 nn.LazyLinear() 的情况下自动执行此操作？

代码示例：

self.conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2)

self.feature_dim = [ADD CODE HERE]

self. fc1 = nn.Linear(self.feature_dim, 256)

谢谢，非常感谢您的帮助！

我尝试过使用公式

[(W−K+2P)/S]+1

并在代码中实现它，但对每一层执行此操作几乎比简单地手动计算要实现更多的工作。 有没有统一的方法直接从 Pytorch 读取输出形状？

Answer 1

您可以使用“元”设备创建“假张量”，无需分配任何内存并使用触发器进行计算。

请参阅以下最小示例：

import torch
from torch import nn

conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=2)
conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2)

x = torch.zeros([1, 1, 8192, 8192], device="meta")

y = conv1(x)
y = conv2(y)
y = conv3(y)

print(y.shape)

这里有一些关于“假张量”的更多注释：https://pytorch.org/torchdistx/latest/fake_tensor_and_deferred_init.html

我希望这有帮助！