LAB 03 · 架构实验室

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')

def rf_table(layers):
    r,j=1,1;rows=[]
    for name,k,s,d in layers:
        r+=(k-1)*d*j;j*=s;rows.append([name,k,s,d,r,j])
    return rows
rf=rf_table([('conv3',3,1,1),('conv3',3,1,1),('pool2',2,2,1),('conv3',3,1,1)])
table(['Layer','k','s','d','RF extent','Jump'],rf)
assert rf[-1][-2:]==[10,2]
three=rf_table([('conv3',3,1,1)]*3)
assert three[-1][-2]==7
print('Equal-channel kernel weights: three 3x3 = 27 C²; one 7x7 = 49 C²')

inp=torch.zeros(1,1,13,13,requires_grad=True)
value=inp
for _ in range(3):value=F.conv2d(value,torch.ones(1,1,3,3),padding=1)
value[0,0,6,6].backward();sensitivity=inp.grad[0,0]
where=torch.nonzero(sensitivity)
extent=(where.max(0).values-where.min(0).values+1).tolist()
assert extent==[7,7]
fig,ax=plt.subplots(figsize=(4.8,4));im=ax.imshow(sensitivity,cmap='Blues')
ax.set(title='Three linear 3x3 layers: center-output gradient',xlabel='Input column',ylabel='Input row')
fig.colorbar(im,ax=ax,label='Derivative');savefig(fig,'lab03-receptive-field-gradient');plt.show()

x=torch.randn(2,3,5,6,dtype=torch.float64);w=torch.randn(4,3,1,1,dtype=torch.float64);b=torch.randn(4,dtype=torch.float64)
y=F.conv2d(x,w,b)
z=torch.einsum('oc,nchw->nohw',w[:,:,0,0],x)+b[None,:,None,None]
torch.testing.assert_close(y,z)
C,H,W,classes=64,7,7,10
head_flat=nn.Linear(C*H*W,classes);head_gap=nn.Linear(C,classes)
head_counts={'flatten_linear':sum(p.numel() for p in head_flat.parameters()),'gap_linear':sum(p.numel() for p in head_gap.parameters())}
table(['Head','Parameters'],list(head_counts.items()))
features=torch.randn(2,C,H,W)
torch.testing.assert_close(head_gap(features.mean((2,3))),F.conv2d(features,head_gap.weight[:,:,None,None],head_gap.bias).mean((2,3)))

direct=nn.Conv2d(192,32,5,padding=2,bias=False)
bottleneck=nn.Sequential(nn.Conv2d(192,32,1,bias=False),nn.ReLU(),nn.Conv2d(32,32,5,padding=2,bias=False))
direct_n=sum(p.numel() for p in direct.parameters());bottle_n=sum(p.numel() for p in bottleneck.parameters())
assert (direct_n,bottle_n)==(153600,31744)
print('Weights:',direct_n,'vs',bottle_n,'/ reduction:',direct_n/bottle_n)
class Inception(nn.Module):
    def __init__(self):
        super().__init__()
        self.a=nn.Sequential(nn.Conv2d(192,64,1),nn.ReLU())
        self.b=nn.Sequential(nn.Conv2d(192,96,1),nn.ReLU(),nn.Conv2d(96,128,3,padding=1),nn.ReLU())
        self.c=nn.Sequential(nn.Conv2d(192,32,1),nn.ReLU(),nn.Conv2d(32,32,5,padding=2),nn.ReLU())
        self.d=nn.Sequential(nn.MaxPool2d(3,stride=1,padding=1),nn.Conv2d(192,32,1),nn.ReLU())
    def forward(self,x):return torch.cat([self.a(x),self.b(x),self.c(x),self.d(x)],dim=1)
assert Inception()(torch.zeros(1,192,14,14)).shape==(1,256,14,14)
print('Concat output channels:',64+128+32+32)

class Residual(nn.Module):
    def __init__(self,cin,cout,stride=1):
        super().__init__()
        self.branch=nn.Sequential(nn.Conv2d(cin,cout,3,stride=stride,padding=1,bias=False),nn.BatchNorm2d(cout),nn.ReLU(),
                                  nn.Conv2d(cout,cout,3,padding=1,bias=False),nn.BatchNorm2d(cout))
        self.shortcut=nn.Identity() if cin==cout and stride==1 else nn.Sequential(nn.Conv2d(cin,cout,1,stride=stride,bias=False),nn.BatchNorm2d(cout))
    def forward(self,x):return F.relu(self.branch(x)+self.shortcut(x))
