LAB 11 · 何时继续计算

LAB_ID='LAB13-D'

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

fig,ax=plt.subplots(figsize=(7.5,3))
ax.plot([r['epoch'] for r in train_history],[r['validation_early_acc'] for r in train_history],label='early')
ax.plot([r['epoch'] for r in train_history],[r['validation_final_acc'] for r in train_history],label='final')
ax.axvline(best_epoch,color='gray',linestyle='--',label='selected by final CE')
ax.set(xlabel='Source training epoch',ylabel='Validation accuracy',ylim=(0,1)); ax.legend()
fig.tight_layout(); fig.savefig(OUTPUT/'training.png'); plt.show()
with torch.inference_mode():
    val_early,val_final=source_model(X[val_ids]); test_early,test_final=source_model(X[test_ids])
for split,ids,e,f in [('validation',val_ids,val_early,val_final),('test',test_ids,test_early,test_final)]:
    rows=[]
    for i,sample_id in enumerate(ids):
        row={'sample_id':int(sample_id),'target':int(y[sample_id])}
        row.update({f'early_logit_{k}':float(e[i,k]) for k in range(10)})
        row.update({f'final_logit_{k}':float(f[i,k]) for k in range(10)})
        rows.append(row)
    save_csv(f'{split}-exit-logits.csv',rows)
print('Both exits computed and saved; no runtime early exit implemented.')

SHARED,EARLY_HEAD,TAIL=2304,640,17728
EARLY_MAC=SHARED+EARLY_HEAD
DYNAMIC_LATE_MAC=SHARED+EARLY_HEAD+TAIL
FULL_MAC=SHARED+TAIL
THRESHOLDS=np.r_[np.arange(.50,1.0,.01),1.01]
def policy(early,final,labels,threshold):
    stopped=early.softmax(1).max(1).values.numpy()>=threshold
    ep,fp=early.argmax(1).numpy(),final.argmax(1).numpy()
    pred=np.where(stopped,ep,fp); truth=labels.numpy()
    return {'threshold':float(threshold),'accuracy':float(np.mean(pred==truth)),
            'early_fraction':float(stopped.mean()),'mean_conv_linear_mac':float(np.where(stopped,EARLY_MAC,DYNAMIC_LATE_MAC).mean()),
            'wrong_early_count':int(np.sum(stopped&(ep!=truth))),
            'avoidable_wrong_early_count':int(np.sum(stopped&(ep!=truth)&(fp==truth)))},stopped,pred
val_rows=[policy(val_early,val_final,y[val_ids],t)[0] for t in THRESHOLDS]
val_full_acc=float((val_final.argmax(1)==y[val_ids]).float().mean())
feasible=[r for r in val_rows if r['accuracy']>=val_full_acc-.01]
chosen=min(feasible,key=lambda r:(r['mean_conv_linear_mac'],-r['accuracy'],-r['threshold']))
chosen_t=chosen['threshold']; random_probability=chosen['early_fraction']
save_csv('validation-thresholds.csv',val_rows)
save_json('threshold-selection.json',{'chosen':chosen,'full_validation_acc':val_full_acc,'max_validation_accuracy_drop':.01,
              'test_labels_used':False,'mac_scope':'conv/linear only; includes early head for dynamic policy; gate omitted',
              'counts':{'shared':SHARED,'early_head':EARLY_HEAD,'tail':TAIL,'direct_full':FULL_MAC}})
print('Selected using validation only:',json.dumps(chosen,indent=2))

test_rows=[]
for t in THRESHOLDS: test_rows.append(policy(test_early,test_final,y[test_ids],t)[0])
chosen_test,stopped,pred=policy(test_early,test_final,y[test_ids],chosen_t)
ep=test_early.argmax(1).numpy(); fp=test_final.argmax(1).numpy(); truth=y[test_ids].numpy()
random_rows=[]
for seed in range(100):
    take_early=np.random.default_rng(SEED+100+seed).random(len(test_ids))<random_probability
    rp=np.where(take_early,ep,fp)
    random_rows.append({'seed':SEED+100+seed,'accuracy':float(np.mean(rp==truth)),'early_fraction':float(take_early.mean()),
                        'mean_conv_linear_mac':float(np.where(take_early,EARLY_MAC,FULL_MAC).mean())})
