LAB 01 · 学会一个卷积核

from pathlib import Path
import copy, hashlib, inspect, json, math, platform, time
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn
from torch.nn import functional as F
from IPython.display import display, Markdown

torch.manual_seed(20260921)
torch.set_num_threads(2)
torch.use_deterministic_algorithms(True)
OUT=Path('results'); OUT.mkdir(exist_ok=True)
FIG=Path('figures'); FIG.mkdir(exist_ok=True)
BLUE,GOLD,TEAL,RED='#002676','#FDB515','#00897B','#B83A3A'
plt.rcParams.update({'font.size':11,'axes.spines.top':False,'axes.spines.right':False,
                     'figure.dpi':110,'savefig.facecolor':'white'})
def savefig(fig,name):
    fig.savefig(FIG/f'{name}.png',dpi=160,bbox_inches='tight')
    fig.savefig(FIG/f'{name}.svg',bbox_inches='tight')
def table(headers,rows):
    display(Markdown('| '+' | '.join(headers)+' |\n| '+ ' | '.join(['---']*len(headers))+' |\n'+
                     '\n'.join('| '+' | '.join(map(str,row))+' |' for row in rows)))
print('PyTorch',torch.__version__,'/ Python',platform.python_version(),'/ device CPU')

from numpy.lib.stride_tricks import sliding_window_view
image=np.arange(9,dtype=np.float64).reshape(3,3)
kernel=np.array([[0.,1.],[2.,3.]])
patches=sliding_window_view(image,kernel.shape)
response=np.einsum('ijab,ab->ij',patches,kernel)
assert np.array_equal(response,[[19.,25.],[37.,43.]])
fig,axes=plt.subplots(1,3,figsize=(9,2.7))
for ax,a,title in zip(axes,[image,kernel,response],['Input','One shared kernel','Response']):
    ax.imshow(a,cmap='Blues')
    for (i,j),v in np.ndenumerate(a):
        ax.text(j,i,f'{v:g}',ha='center',va='center',color='white' if v>a.max()*.6 else BLUE)
    ax.set(title=title,xticks=[],yticks=[])
savefig(fig,'lab01-shared-correlation');plt.show()
print(response)

rng=np.random.default_rng(20260921)
X=rng.normal(size=(2,3,6,7));K=rng.normal(size=(4,3,3,2));b=rng.normal(size=4)
P=sliding_window_view(X,(3,2),axis=(-2,-1))
Y_numpy=np.einsum('ncijab,ocab->noij',P,K)+b[None,:,None,None]
Y_torch=F.conv2d(torch.from_numpy(X),torch.from_numpy(K),torch.from_numpy(b))
np.testing.assert_allclose(Y_numpy,Y_torch.numpy(),rtol=1e-12,atol=1e-12)
conv=nn.Conv2d(3,16,3,padding=1)
assert sum(p.numel() for p in conv.parameters())==448
print('NumPy vs PyTorch max error:',np.abs(Y_numpy-Y_torch.numpy()).max())
print('weight:',tuple(conv.weight.shape),'bias:',tuple(conv.bias.shape),'total parameters:',448)
print('32x32 output:',tuple(conv(torch.zeros(1,3,32,32)).shape))
print('64x64 output:',tuple(conv(torch.zeros(1,3,64,64)).shape))

torch.manual_seed(20260921)
x=torch.randn(2,2,5,6,dtype=torch.float64)
w=torch.randn(3,2,3,2,dtype=torch.float64,requires_grad=True)
y=F.conv2d(x,w);delta=torch.randn_like(y)
(y*delta).sum().backward()
manual=torch.zeros_like(w)
for n in range(x.shape[0]):
    for o in range(w.shape[0]):
        for c in range(w.shape[1]):
            for a in range(w.shape[2]):
                for b_ in range(w.shape[3]):
                    manual[o,c,a,b_] += (delta[n,o]*x[n,c,a:a+y.shape[2],b_:b_+y.shape[3]]).sum()
torch.testing.assert_close(manual,w.grad)
contributions=delta[:,0]*x[:,0,:y.shape[2],:y.shape[3]]
fig,axes=plt.subplots(1,2,figsize=(9,3.4),layout='constrained')
limit=float(contributions.abs().max())
for n,ax in enumerate(axes):
    im=ax.imshow(contributions[n],cmap='RdBu_r',vmin=-limit,vmax=limit)
    ax.set(title=f'Sample {n}',xlabel='Output column',ylabel='Output row')
