LAB 02 · 小 CNN 与图像平移

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

SPLIT_SEED,MODEL_SEED=20260908,20260921
TRAIN_N,EPOCHS,BATCH_SIZE,LEARNING_RATE=12000,6,128,.001

import gzip, struct, urllib.request

# 读取完整标准 MNIST 文件；先用本地副本，没有时从 PyTorch 使用的公开镜像下载。
DATA = Path("data"); DATA.mkdir(exist_ok=True)
MNIST_FILES = {
    "train-images-idx3-ubyte.gz": "f68b3c2dcbeaaa9fbdd348bbdeb94873",
    "train-labels-idx1-ubyte.gz": "d53e105ee54ea40749a09fcbcd1e9432",
    "t10k-images-idx3-ubyte.gz": "9fb629c4189551a2d022fa330f9573f3",
    "t10k-labels-idx1-ubyte.gz": "ec29112dd5afa0611ce80d1b7f02629c",
}
data_identity = []
for filename, expected in MNIST_FILES.items():
    path = DATA / filename
    if not path.exists():
        url = "https://ossci-datasets.s3.amazonaws.com/mnist/" + filename
        try:
            with urllib.request.urlopen(url, timeout=60) as source:
                path.write_bytes(source.read())
        except Exception as exc:
            raise RuntimeError("MNIST 下载失败。请将课程材料包中 data/ 的四个 .gz 文件上传到项目 data/，再重跑。本实验不会退回伪造数据。") from exc
    raw = path.read_bytes()
    if hashlib.md5(raw).hexdigest() != expected:
        raise ValueError(f"MNIST 文件校验失败：{path}，请重新取得该文件。")
    data_identity.append({"file":filename, "bytes":len(raw), "md5":expected,
                          "sha256":hashlib.sha256(raw).hexdigest()})

def read_idx(filename):
    with gzip.open(DATA / filename, "rb") as source:
        magic, n = struct.unpack(">II", source.read(8))
        if magic == 2051:
            h, w = struct.unpack(">II", source.read(8))
            assert (h, w) == (28, 28)
            return np.frombuffer(source.read(), np.uint8).reshape(n, 1, h, w).copy()
        assert magic == 2049
        return np.frombuffer(source.read(), np.uint8).reshape(n).copy()

raw_train = read_idx("train-images-idx3-ubyte.gz")
raw_train_y = read_idx("train-labels-idx1-ubyte.gz")
raw_test = read_idx("t10k-images-idx3-ubyte.gz")
raw_test_y = read_idx("t10k-labels-idx1-ubyte.gz")
assert raw_train.shape == (60000, 1, 28, 28) and raw_test.shape == (10000, 1, 28, 28)

# 沿用第一讲划分 seed；本次只从原先 50k train 取固定 12k 子集。
partition = np.random.default_rng(SPLIT_SEED).permutation(60000)
train_ids, val_ids = partition[:TRAIN_N], partition[50000:]
assert not set(train_ids).intersection(val_ids)
train_x = torch.tensor(raw_train[train_ids], dtype=torch.float32) / 255
train_y = torch.tensor(raw_train_y[train_ids], dtype=torch.long)
val_x = torch.tensor(raw_train[val_ids], dtype=torch.float32) / 255
val_y = torch.tensor(raw_train_y[val_ids], dtype=torch.long)
test_x = torch.tensor(raw_test, dtype=torch.float32) / 255
test_y = torch.tensor(raw_test_y, dtype=torch.long)
print(f"MNIST: train={len(train_y):,} / validation={len(val_y):,} / test={len(test_y):,}")
print("12k training labels:", np.bincount(raw_train_y[train_ids]).tolist())

from collections import OrderedDict
def make_model(kind):
    if kind=='MLP':
        return nn.Sequential(OrderedDict([('flatten',nn.Flatten()),('hidden',nn.Linear(784,32)),
                                         ('relu',nn.ReLU()),('logits',nn.Linear(32,10))]))
    return nn.Sequential(OrderedDict([
        ('conv1',nn.Conv2d(1,8,3,padding=1)),('relu1',nn.ReLU()),('pool1',nn.MaxPool2d(2)),
        ('conv2',nn.Conv2d(8,16,3,padding=1)),('relu2',nn.ReLU()),('pool2',nn.MaxPool2d(2)),
        ('flatten',nn.Flatten()),('hidden',nn.Linear(784,32)),('relu3',nn.ReLU()),('logits',nn.Linear(32,10))]))
