import torch
from torch import nn


# 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
# class Con2d(nn.Module):
#     def __init__(self,kernel_size):
#         super().__init__()
#         self.weight=nn.Parameter(torch.ones(kernel_size))
#         self.bias=nn.Parameter(torch.zeros(1))
#     def forward(self,X):
#         return corr2d(X,self.weight)+self.bias
X = torch.ones((6, 8))
X[:, 2:6] = 0
K = torch.tensor([[1.0, -1.0]])
Y= torch.tensor([[ 0.,  1.,  0.,  0.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.,  0., -1.,  0.]])
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)
    l=(y_hat-Y)**2
    conv2d.zero_grad()
    l.sum().backward()
    conv2d.weight.data[:]-=3e-2*conv2d.weight.grad
    if (i+1)%2==0:
        print(f'batch{i+1},loss{l.sum():.3f}')