fig.suptitle('Per-position contributions to the shared weight K[0,0,0,0]')
fig.colorbar(im,ax=axes,shrink=.8);savefig(fig,'lab01-gradient-contributions');plt.show()
print('Sum over samples and positions:',float(contributions.sum()))
print('Autograd weight gradient:',float(w.grad[0,0,0,0]))
gradient_error=float((manual-w.grad).abs().max())
print('Full gradient maximum difference:',gradient_error)

columns=F.unfold(x,kernel_size=(3,2))
Y_matrix=(w.detach().reshape(3,-1)@columns).reshape_as(y)
torch.testing.assert_close(Y_matrix,y)
matrix_gradient=torch.einsum('nol,nkl->ok',delta.flatten(2),columns).reshape_as(w)
torch.testing.assert_close(matrix_gradient,w.grad)
print('Input:',tuple(x.shape),'unfold:',tuple(columns.shape),'kernel matrix:',tuple(w.reshape(3,-1).shape))
print('Input entries:',x.numel(),'unfold entries:',columns.numel())

torch.manual_seed(20260921)
inputs=torch.randn(32,1,8,8)
targets=F.conv2d(inputs,torch.tensor([[[[1.,-1.]]]]))
learner=nn.Conv2d(1,1,(1,2),bias=False)
losses=[]
for step in range(60):
    learner.zero_grad(set_to_none=True)
    loss=F.mse_loss(learner(inputs),targets);loss.backward()
    with torch.no_grad(): learner.weight-=.15*learner.weight.grad
    losses.append(float(loss.detach()))
print('Learned kernel:',learner.weight.detach().flatten().tolist())
assert torch.max(torch.abs(learner.weight.detach().flatten()-torch.tensor([1.,-1.])))<1e-5
fig,ax=plt.subplots(figsize=(6,3));ax.semilogy(losses,color=BLUE)
ax.set(xlabel='Update',ylabel='MSE',title='Recover a known operator');savefig(fig,'lab01-learned-kernel');plt.show()

H,p,d,k,s=8,1,1,3,3
expected=(H+2*p-d*(k-1)-1)//s+1
actual=F.conv2d(torch.randn(1,1,H,H),torch.ones(1,1,k,k),stride=s,padding=p).shape[-1]
assert expected==actual==3
torch.manual_seed(20260921)
z=torch.randn(1,1,8,8);weight=torch.randn(1,1,3,3)
def circular(z,s=1): return F.conv2d(F.pad(z,(1,1,1,1),mode='circular'),weight,stride=s)
errors={
 'periodic stride1 / input shift1 / output shift1':float((circular(z.roll(1,-1))-circular(z).roll(1,-1)).abs().max()),
 'periodic stride2 / input shift2 / output shift1':float((circular(z.roll(2,-1),2)-circular(z,2).roll(1,-1)).abs().max()),
 'periodic stride2 / input shift1 / output shift1':float((circular(z.roll(1,-1),2)-circular(z,2).roll(1,-1)).abs().max()),
 'zero pad stride1 / circular input shift1':float((F.conv2d(z.roll(1,-1),weight,padding=1)-F.conv2d(z,weight,padding=1).roll(1,-1)).abs().max()),
}
assert list(errors.values())[0]<1e-10 and list(errors.values())[1]<1e-10
table(['Condition','Maximum absolute difference'],[[k,f'{v:.6g}'] for k,v in errors.items()])

source_record={'torch':torch.__version__,'git_version':torch.version.git_version,
               'file':inspect.getsourcefile(nn.Conv2d)}
print(json.dumps(source_record,indent=2))
print(inspect.getsource(nn.Conv2d.forward))
print(inspect.getsource(nn.Conv2d._conv_forward))
summary={'hand_output':response.tolist(),'rgb_conv_parameters':448,'shared_gradient_max_error':gradient_error,
         'equivariance_errors':errors,'learned_kernel':learner.weight.detach().flatten().tolist(),
         'source':source_record,'execution':'local CPU; complete notebook'}
(OUT/'lab01-results.json').write_text(json.dumps(summary,indent=2))

import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(2)
C_IN, C_OUT, KH, KW = 2, 3, 2, 2  # 再试：3, 1, 1, 3
assert 1 <= C_IN <= 4 and 1 <= C_OUT <= 4
assert 1 <= KH <= 4 and 1 <= KW <= 5

def corr2d_loop(x, kernel):
    kh, kw = kernel.shape
    h, w = x.shape[0] - kh + 1, x.shape[1] - kw + 1
    result = x.new_zeros((h, w))
    for i in range(h):
        for j in range(w):
            result[i, j] = (x[i:i+kh, j:j+kw] * kernel).sum()
    return result

def multi_in(x, kernels):
    assert x.shape[0] == kernels.shape[0]  # 防止 zip 静默截断通道
    return sum(corr2d_loop(xc, kc) for xc, kc in zip(x, kernels))

