import torch
from torch import nn
from d2l import torch as d2l

class Conv2D(nn.Module):
    def __init__(self,kernel_size):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(kernel_size)) #如kernel_size= (2,2),则随机初始化一个2x2的卷积
        self.bias = nn.Parameter(torch.zeros(1)) #bias初始化为0
    def forward(self,X):
        return self.corr2d(X,self.weight) + self.bias
    
    # 卷积操作函数
    def corr2d(self,X,K):
        h,w = K.shape
        # 卷积输出大小
        Y = torch.zeros((X.shape[0] - h + 1),X.shape[1] - w + 1)
        for i in range(Y.shape[0]):
            for j in range(Y.shape[1]):
                # (i,j)是目标区域的左上角坐标
                Y[i,j] = (X[i:i+h,j:j+w] * K).sum()
        return Y

# 测试
X = torch.tensor([[0.0,1.0,2.0],
                  [3.0,4.0,5.0],
                  [6.0,7.0,8.0]])
model = Conv2D(kernel_size=(2,2))
Y = model.forward(X)
print(Y)