model_info={}
for kind in ('MLP','CNN'):
    torch.manual_seed(MODEL_SEED);m=make_model(kind)
    v=torch.zeros(1,1,28,28);macs=0;shapes=[]
    for name,layer in m.named_children():
        v=layer(v);shapes.append([name,str(tuple(v.shape))])
        if isinstance(layer,nn.Conv2d):macs+=v.numel()*layer.in_channels*math.prod(layer.kernel_size)//layer.groups
        if isinstance(layer,nn.Linear):macs+=layer.in_features*layer.out_features
    model_info[kind]={'parameters':sum(p.numel() for p in m.parameters()),'conv_linear_macs':macs,'shapes':shapes}
assert model_info['CNN']['parameters']==26698
table(['Model','Parameters','Forward MACs / sample'],[[k,v['parameters'],v['conv_linear_macs']] for k,v in model_info.items()])

def translate(batch,dx,dy):
    result=torch.zeros_like(batch)
    H,W=batch.shape[-2:]
    x0,x1=max(0,-dx),min(W,W-dx)
    y0,y1=max(0,-dy),min(H,H-dy)
    result[:,:,y0+dy:y1+dy,x0+dx:x1+dx]=batch[:,:,y0:y1,x0:x1]
    return result

def augment(batch,generator):
    shifts=torch.randint(-2,3,(len(batch),2),generator=generator)
    result=torch.zeros_like(batch)
    for dx in range(-2,3):
        for dy in range(-2,3):
            mask=(shifts[:,0]==dx)&(shifts[:,1]==dy)
            if mask.any():result[mask]=translate(batch[mask],dx,dy)
    return result

shifted_test=translate(test_x,2,0)
assert torch.equal(shifted_test[:,:,:,2:],test_x[:,:,:,:-2])
assert shifted_test[:,:,:,:2].count_nonzero()==0
fig,axes=plt.subplots(2,6,figsize=(10,3.8))
sample_ids=[0,1,2,3,4,5]
for j,i in enumerate(sample_ids):
    for row,arr in enumerate([test_x,shifted_test]):
        axes[row,j].imshow(arr[i,0],cmap='gray',vmin=0,vmax=1);axes[row,j].axis('off')
        axes[row,j].set_title(f'id {i}, label {int(test_y[i])}' if row==0 else 'Right 2 px')
savefig(fig,'lab02-fixed-test-shift');plt.show()

# 将来源记录绑定到这次真实训练与保存的文件，防止与旧 report 混用。
assert TRAIN_N==len(train_y)==len(train_ids)
assert torch.equal(train_x,torch.tensor(raw_train[train_ids],dtype=torch.float32)/255)
assert torch.equal(train_y,torch.tensor(raw_train_y[train_ids],dtype=torch.long))
assert torch.equal(val_x,torch.tensor(raw_train[val_ids],dtype=torch.float32)/255)
assert torch.equal(val_y,torch.tensor(raw_train_y[val_ids],dtype=torch.long))
training_source_protocol={
    'split_seed':int(SPLIT_SEED),'train_n':len(train_ids),'validation_n':len(val_ids),
    'partition_sha256':hashlib.sha256(partition.tobytes()).hexdigest(),
    'train_ids_sha256':hashlib.sha256(train_ids.tobytes()).hexdigest(),
    'validation_ids_sha256':hashlib.sha256(val_ids.tobytes()).hexdigest(),
    'files':copy.deepcopy(data_identity)}
checkpoint_provenance={}

@torch.no_grad()
def evaluate(model,x,y,batch_size=512):
    model.eval()
    logits=torch.cat([model(x[i:i+batch_size]) for i in range(0,len(x),batch_size)])
    pred=logits.argmax(1);correct=int((pred==y).sum())
    return {'n':len(y),'correct':correct,'accuracy':correct/len(y),
            'loss':float(F.cross_entropy(logits,y))},pred

