import torch
from torch import nn
# 此函数初始化卷积层权重，并对输入和输出提高和缩减相应的维数
def comp_conv2d(conv2d, X):
# 这里的（1，1）表示批量大小和通道数都是1
#将输入张量 X 的形状调整为 (1, 1, height, width)
X = X.reshape((1,1) + X.shape)
Y = conv2d(X) #张量X为 8 * 8，经过conv2d，填充为1，变为10 * 10
#卷积核为3 * 3， 得到Y为 8 * 8
return Y.reshape(Y.shape[2:])
#将输出张量 Y 的形状从 (1, 1, new_height, new_width)
#变换为 (new_height, new_width)，去掉批量大小和通道数的维度。
# 请注意,这里每边都填充了1行或1列，因此总共添加了2行或2列
conv2d = nn.Conv2d(1, 1, kernel_size=3, padding=1)
X = torch.rand(size=(8, 8))
print(comp_conv2d(conv2d, X).shape)
# torch.Size([8, 8])
# 步幅
# 高度和宽度的步幅设置为2，从而将输入的高度和宽度减半
conv2d = nn.Conv2d(1, 1, kernel_size=3, padding=1, stride=2)
print(comp_conv2d(conv2d, X).shape)
# torch.Size([4, 4])
conv2d = nn.Conv2d(1, 1, kernel_size=(3, 5), padding=(0, 1), stride=(3, 4))
print(comp_conv2d(conv2d, X).shape)
# torch.Size([2, 2])
