LAB 12 · Source、Norm 与 Tent

LAB_ID='LAB14-E'

from pathlib import Path
from copy import deepcopy
import csv, hashlib, json, platform, sys, time
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import sklearn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
SEED = 20260918
np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(1)
torch.use_deterministic_algorithms(True)
OUTPUT = Path('outputs') / LAB_ID
OUTPUT.mkdir(parents=True, exist_ok=True)
plt.rcParams.update({'figure.dpi':120, 'font.size':10, 'axes.spines.top':False, 'axes.spines.right':False})
def save_json(name, obj):
    (OUTPUT/name).write_text(json.dumps(obj,ensure_ascii=False,indent=2))
def save_csv(name, rows):
    with (OUTPUT/name).open('w',newline='') as f:
        w=csv.DictWriter(f, fieldnames=list(rows[0])); w.writeheader(); w.writerows(rows)
def digest_array(a): return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()
def print_table(rows, cols=None):
    cols=cols or list(rows[0]); print(' | '.join(cols)); print(' | '.join(['---']*len(cols)))
    for row in rows:
        print(' | '.join(f'{row[c]:.5f}' if isinstance(row[c],float) else str(row[c]) for c in cols))
cpu_lines=Path('/proc/cpuinfo').read_text().splitlines() if Path('/proc/cpuinfo').exists() else []
cpu_name=next((line.split(':',1)[1].strip() for line in cpu_lines if line.startswith('model name')),platform.processor())
environment={'python':sys.version.split()[0],'torch':torch.__version__,'numpy':np.__version__,
             'sklearn':sklearn.__version__,'cpu':cpu_name,'platform':platform.platform(),
             'torch_threads':torch.get_num_threads(),'device':'cpu','dtype':'float32',
             'seed':SEED,'exclusive_hardware':False}
save_json('environment.json',environment)
print(json.dumps(environment,ensure_ascii=False,indent=2))

dataset = load_digits()
X = torch.tensor(dataset.images[:,None,:,:] / 16.0, dtype=torch.float32)
y = torch.tensor(dataset.target, dtype=torch.long)
all_ids = np.arange(len(y))
trainval_ids, test_ids = train_test_split(all_ids, test_size=0.2, random_state=SEED, stratify=y.numpy())
train_ids, val_ids = train_test_split(trainval_ids, test_size=0.25, random_state=SEED, stratify=y.numpy()[trainval_ids])
assert not (set(train_ids)&set(val_ids) or set(train_ids)&set(test_ids) or set(val_ids)&set(test_ids))
assert len(train_ids)+len(val_ids)+len(test_ids)==len(y)
split_manifest={'dataset':'sklearn.datasets.load_digits; 1797 real 8x8 images; pixel/16',
                'counts':{'train':len(train_ids),'validation':len(val_ids),'test':len(test_ids)},
                'train_ids':train_ids.tolist(),'validation_ids':val_ids.tolist(),'test_ids':test_ids.tolist(),
                'input_sha256':digest_array(X.numpy()),'labels_sha256':digest_array(y.numpy()),
                'split_policy':'stratified random by sample, NOT writer-disjoint', 'test_labels':'evaluation only'}
save_json('split-manifest.json',split_manifest)
print('Split:',split_manifest['counts'])

MODEL_ID = 'DigitsMultiExitBN-integration-20260918-lcde-r1'
class DigitsMultiExitBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(1,4,3,padding=1,bias=False)
        self.bn1=nn.BatchNorm2d(4)
        self.early=nn.Linear(4*4*4,10)
        self.conv2=nn.Conv2d(4,16,3,padding=1,bias=False)
        self.bn2=nn.BatchNorm2d(16)
        self.fc=nn.Linear(16*4*4,32)
        self.final=nn.Linear(32,10)
    def forward(self,x):
        h=F.avg_pool2d(F.relu(self.bn1(self.conv1(x))),2)
        early=self.early(h.flatten(1))
        h=F.relu(self.bn2(self.conv2(h)))
        final=self.final(F.relu(self.fc(h.flatten(1))))
        return early,final

