打印卷积层和最大池化层输出维度

问题描述

当阅读神经网络代码时，经常会有很多各种各样的层，比如有一个这样的卷积层 nn.Conv2d(3, 48, kernel_size=5, stride=2, padding=0)，我想知道在执行完这行代码后，输出的张量维度为多少。

已知输入维度是 (224, 224, 3)，分别表示 (width, height, channel)

问题解决

import torch.nn as nn
import torch
conv = nn.Conv2d(3, 48, kernel_size=5, stride=2, padding=0) # 卷积层
# 构建一个维度正确的输入
x = torch.rand(1, 3, 224, 224)
res = conv(x)
print(res.shape) # 输出 torch.Size([1, 48, 110, 110])

附

卷积层输出张量维度计算公式

Conv2d：output = ((input_size - kernel_size + 2 * padding) / stride ) + 1 向下取整

举例：

如果输入维度是 (224, 224, 3)，卷积层为 nn.Conv2d(in_channels=3, out_channels=48, kernel_size=5, stride=2, padding=0)，计算公式如下：

(224-5+2*0)/2+1 =110.5

所以输出的维度就是 (110, 110, 48)

最大池化层输出张量维度计算公式

MaxPool2d：(input_size - kernel_size) / stride + 1 向下取整

举例

如果输入维度为 (106, 106, 96)，最大池化层为 nn.MaxPool2d(kernel_size=3, stride=2)，计算公式如下：

(106-3)/2+1=52.5

所以输出的维度就是 (52, 52, 96)

代码验证如下：

x = torch.rand(96, 106, 106)
maxpooling = nn.MaxPool2d(kernel_size=3, stride=2)
res = maxpooling(x)
print(res.shape)	# 输出结果为 torch.Size([96, 52, 52])

多个卷积或池化层的网络

import torch.nn as nn
import torch

# 网络结构
net = nn.Sequential(
	nn.Conv2d(3, 48, kernel_size=5, stride=2, padding=0), # 卷积层 (224, 224, 3) -> (110, 110, 48)
	nn.MaxPool2d(kernel_size=3, stride=2)	# (110, 110, 48) -> (54, 54, 48)
)
# 构建输入
x = torch.rand(1, 3, 224, 224)
# 将输入特征向量送入网络
res = net (x)	# torch.Size([1, 48, 54, 54])