runs={};trained_models={};predictions={};initial_hashes={}
for kind in ('MLP','CNN'):
    for use_aug in (False,True):
        label=kind+(' + shift augmentation' if use_aug else ' / no augmentation')
        torch.manual_seed(MODEL_SEED);model=make_model(kind)
        init_hash=hashlib.sha256(b''.join(p.detach().numpy().tobytes() for p in model.parameters())).hexdigest()
        initial_hashes[label]=init_hash
        optimizer=torch.optim.Adam(model.parameters(),lr=LEARNING_RATE)
        shuffle=torch.Generator().manual_seed(MODEL_SEED+1)
        augmentation_rng=torch.Generator().manual_seed(MODEL_SEED+2)
        history=[];best_loss=float('inf');best_state=None
        start=time.perf_counter()
        for epoch in range(1,EPOCHS+1):
            model.train();order=torch.randperm(TRAIN_N,generator=shuffle);total_loss=0.
            for idx in order.split(BATCH_SIZE):
                xb=augment(train_x[idx],augmentation_rng) if use_aug else train_x[idx]
                optimizer.zero_grad(set_to_none=True)
                loss=F.cross_entropy(model(xb),train_y[idx]);loss.backward();optimizer.step()
                total_loss+=float(loss.detach())*len(idx)
            val,_=evaluate(model,val_x,val_y)
            history.append({'epoch':epoch,'train_online_loss':total_loss/TRAIN_N,
                            'val_loss':val['loss'],'val_accuracy':val['accuracy']})
            if val['loss']<best_loss:
                best_loss=val['loss'];best_state=copy.deepcopy(model.state_dict());best_epoch=epoch
            print(f'{label} | epoch {epoch}: val loss={val["loss"]:.4f}, val acc={val["accuracy"]:.2%}')
        seconds=time.perf_counter()-start
        model.load_state_dict(best_state)
        clean,p0=evaluate(model,test_x,test_y);shifted,p1=evaluate(model,shifted_test,test_y)
        original_ok=p0==test_y;shifted_ok=p1==test_y
        transitions={'correct_to_correct':int((original_ok&shifted_ok).sum()),
                     'correct_to_wrong':int((original_ok&~shifted_ok).sum()),
                     'wrong_to_correct':int((~original_ok&shifted_ok).sum()),
                     'wrong_to_wrong':int((~original_ok&~shifted_ok).sum())}
        runs[label]={'architecture':kind,'augmentation':use_aug,'parameters':model_info[kind]['parameters'],
                     'conv_linear_macs':model_info[kind]['conv_linear_macs'],'history':history,
                     'selected_epoch':best_epoch,'seconds':seconds,'initial_state_sha256':init_hash,
                     'clean':clean,'right_2px':shifted,'prediction_changed':int((p0!=p1).sum()),'paired':transitions}
        trained_models[label]=model;predictions[label]=(p0,p1)
        checkpoint_name=kind.lower()+('-aug' if use_aug else '-plain')+'-state.pt'
        checkpoint_path=OUT/checkpoint_name
        torch.save(best_state,checkpoint_path)
        checkpoint_provenance[checkpoint_name]={
            'run_label':label,'sha256':hashlib.sha256(checkpoint_path.read_bytes()).hexdigest(),
            'source_protocol':copy.deepcopy(training_source_protocol)}
for kind in ('MLP','CNN'):
    hashes=[v for k,v in initial_hashes.items() if k.startswith(kind)]
    assert len(set(hashes))==1
print('All four runs complete; identical initial weights confirmed within each architecture.')

labels=list(runs)
table(['Condition','Selected epoch','Clean correct / 10000','Clean %','Right 2px %','Correct → wrong'],[
    [k,v['selected_epoch'],v['clean']['correct'],f"{100*v['clean']['accuracy']:.2f}",
     f"{100*v['right_2px']['accuracy']:.2f}",v['paired']['correct_to_wrong']] for k,v in runs.items()])