class LoopConv2d(nn.Module):
    def __init__(self, cin, cout, kh, kw):
        super().__init__()
        values = torch.arange(cout*cin*kh*kw, dtype=torch.float64)
        self.weight = nn.Parameter(values.reshape(cout, cin, kh, kw) / 20)
        self.bias = nn.Parameter(torch.linspace(-0.2, 0.2, cout, dtype=torch.float64))
    def forward(self, x):
        return torch.stack([multi_in(x, k) for k in self.weight]) + self.bias[:, None, None]

x = torch.arange(C_IN*4*5, dtype=torch.float64).reshape(C_IN, 4, 5) / 10
layer = LoopConv2d(C_IN, C_OUT, KH, KW)
y = layer(x)
ref_weight = layer.weight.detach().clone().requires_grad_()
ref_bias = layer.bias.detach().clone().requires_grad_()
y_ref = F.conv2d(x[None], ref_weight, ref_bias)[0]
torch.testing.assert_close(y, y_ref)
upstream = torch.linspace(-1, 1, y.numel(), dtype=x.dtype).reshape_as(y)
(y * upstream).sum().backward()
(y_ref * upstream).sum().backward()
torch.testing.assert_close(layer.weight.grad, ref_weight.grad)
torch.testing.assert_close(layer.bias.grad, ref_bias.grad)
parts = torch.stack([corr2d_loop(xc, kc.detach()) for xc, kc in zip(x, layer.weight[0])])
print('Input:', tuple(x.shape), 'Weight:', tuple(layer.weight.shape), 'Output:', tuple(y.shape))
print('Registered parameters:', [(name, tuple(p.shape)) for name, p in layer.named_parameters()])
print('For output[0,0,0], one contribution per input channel:', parts[:, 0, 0].tolist())
print('Their sum + bias:', float(parts[:, 0, 0].sum() + layer.bias[0].detach()))
print('Full output:\n', y.detach())
HISTORY_RESULT = {
    'output_shape': list(y.shape),
    'parameters': sum(p.numel() for p in layer.parameters()),
    'forward_max_error': float((y.detach()-y_ref.detach()).abs().max()),
    'weight_grad_max_error': float((layer.weight.grad-ref_weight.grad).abs().max()),
    'bias_grad_max_error': float((layer.bias.grad-ref_bias.grad).abs().max()),
}
print('Checks:', HISTORY_RESULT)

import torch
from torch.nn import functional as F

torch.set_num_threads(2)
DELTA, BRIGHTNESS = 1.0, 1.0  # 再试：2.5, 3.0
assert BRIGHTNESS != 0 and abs(DELTA) <= 100
constant = torch.full((1, 1, 2, 5), BRIGHTNESS, dtype=torch.float64)
probe = torch.tensor([[[[0., 1., 0., 1., 0.], [1., 0., 1., 0., 1.]]]], dtype=torch.float64)
k1 = torch.tensor([[[[1., -1.]]]], dtype=torch.float64)
k2 = k1 + DELTA * torch.tensor([[[[1., -1.]]]], dtype=torch.float64)
y1, y2 = F.conv2d(constant, k1), F.conv2d(constant, k2)
p1, p2 = F.conv2d(probe, k1), F.conv2d(probe, k2)
A_constant = F.unfold(constant, (1, 2))[0].T
A_probe = F.unfold(probe, (1, 2))[0].T
null_direction = torch.tensor([1., -1.], dtype=torch.float64)
torch.testing.assert_close(y1, y2)
torch.testing.assert_close(A_constant @ null_direction, torch.zeros(A_constant.shape[0], dtype=torch.float64))
rank_constant = int(torch.linalg.matrix_rank(A_constant))
rank_combined = int(torch.linalg.matrix_rank(torch.cat([A_constant, A_probe])))
assert rank_constant == 1 and rank_combined == 2
if DELTA != 0:
    assert not torch.allclose(p1, p2)
print('K1:', k1.flatten().tolist(), 'K2:', k2.flatten().tolist())
for name, inp, a, b in [('Constant observations', constant, y1, y2), ('New alternating probe', probe, p1, p2)]:
    print('\n' + name, '\ninput:\n', inp[0, 0])
    print('K1 response:\n', a[0, 0], '\nK2 response:\n', b[0, 0])
HISTORY_RESULT = {
    'constant_response_max_gap': float((y1-y2).abs().max()),
    'probe_response_max_gap': float((p1-p2).abs().max()),
    'constant_design_rank': rank_constant,
    'combined_design_rank': rank_combined,
    'unknown_kernel_values': k1.numel(),
}
print('\nChecks:', HISTORY_RESULT)