# Random routing is known before the early head; deep-routed samples skip early head.
summary=[{'policy':'full depth','accuracy':float(np.mean(fp==truth)),'early_fraction':0.0,'mean_conv_linear_mac':float(FULL_MAC)},
         {'policy':'fixed shallow','accuracy':float(np.mean(ep==truth)),'early_fraction':1.0,'mean_conv_linear_mac':float(EARLY_MAC)},
         {'policy':'selected threshold','accuracy':chosen_test['accuracy'],'early_fraction':chosen_test['early_fraction'],
          'mean_conv_linear_mac':chosen_test['mean_conv_linear_mac']},
         {'policy':'random routing mean(100 seeds)','accuracy':float(np.mean([r['accuracy'] for r in random_rows])),
          'early_fraction':float(np.mean([r['early_fraction'] for r in random_rows])),
          'mean_conv_linear_mac':float(np.mean([r['mean_conv_linear_mac'] for r in random_rows]))}]
save_csv('test-thresholds.csv',test_rows); save_csv('random-policies.csv',random_rows);save_csv('policy-summary.csv',summary)
rows=[{'sample_id':int(s),'target':int(truth[i]),'early_prediction':int(ep[i]),'final_prediction':int(fp[i]),
       'early_confidence':float(test_early.softmax(1)[i].max()),'stopped':bool(stopped[i]),'policy_prediction':int(pred[i])} for i,s in enumerate(test_ids)]
save_csv('selected-policy-per-sample.csv',rows)
print_table(summary)
print('Chosen policy wrong early exits:',chosen_test['wrong_early_count'],
      '; final exit would fix:',chosen_test['avoidable_wrong_early_count'])
print('Random accuracy SD over predefined seeds:',float(np.std([r['accuracy'] for r in random_rows],ddof=1)))
fig,ax=plt.subplots(figsize=(8,4))
ax.plot([r['mean_conv_linear_mac']/FULL_MAC for r in test_rows], [100*r['accuracy'] for r in test_rows],'-',label='threshold sweep; test shown after selection',alpha=.6)
for r,m in zip(summary,['o','s','*','D']):
    ax.scatter(r['mean_conv_linear_mac']/FULL_MAC,100*r['accuracy'],marker=m,s=110,label=r['policy'])
ax.set(xlabel='Conv/linear MAC relative to direct full depth',ylabel='Held-out test accuracy (%)',
       title='Offline exit decisions; no latency claim')
ax.legend(fontsize=8,loc='best');fig.tight_layout(); fig.savefig(OUTPUT/'accuracy-mac.png');plt.show()

wrong=np.flatnonzero(stopped&(ep!=truth))
# If none exist at selected threshold, show actual high-confidence shallow errors without pretending they stopped.
show_ids=wrong[:6]
plot_label='Wrong early exits at selected threshold'
if not len(show_ids):
    errors=np.flatnonzero(ep!=truth)
    show_ids=errors[np.argsort(-test_early.softmax(1).max(1).values.numpy()[errors])[:6]]
    plot_label='No wrong selected exits; these shallow errors did NOT necessarily stop'
fig,axs=plt.subplots(1,max(1,len(show_ids)),figsize=(max(5,1.7*len(show_ids)),3.0),squeeze=False)
for ax,j in zip(axs[0],show_ids):
    ax.imshow(X[test_ids[j],0],cmap='gray',vmin=0,vmax=1);ax.axis('off')
    ax.set_title(f'id={test_ids[j]} y={truth[j]}\nearly={ep[j]} p={float(test_early.softmax(1)[j].max()):.2f}\nfinal={fp[j]} stop={bool(stopped[j])}',fontsize=8)
fig.suptitle(plot_label,fontsize=10);fig.tight_layout(rect=[0,0,1,.84]);fig.savefig(OUTPUT/'wrong-exits.png');plt.show()