def train_source():
    torch.manual_seed(SEED)
    model=DigitsMultiExitBN()
    optimizer=torch.optim.Adam(model.parameters(),lr=0.003)
    order_generator=torch.Generator().manual_seed(SEED+1)
    best_loss=float('inf'); best_state=None; best_epoch=None; history=[]
    for epoch in range(1,51):
        model.train()
        order=torch.as_tensor(train_ids)[torch.randperm(len(train_ids),generator=order_generator)]
        total=0.0
        for ids in order.split(64):
            optimizer.zero_grad(set_to_none=True)
            early,final=model(X[ids]); loss=0.35*F.cross_entropy(early,y[ids])+F.cross_entropy(final,y[ids])
            loss.backward(); optimizer.step(); total += loss.item()*len(ids)
        model.eval()
        with torch.no_grad():
            ve,vf=model(X[val_ids]); val_loss=F.cross_entropy(vf,y[val_ids]).item()
        history.append({'epoch':epoch,'train_weighted_loss':total/len(train_ids),
                        'validation_final_ce':val_loss,'validation_early_acc':(ve.argmax(1)==y[val_ids]).float().mean().item(),
                        'validation_final_acc':(vf.argmax(1)==y[val_ids]).float().mean().item()})
        if val_loss < best_loss:
            best_loss=val_loss; best_state=deepcopy(model.state_dict()); best_epoch=epoch
    model.load_state_dict(best_state); model.eval()
    return model,history,best_epoch
source_model,train_history,best_epoch=train_source()
source_state=deepcopy(source_model.state_dict())
state_hash=hashlib.sha256(b''.join(k.encode()+v.cpu().numpy().tobytes() for k,v in sorted(source_state.items()))).hexdigest()
save_csv('source-training.csv',train_history)
# Transparent JSON state is a local convenience; notebook always reconstructs from readable code.
save_json('source-state.json',{k:v.tolist() for k,v in source_state.items()})
model_manifest={'model_id':MODEL_ID,'state_sha256':state_hash,'best_epoch':best_epoch,
                'checkpoint_selection':'minimum validation final-exit CE over 50 epochs',
                'train_loss':'0.35*CE(early)+CE(final)','optimizer':'Adam lr=0.003 batch=64',
                'initialization_seed':SEED,'params':sum(p.numel() for p in source_model.parameters()),
                'old_LAB02_dependency':False,'data_split_sha256':hashlib.sha256(json.dumps(split_manifest,sort_keys=True).encode()).hexdigest()}
save_json('model-manifest.json',model_manifest)
print(json.dumps(model_manifest,indent=2))

BATCH=32; TTA_LR=.01
sample_noise=np.stack([np.random.default_rng(SEED+int(s)*17).normal(0,.30,(1,8,8)) for s in test_ids])
target_inputs={'contrast':torch.clamp(.55*X[test_ids]+.25,0,1),
               'noise':torch.clamp(X[test_ids]+torch.tensor(sample_noise,dtype=torch.float32),0,1)}
orders={'shuffled':np.random.default_rng(SEED+77).permutation(len(test_ids)),
        'class_sorted':np.argsort(y[test_ids].numpy(),kind='stable')}
save_json('tta-protocol.json',{'batch':BATCH,'learning_rate':TTA_LR,'steps_per_batch':1,'optimizer':'SGD; momentum=0',
 'corruptions':{'contrast':'clip(0.55*x+0.25,0,1)','noise':'clip(x+N(0,0.30),0,1); seed=SEED+sample_id*17'},
 'prediction_timing':'logits from forward before this batch optimizer.step()',
 'reset':'same source checkpoint for each scenario and arm; episodic resets each batch',
 'label_use':'only post-prediction evaluation; class_sorted stress order constructed offline with labels',
 'target_input_hashes':{k:digest_array(v.numpy()) for k,v in target_inputs.items()},
 'orders':{k:v.tolist() for k,v in orders.items()},'test_hyperparameter_tuning':False})
fig,axs=plt.subplots(3,6,figsize=(9,4.3))
for row,(name,inputs) in enumerate([('clean',X[test_ids]),*target_inputs.items()]):
    for j in range(6):
        axs[row,j].imshow(inputs[j,0],cmap='gray',vmin=0,vmax=1);axs[row,j].set_xticks([]);axs[row,j].set_yticks([])
        if row==0:axs[row,j].set_title(f'id={test_ids[j]} y={int(y[test_ids[j]])}',fontsize=9)
    axs[row,0].set_ylabel(name)
fig.suptitle('Same inputs; fixed display range [0,1]');fig.tight_layout();fig.savefig(OUTPUT/'tta-inputs.png');plt.show()