same=Residual(16,16);down=Residual(16,32,2)
assert same(torch.randn(2,16,14,14)).shape==(2,16,14,14)
assert down(torch.randn(2,16,14,14)).shape==(2,32,7,7)
v=torch.tensor([-1.,1.],dtype=torch.float64)
j1=torch.autograd.functional.jacobian(lambda v:v+.1*v,v)
j2=torch.autograd.functional.jacobian(lambda v:F.relu(v+.1*v),v)
j3=torch.autograd.functional.jacobian(lambda v:v-v,v)
print('I + J_F:\n',j1,'\nWith post-ReLU:\n',j2,'\nCancellation F=-x:\n',j3)

batch=torch.tensor([[[[1.,3.]]],[[[5.,7.]]]],dtype=torch.float64)
mode_rows=[]
for training in (True,False):
    for enabled in (True,False):
        bn=nn.BatchNorm2d(1,momentum=.1).double();bn.train(training)
        inp=batch.clone().requires_grad_()
        with torch.set_grad_enabled(enabled):
            value=bn(inp)
            if enabled:value.sum().backward()
        mode_rows.append([str(training),str(enabled),int(bn.num_batches_tracked),
                          round(float(bn.running_mean[0]),6),round(float(bn.running_var[0]),6),inp.grad is not None])
table(['training','grad enabled','Batches tracked','Running mean','Running variance','Input gradient exists'],mode_rows)
assert [row[2] for row in mode_rows]==[1,1,0,0]
bn=nn.BatchNorm2d(1,affine=False,momentum=.1).double().train()
with torch.no_grad():observed=bn(batch)
biased=batch.var(unbiased=False);unbiased=batch.var(unbiased=True)
torch.testing.assert_close(observed,(batch-batch.mean())/(biased+bn.eps).sqrt())
torch.testing.assert_close(bn.running_var,(.9+.1*unbiased).reshape(1))
print('Biased variance:',float(biased),'unbiased:',float(unbiased),'running:',float(bn.running_var[0]))
print('Parameters:',[n for n,_ in nn.BatchNorm2d(3).named_parameters()])
print('Buffers:',[n for n,_ in nn.BatchNorm2d(3).named_buffers()])

from torch.nn.modules.batchnorm import _BatchNorm
print('PyTorch',torch.__version__,'git',torch.version.git_version)
print('File:',inspect.getsourcefile(_BatchNorm))
print(inspect.getsource(_BatchNorm.forward))
bn_untracked=nn.BatchNorm2d(1,affine=False,track_running_stats=False).double().eval()
torch.testing.assert_close(bn_untracked(batch),(batch-batch.mean())/(batch.var(unbiased=False)+bn_untracked.eps).sqrt())
print('track_running_stats=False in eval uses current batch statistics: verified.')

summary={'receptive_field_rows':rf,'three_conv_gradient_extent':extent,'heads':head_counts,
         'inception_5x5':{'direct_weights':direct_n,'bottleneck_weights':bottle_n,'ratio':direct_n/bottle_n,'definition':'192->32->32; no bias'},
         'residual_jacobians':{'plain_residual':j1.tolist(),'post_relu':j2.tolist(),'cancellation':j3.tolist()},
         'batchnorm_modes':mode_rows,'batchnorm_variances':{'biased':float(biased),'unbiased':float(unbiased),'running_after_one_batch':float(bn.running_var[0])},
         'execution':'local CPU; full notebook; actual PyTorch source inspected','torch':torch.__version__}
(OUT/'lab03-results.json').write_text(json.dumps(summary,indent=2))
print('Saved results/lab03-results.json')

import copy
import torch
from torch import nn

