pytorch卷积层代码原理实现

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

def corr2d(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]):
            Y[i,j]=(X[i:i+h,j:j+w]*K).sum()
    return Y

X=torch.tensor([[1,2,3],[2,3,4],[4,5,6,]])
k=torch.tensor([[1,2],[2,3]])
print(corr2d(X,k))

# 二维卷积层的实现
class Conv2D(nn.Module):
    def __init__(self,kernel_size):
        super().__init__()
        self.weight=nn.Parameter(torch.rand(kernel_size))
        self.bias=nn.Parameter(torch.zeros(1))

    def foward(self,x):
        return corr2d(x,self.weight)+self.bias

x=torch.ones((6,8))
x[:,2:6]=0
print(x)

k=torch.tensor([[1,-1]])
y=corr2d(x,k)
print(y)

print(corr2d(x.t(),k))


conv2d=nn.Conv2d(1,1,kernel_size=(1,2),bias=False)
X=x.reshape((1,1,6,8))
Y=y.reshape((1,1,6,7))
for i in range(10):
    Y_hat=conv2d(X)
    loss=(Y_hat-Y)**2
    conv2d.zero_grad()
    loss.sum().backward()
    conv2d.weight.data[:]-=3e-2*conv2d.weight.grad
    if (i+1)%2==0:
        print(f'batch{i+1},loss{loss.sum():.3f}')

# 卷积核权重张量
print(conv2d.weight.data.reshape((1,2)))