def entropy(logits): return -(logits.softmax(1)*logits.log_softmax(1)).sum(1)
def new_arm(arm):
    model=DigitsMultiExitBN();model.load_state_dict(source_state);model.requires_grad_(False)
    if arm=='source': model.eval(); return model,None
    model.train()
    affine=[]
    for module in model.modules():
        if isinstance(module,nn.BatchNorm2d):
            module.track_running_stats=False;module.running_mean=None;module.running_var=None
            if arm=='tent':
                module.weight.requires_grad_(True);module.bias.requires_grad_(True)
                affine += [module.weight,module.bias]
    opt=torch.optim.SGD(affine,lr=TTA_LR,momentum=0) if arm=='tent' else None
    return model,opt

def adapt_stream(images,order,arm,episodic=False):
    # No labels argument. Source state read-only; each invocation starts independently.
    model,opt=new_arm(arm)
    initial_parameters={n:p.detach().clone() for n,p in model.named_parameters()}
    initial_buffers={n:b.detach().clone() for n,b in model.named_buffers()}
    logits=np.zeros((len(images),10),dtype=np.float32); batch_records=[]
    allowed={'bn1.weight','bn1.bias','bn2.weight','bn2.bias'} if arm=='tent' else set()
    for batch_index,start in enumerate(range(0,len(order),BATCH)):
        if episodic: model,opt=new_arm(arm)
        local=order[start:start+BATCH]; batch=images[local]
        if arm=='tent':
            opt.zero_grad(set_to_none=True)
            _,out=model(batch)
            loss=entropy(out).mean()
            # Save the prediction from before the current update; no extra forward.
            logits[local]=out.detach().numpy()
            loss.backward(); opt.step()
        else:
            with torch.no_grad(): _,out=model(batch)
            logits[local]=out.numpy();loss=entropy(out).mean()
        changed_this_batch=[n for n,p in model.named_parameters() if not torch.equal(p.detach(),initial_parameters[n])]
        assert set(changed_this_batch)<=allowed
        batch_records.append({'batch_index':batch_index,'n':len(local),'entropy_pre_update':float(loss.detach()),
            'affine_delta_l2':float(sum((p.detach()-initial_parameters[n]).square().sum() for n,p in model.named_parameters()).sqrt())})
    deltas={n:float((p.detach()-initial_parameters[n]).abs().max()) for n,p in model.named_parameters()}
    changed=[n for n,d in deltas.items() if d>0]
    buffer_changes=[n for n,b in model.named_buffers() if not torch.equal(b,initial_buffers[n])]
    assert set(changed)<=allowed and not buffer_changes
    if arm=='tent': assert changed and sum(p.numel() for p in model.parameters() if p.requires_grad)==40
    audit={'arm':arm,'episodic':episodic,'changed_parameters':changed,'max_absolute_parameter_deltas':deltas,
           'changed_buffers_during_stream':buffer_changes,
           'running_means_none':{n:m.running_mean is None for n,m in model.named_modules() if isinstance(m,nn.BatchNorm2d)},
           'num_batches_tracked':{n:int(m.num_batches_tracked) for n,m in model.named_modules() if isinstance(m,nn.BatchNorm2d)},
           'trainable_parameter_count':sum(p.numel() for p in model.parameters() if p.requires_grad)}
    return logits,batch_records,audit

# Compute every scenario before calling target labels for evaluation below.
runs={};audits={};batch_rows=[]
for corruption,images in target_inputs.items():
    for order_name,order in orders.items():
        for arm,episodic in [('source',False),('norm',False),('tent',False),('tent',True)]:
            name=f'{corruption}/{order_name}/{arm}'+('-episodic' if episodic else '')
            result,batches,audit=adapt_stream(images,order,arm,episodic)
            runs[name]=result;audits[name]=audit
            batch_rows += [{'run':name,**row} for row in batches]
# At the first batch all affine values coincide: Tent prediction must equal Norm prediction.
for corruption in target_inputs:
    for order_name,order in orders.items():
        idx=order[:BATCH]
        assert np.allclose(runs[f'{corruption}/{order_name}/tent'][idx],runs[f'{corruption}/{order_name}/norm'][idx],atol=1e-6)
        # With per-batch reset and before-update prediction, episodic Tent outputs equal Norm for every batch.
        assert np.allclose(runs[f'{corruption}/{order_name}/tent-episodic'],runs[f'{corruption}/{order_name}/norm'],atol=1e-6)