torch.set_num_threads(2)
PEER_SHIFT, MOMENTUM = 8.0, 0.1  # 再试：0.0, 0.1；或 -4.0, 0.1
assert 0 <= MOMENTUM <= 1
anchor = torch.tensor([[[[1., 3.]]]], dtype=torch.float64)
peer_a = anchor.clone()
peer_b = peer_a + PEER_SHIFT
batch_a = torch.cat([anchor, peer_a], dim=0)
batch_b = torch.cat([anchor, peer_b], dim=0)
base = nn.BatchNorm2d(1, affine=True, momentum=MOMENTUM, track_running_stats=True).double()
with torch.no_grad():
    base.weight.fill_(1.0)
    base.bias.zero_()
    base.running_mean.fill_(2.0)
    base.running_var.fill_(4.0)

HISTORY_RESULT = {}
print('Anchor in both batches:', anchor.flatten().tolist())
print('Peer A:', peer_a.flatten().tolist(), '| Peer B:', peer_b.flatten().tolist())
for mode in ('train', 'eval'):
    paired = []
    for label, batch in [('A', batch_a), ('B', batch_b)]:
        bn = copy.deepcopy(base)
        bn.train(mode == 'train')
        with torch.no_grad():
            output = bn(batch)[0].clone()
        paired.append(output)
        print(mode, label, '| anchor output:', [round(v, 6) for v in output.flatten().tolist()],
              '| running mean/var:', round(bn.running_mean.item(), 6), round(bn.running_var.item(), 6),
              '| batches tracked:', bn.num_batches_tracked.item())
    gap = float((paired[0]-paired[1]).abs().max())
    HISTORY_RESULT[mode] = {'anchor_A': paired[0].flatten().tolist(), 'anchor_B': paired[1].flatten().tolist(), 'max_gap': gap}
    if mode == 'eval' or PEER_SHIFT == 0:
        torch.testing.assert_close(paired[0], paired[1])
    else:
        assert not torch.allclose(paired[0], paired[1])
print('Checks:', HISTORY_RESULT)

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

torch.set_num_threads(2)
BRANCH_GAIN = 0.0  # 再试：-0.5；-1.0
assert -2 <= BRANCH_GAIN <= 2

class SwitchableResidual(nn.Module):
    def __init__(self, gain):
        super().__init__()
        self.branch = nn.Conv2d(1, 1, 1, bias=False).double()
        with torch.no_grad():
            self.branch.weight.fill_(gain)
        self.branch.weight.requires_grad_(False)  # 只对输入求导，不改变分支
    def forward(self, x, skip, post_relu):
        pre_activation = self.branch(x) + (x if skip else torch.zeros_like(x))
        return F.relu(pre_activation) if post_relu else pre_activation

block = SwitchableResidual(BRANCH_GAIN).eval()
base_input = torch.tensor([-2., -0.5, 0.5, 2.], dtype=torch.float64).reshape(1, 1, 2, 2)
HISTORY_RESULT = {'branch_gain': BRANCH_GAIN, 'rows': []}
print('Fixed input:', base_input.flatten().tolist(), '| fixed branch gain:', BRANCH_GAIN)
print(f"{'skip':<7} {'post-ReLU':<11} {'output':<32} {'max |y-x|':<12} derivative")
for skip in (False, True):
    for post_relu in (False, True):
        x = base_input.clone().requires_grad_()
        y = block(x, skip=skip, post_relu=post_relu)
        dx, = torch.autograd.grad(y.sum(), x)
        values, derivatives = y.detach().flatten().tolist(), dx.flatten().tolist()
        error = float((y.detach()-base_input).abs().max())
        print(f'{str(skip):<7} {str(post_relu):<11} {str(values):<32} {error:<12.6g} {derivatives}')
        HISTORY_RESULT['rows'].append({'skip': skip, 'post_relu': post_relu, 'output': values, 'max_identity_error': error, 'd_sum_output_d_input': derivatives})
with torch.no_grad():
    off = block(base_input, skip=False, post_relu=False)
    on = block(base_input, skip=True, post_relu=False)
    torch.testing.assert_close(on-off, base_input)  # 真正隔离开关增加的项
    if BRANCH_GAIN == 0:
        torch.testing.assert_close(on, base_input)
        torch.testing.assert_close(block(base_input, True, True), F.relu(base_input))
    if BRANCH_GAIN == -1:
        torch.testing.assert_close(on, torch.zeros_like(base_input))
print('Skip-on minus skip-off, before activation:', (on-off).flatten().tolist())