fig,axes=plt.subplots(1,2,figsize=(12,4))
pos=np.arange(4);width=.35
axes[0].bar(pos-width/2,[100*runs[k]['clean']['accuracy'] for k in labels],width,label='Clean',color=BLUE)
axes[0].bar(pos+width/2,[100*runs[k]['right_2px']['accuracy'] for k in labels],width,label='Right 2 px',color=TEAL)
axes[0].set(xticks=pos,xticklabels=['MLP','MLP + aug','CNN','CNN + aug'],ylabel='Accuracy (%)',ylim=(0,120),yticks=[0,20,40,60,80,100],title='Same 10,000 test images');axes[0].legend(loc='upper left',ncol=2)
for i,k in enumerate(labels):
    for offset,metric in [(-width/2,'clean'),(width/2,'right_2px')]:
        v=100*runs[k][metric]['accuracy'];axes[0].text(i+offset,v+1,f'{v:.1f}',ha='center',fontsize=9)
for k,color,marker in zip(labels,[BLUE,TEAL,RED,'#6957A0'],['o','s','^','D']):
    axes[1].plot([v['epoch'] for v in runs[k]['history']],[v['val_loss'] for v in runs[k]['history']],marker=marker,color=color,label=k.replace('shift augmentation','aug').replace(' / no augmentation',''))
axes[1].set(xlabel='Epoch',ylabel='Validation cross-entropy',title='Clean validation selects each checkpoint');axes[1].legend(fontsize=8)
fig.tight_layout();savefig(fig,'lab02-controlled-2x2');plt.show()
shift_acc=[100*runs[k]['right_2px']['accuracy'] for k in labels]
interaction_pp=(shift_acc[3]-shift_acc[2])-(shift_acc[1]-shift_acc[0])
display(Markdown(f'本次运行：MLP 的增强收益为 **{shift_acc[1]-shift_acc[0]:+.2f} pp**，CNN 为 **{shift_acc[3]-shift_acc[2]:+.2f} pp**；差的差（CNN − MLP）为 **{interaction_pp:+.2f} pp**。这是单 seed 的观测，不附会统计显著性。'))

label='CNN / no augmentation';p0,p1=predictions[label]
ids=torch.where((p0==test_y)&(p1!=test_y))[0][:6]
selection='first six test indices where clean is correct and shifted is wrong'
if not len(ids):ids=torch.arange(6);selection='no clean-correct/shift-wrong case; first six test indices'
fig,axes=plt.subplots(2,len(ids),figsize=(max(6,len(ids)*1.6),3.7),squeeze=False)
for j,i in enumerate(ids.tolist()):
    axes[0,j].imshow(test_x[i,0],cmap='gray');axes[0,j].set_title(f'id {i}\ny={int(test_y[i])}, p={int(p0[i])}',fontsize=10)
    axes[1,j].imshow(shifted_test[i,0],cmap='gray');axes[1,j].set_title(f'Shifted: p={int(p1[i])}',fontsize=10)
    for ax in axes[:,j]:ax.axis('off')
fig.tight_layout();savefig(fig,'lab02-paired-failures');plt.show()

chosen=int(ids[0]);m=trained_models[label]
with torch.no_grad():a0=F.relu(m.conv1(test_x[chosen:chosen+1]))[0];a1=F.relu(m.conv1(shifted_test[chosen:chosen+1]))[0]
fig,axes=plt.subplots(2,8,figsize=(13,3.8));limit=float(torch.cat([a0,a1]).max())
for c in range(8):
    for row,a in enumerate([a0,a1]):
        axes[row,c].imshow(a[c],cmap='magma',vmin=0,vmax=limit);axes[row,c].axis('off');axes[row,c].set_title(f'Channel {c}')
fig.suptitle(f'Actual first-layer ReLU activations: test index {chosen}; top clean, bottom right 2px')
savefig(fig,'lab02-actual-activations');plt.show()

# 只保存本轮训练时登记的身份；不替换或补签未知文件。
if 'checkpoint_provenance' not in globals() or not checkpoint_provenance:
    raise RuntimeError('请先运行第4节训练单元，再保存本轮结果与模型身份。')