save_json('parameter-buffer-audit.json',audits);save_csv('batch-trajectories.csv',batch_rows)
print('Passed: only allowed BN affine changed; no running buffers updated; first-batch predictions match Norm;')
print('episodic + one step + before-update predictions match Norm across the whole stream.')

truth=y[test_ids].numpy();summary=[];prediction_rows=[]
for name,logits in runs.items():
    probs=torch.tensor(logits).softmax(1).numpy();pred=probs.argmax(1)
    ents=entropy(torch.tensor(logits)).numpy(); wrong=pred!=truth
    summary.append({'run':name,'accuracy':float(np.mean(~wrong)),'errors':int(wrong.sum()),
                    'mean_entropy':float(ents.mean()),'wrong_confident_p_ge_09':int(np.sum(wrong&(probs.max(1)>=.9)))})
    for i,sample_id in enumerate(test_ids):
        prediction_rows.append({'run':name,'sample_id':int(sample_id),'target':int(truth[i]),'prediction':int(pred[i]),
              'confidence':float(probs[i].max()),'entropy':float(ents[i]),**{f'logit_{k}':float(logits[i,k]) for k in range(10)}})
save_csv('tta-summary.csv',summary);save_csv('tta-predictions.csv',prediction_rows)
print_table(summary)
fig,axs=plt.subplots(1,2,figsize=(11,4))
labels=[c+'/'+o for c in target_inputs for o in orders]
xx=np.arange(len(labels));colors=['#52616b','#2f80a0','#d07b27']
for i,arm in enumerate(['source','norm','tent']):
    selected=[next(r for r in summary if r['run']==f'{label}/{arm}') for label in labels]
    axs[0].bar(xx+(i-1)*.24,[100*(1-r['accuracy']) for r in selected],width=.24,label=arm,color=colors[i])
    axs[1].bar(xx+(i-1)*.24,[r['mean_entropy'] for r in selected],width=.24,label=arm,color=colors[i])
for ax in axs: ax.set_xticks(xx,labels,rotation=18); ax.legend(fontsize=8)
axs[0].set(ylabel='Error rate (%)',title='Target labels used only for evaluation')
axs[1].set(ylabel='Mean prediction entropy (nats)',title='Proxy objective is not correctness')
fig.tight_layout();fig.savefig(OUTPUT/'tta-comparison.png');plt.show()
comparisons=[]
for label in labels:
    s,n,t=[next(r for r in summary if r['run']==f'{label}/{arm}') for arm in ['source','norm','tent']]
    comparisons.append({'scenario':label,'tent_minus_source_error_pp':100*(s['accuracy']-t['accuracy']),
                        'tent_minus_norm_error_pp':100*(n['accuracy']-t['accuracy']),
                        'tent_minus_norm_entropy':t['mean_entropy']-n['mean_entropy']})
save_csv('tta-differences.csv',comparisons);print_table(comparisons)

# Inspect actual model mistakes. Prefer cases where Tent is wrong with high confidence.
worst_name=max((n for n in runs if n.endswith('/tent')),key=lambda n:int((runs[n].argmax(1)!=truth).sum()))
logits=runs[worst_name];p=torch.tensor(logits).softmax(1).numpy();pred=p.argmax(1)
wrong=np.flatnonzero(pred!=truth); chosen=wrong[np.argsort(-p[wrong].max(1))[:6]]
corruption=worst_name.split('/')[0]
fig,axs=plt.subplots(1,6,figsize=(10,2.4))
for ax,i in zip(axs,chosen):
    ax.imshow(target_inputs[corruption][i,0],cmap='gray',vmin=0,vmax=1);ax.axis('off')
    ax.set_title(f'id={test_ids[i]} y={truth[i]}\npred={pred[i]}, p={p[i].max():.2f}',fontsize=9)
fig.suptitle('Descriptive failure gallery: '+worst_name,fontsize=11)
fig.tight_layout();fig.savefig(OUTPUT/'tta-failures.png');plt.show()
print('Gallery chosen after evaluation to describe failure, not to select algorithm or report a best run.')
# Logical counterexample; separate from empirical frequency above.
wrong_logits=torch.tensor([[0.,2.,0.]])  # true class 0, predicted class 1
counterexample=[]
for scale in [1.,2.,5.]:
    z=scale*wrong_logits
    counterexample.append({'logit_scale':scale,'prediction':int(z.argmax(1)),
                          'entropy':float(entropy(z)),'correct':bool(z.argmax(1)==0)})
print_table(counterexample);save_csv('exact-entropy-counterexample.csv',counterexample)