给定输入空间维度 w,2d 卷积层将在该维度上输出具有以下大小的张量:
int((w + 2*p - d*(k - 1) - 1)/s + 1)
完全相同的情况也适用于nn.MaxPool2d https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html。作为参考,您可以在此处查找PyTorch 文档 https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html.
模型的卷积部分由三个 (Conv2d + MaxPool2d) 块组成。您可以使用此辅助函数轻松推断输出的空间维度大小:
def conv_shape(x, k=1, p=0, s=1, d=1):
return int((x + 2*p - d*(k - 1) - 1)/s + 1)
递归调用它,您将得到最终的空间维度:
>>> w = conv_shape(conv_shape(32, k=3, p=1), k=2, s=2)
>>> w = conv_shape(conv_shape(w, k=3), k=2, s=2)
>>> w = conv_shape(conv_shape(w, k=3), k=2, s=2)
>>> w
2
由于您的卷积具有平方内核和相同的步幅、填充(水平等于垂直),因此上述计算适用于张量的宽度和高度尺寸。最后看最后一个卷积层conv3
,它有 64 个过滤器,在全连接层之前每批元素的最终元素数量为:w*w*64
, i.e. 256
.
但是,没有什么可以阻止您调用图层来找出输出形状!
class Net(nn.Module):
def __init__(self):
super().__init__()
self.feature_extractor = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2,2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2,2),
nn.Flatten())
n_channels = self.feature_extractor(torch.empty(1, 3, 32, 32)).size(-1)
self.classifier = nn.Sequential(
nn.Linear(n_channels, 200),
nn.ReLU(),
nn.Dropout(0.25),
nn.Linear(200, 100),
nn.ReLU(),
nn.Dropout(0.25),
nn.Linear(100, 10))
def forward(self, x):
features = self.feature_extractor(x)
out = self.classifier(features)
return out
model = Net()