for checkpoint_name,identity in checkpoint_provenance.items():
    path=OUT/checkpoint_name
    if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest()!=identity['sha256']:
        raise RuntimeError('模型文件与本轮训练登记不一致：'+str(path)+'；请重新运行第4节，不能给混用文件补签。')
    saved_state=torch.load(path,map_location='cpu',weights_only=True)
    live_state=trained_models[identity['run_label']].state_dict()
    if set(saved_state)!=set(live_state) or any(not torch.equal(saved_state[k],live_state[k].cpu()) for k in saved_state):
        raise RuntimeError('内存模型与本轮保存文件不一致，请重新运行第4节。')
assert SPLIT_SEED==training_source_protocol['split_seed']
assert TRAIN_N==training_source_protocol['train_n']
assert hashlib.sha256(train_ids.tobytes()).hexdigest()==training_source_protocol['train_ids_sha256']
assert hashlib.sha256(val_ids.tobytes()).hexdigest()==training_source_protocol['validation_ids_sha256']

summary={'experiment':'Controlled 2x2: MLP/CNN x no augmentation/random integer translation',
 'execution':'all four runs executed in the current Python environment on CPU; single seed; no synthetic fallback',
 'environment':{'python':platform.python_version(),'torch':torch.__version__,'numpy':np.__version__,'torch_threads':torch.get_num_threads(),'deterministic_algorithms':torch.are_deterministic_algorithms_enabled()},
 'protocol':{'train_n':TRAIN_N,'validation_n':len(val_y),'test_n':len(test_y),'split_seed':SPLIT_SEED,
             'model_seed':MODEL_SEED,'shuffle_seed':MODEL_SEED+1,'augmentation_seed':MODEL_SEED+2,
             'epochs':EPOCHS,'batch_size':BATCH_SIZE,'optimizer':'Adam','learning_rate':LEARNING_RATE,
             'selection':'minimum clean validation cross-entropy; test not used',
             'train_augmentation':'independent per-image dx,dy uniform integers [-2,2], zero-fill/crop',
             'test_transform':'fixed right 2 px, zero-fill left and crop right',
             'split':f'NumPy permutation; first {TRAIN_N} train, positions 50000:60000 validation; official test',
             'partition_sha256':hashlib.sha256(partition.tobytes()).hexdigest(),
             'train_ids_sha256':hashlib.sha256(train_ids.tobytes()).hexdigest(),'files':data_identity},
 'models':model_info,'runs':runs,'shift_augmentation_interaction_pp':interaction_pp,
 'failure_examples':{'selection':selection,'test_indices':ids.tolist(),'activation_example':chosen},
 'limitations':['single seed','architectures have different parameters and compute','MNIST only','cropping may remove information']}
summary['checkpoint_provenance']=copy.deepcopy(checkpoint_provenance)
(OUT/'lab02-results.json').write_text(json.dumps(summary,indent=2))
print('Saved results/lab02-results.json and four selected model states.')

import torch
from torch import nn
from collections import OrderedDict

torch.set_num_threads(2)
INPUT_SIDE, FIRST_PADDING = 28, 2  # 再试：28, 0；或 32, 2
assert INPUT_SIDE in (28, 32) and FIRST_PADDING in (0, 2)
probe = torch.zeros(1, 1, INPUT_SIDE, INPUT_SIDE)
with torch.random.fork_rng(devices=[]):
    torch.manual_seed(26)
    old_features = nn.Sequential(OrderedDict([
        ('conv1', nn.Conv2d(1, 6, 5, padding=FIRST_PADDING)), ('sigmoid1', nn.Sigmoid()),
        ('avgpool1', nn.AvgPool2d(2)), ('conv2', nn.Conv2d(6, 16, 5)),
        ('sigmoid2', nn.Sigmoid()), ('avgpool2', nn.AvgPool2d(2)),
    ]))
    current_features = nn.Sequential(OrderedDict([
        ('conv1', nn.Conv2d(1, 8, 3, padding=1)), ('relu1', nn.ReLU()),
        ('maxpool1', nn.MaxPool2d(2)), ('conv2', nn.Conv2d(8, 16, 3, padding=1)),
        ('relu2', nn.ReLU()), ('maxpool2', nn.MaxPool2d(2)),
    ]))
    with torch.no_grad():
        old_width = old_features(probe).numel()
        current_width = current_features(probe).numel()
    old = nn.Sequential(OrderedDict(list(old_features.named_children()) + [
        ('flatten', nn.Flatten()), ('fc1', nn.Linear(old_width, 120)), ('sigmoid3', nn.Sigmoid()),
        ('fc2', nn.Linear(120, 84)), ('sigmoid4', nn.Sigmoid()), ('logits', nn.Linear(84, 10)),
    ]))
    current = nn.Sequential(OrderedDict(list(current_features.named_children()) + [
        ('flatten', nn.Flatten()), ('hidden', nn.Linear(current_width, 32)),
        ('relu3', nn.ReLU()), ('logits', nn.Linear(32, 10)),
    ]))

HISTORY_RESULT = {}
for name, model, width in [('Historical teaching LeNet', old, old_width), ('Current LAB02 CNN', current, current_width)]:
    value, rows = probe, []
    print('\n' + name)
    print(f"{'layer':<14} {'output shape':<22} {'parameters':>10}")
    with torch.no_grad():
        for layer_name, layer in model.named_children():
            value = layer(value)
            count = sum(p.numel() for p in layer.parameters())
            rows.append({'layer': layer_name, 'shape': list(value.shape), 'parameters': count})
            print(f'{layer_name:<14} {str(tuple(value.shape)):<22} {count:>10,}')
    assert tuple(value.shape) == (1, 10)
    total = sum(p.numel() for p in model.parameters())
    print('Flatten width:', width, '| total parameters:', total)
    HISTORY_RESULT[name] = {'flatten_width': width, 'parameters': total, 'layers': rows}
if (INPUT_SIDE, FIRST_PADDING) == (28, 2):
    assert (old_width, current_width) == (400, 784)
    assert HISTORY_RESULT['Historical teaching LeNet']['parameters'] == 61706
    assert HISTORY_RESULT['Current LAB02 CNN']['parameters'] == 26698

import torch
from torch.nn import functional as F

TRUE_CLASS, SCALE, COMMON_SHIFT = 0, 1.0, 1000.0  # 再试：2, 3.0, -1000.0
assert TRUE_CLASS in (0, 1, 2) and 0 < SCALE <= 10
z = (torch.tensor([[2., 1., 0.]], dtype=torch.float64) * SCALE).requires_grad_()
y = torch.tensor([TRUE_CLASS], dtype=torch.long)
loss = F.cross_entropy(z, y)
grad, = torch.autograd.grad(loss, z)
p = z.detach().softmax(dim=1)
manual = p - F.one_hot(y, num_classes=3).to(p.dtype)
torch.testing.assert_close(grad, manual)

shifted = (z.detach() + COMMON_SHIFT).requires_grad_()
shifted_loss = F.cross_entropy(shifted, y)
shifted_grad, = torch.autograd.grad(shifted_loss, shifted)
torch.testing.assert_close(shifted.detach().softmax(1), p)
torch.testing.assert_close(shifted_loss, loss)
torch.testing.assert_close(shifted_grad, grad)

wrong_z = z.detach().clone().requires_grad_()
wrong_loss = F.cross_entropy(wrong_z.softmax(dim=1), y)
wrong_grad, = torch.autograd.grad(wrong_loss, wrong_z)
print(f"{'class':<7} {'logit':>8} {'p':>10} {'p-target':>12} {'autograd':>12} {'double-softmax grad':>20}")
for c in range(3):
    print(f'{c:<7} {z.detach()[0,c].item():8.3f} {p[0,c].item():10.6f} '
          f'{manual[0,c].item():12.6f} {grad[0,c].item():12.6f} {wrong_grad[0,c].item():20.6f}')
HISTORY_RESULT = {
    'true_class': TRUE_CLASS,
    'correct_ce': float(loss.detach()),
    'double_softmax_ce': float(wrong_loss.detach()),
    'manual_grad_max_error': float((grad-manual).abs().max()),
    'shift_loss_gap': float((shifted_loss.detach()-loss.detach()).abs()),
    'shift_grad_max_gap': float((shifted_grad-grad).abs().max()),
    'double_softmax_grad_max_gap': float((wrong_grad-grad).abs().max()),
    'gradient_sum': float(grad.sum()),
}
print('Checks:', HISTORY_RESULT)