LAB 10 · 成本目标与结构选择

LAB_ID='LAB12-C'

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

WARMUP,REPEATS,ROUNDS=30,200,3
rng=np.random.default_rng(SEED)
def modules_for(m,n):
    standard=nn.Conv2d(m,n,3,padding=1,bias=False).eval()
    decomposed=nn.Sequential(nn.Conv2d(m,m,3,padding=1,groups=m,bias=False),nn.Conv2d(m,n,1,bias=False)).eval()
    return {'standard':standard,'dw_pw':decomposed}
# Exact constructed factorable kernel, not a claim about arbitrary standard kernels.
small=modules_for(2,3); sx=torch.randn(1,2,4,4)
with torch.no_grad():
    dw,pw=small['dw_pw'][0].weight,small['dw_pw'][1].weight
    small['standard'].weight.copy_(pw[:,:,0,0][:,:,None,None]*dw[:,0][None,:,:,:])
    max_error=(small['standard'](sx)-small['dw_pw'](sx)).abs().max().item()
assert max_error < 1e-5
print('Factorable-kernel numerical check; max absolute error =',max_error)

def benchmark_group(modules,x):
    samples={name:[] for name in modules}
    with torch.inference_mode():
        for module in modules.values():
            for _ in range(WARMUP): module(x)
        for run in range(ROUNDS):
            for name in rng.permutation(list(modules)):
                module=modules[name]
                for repeat in range(REPEATS):
                    start=time.perf_counter_ns(); module(x); duration=time.perf_counter_ns()-start
                    samples[name].append({'round':run,'repeat':repeat,'latency_us':duration/1000})
    return samples
operator_rows=[]; raw_operator=[]
for case,m,n,h in [('large',32,64,32),('small',4,16,4)]:
    x=torch.randn(1,m,h,h); modules=modules_for(m,n); samples=benchmark_group(modules,x)
    for name,module in modules.items():
        values=[r['latency_us'] for r in samples[name]]
        params=sum(p.numel() for p in module.parameters())
        row={'case':case,'module':name,'input_shape':str(list(x.shape)),'output_shape':str([1,n,h,h]),
             'params':params,'mac':params*h*h,'activation_ledger_bytes':4*h*h*(m+n+(m if name=='dw_pw' else 0)),
             'p10_us':float(np.percentile(values,10)),'median_us':float(np.median(values)),'p90_us':float(np.percentile(values,90))}
        operator_rows.append(row)
        for r in samples[name]: raw_operator.append({'case':case,'module':name,**r})
save_csv('operator-costs.csv',operator_rows); save_csv('operator-timings.csv',raw_operator)
save_json('benchmark-protocol.json',{'warmup':WARMUP,'repetitions':REPEATS,'rounds':ROUNDS,'layout':'NCHW',
  'scope':'Python call, operator and output allocation; inputs and weights precreated','synchronization':'CPU synchronous',
  'activation_ledger':'input+output+depthwise intermediate, NOT peak-memory measurement'})
print_table(operator_rows,['case','module','params','mac','activation_ledger_bytes','median_us','p10_us','p90_us'])

fig,axs=plt.subplots(1,2,figsize=(10,3.4))
for ax,case in zip(axs,['large','small']):
    vals=[[r['latency_us'] for r in raw_operator if r['case']==case and r['module']==name] for name in ['standard','dw_pw']]
    ax.boxplot(vals,tick_labels=['standard','depthwise +\npointwise'],showfliers=False)
    ax.set(ylabel='CPU latency (microseconds)',title=case+' shape; 600 measured calls each')
fig.suptitle('Less arithmetic need not imply less elapsed time',fontsize=12)
fig.tight_layout(); fig.savefig(OUTPUT/'operator-latencies.png'); plt.show()
for case in ['large','small']:
    a,b=[r for r in operator_rows if r['case']==case]
    print(case,': MAC ratio DW/PW / standard =',round(b['mac']/a['mac'],3),
          '; measured median ratio =',round(b['median_us']/a['median_us'],3))

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

class Candidate(nn.Module):
    def __init__(self,kind,width):
        super().__init__(); self.stem=nn.Conv2d(1,8,3,padding=1,bias=False)
        self.module=(nn.Conv2d(8,width,3,padding=1,bias=False) if kind=='standard' else
                     nn.Sequential(nn.Conv2d(8,8,3,padding=1,groups=8,bias=False),nn.Conv2d(8,width,1,bias=False)))
        self.head=nn.Linear(width*2*2,10)
    def forward(self,x):
        x=F.avg_pool2d(F.relu(self.stem(x)),2)
        x=F.avg_pool2d(F.relu(self.module(x)),2)
        return self.head(x.flatten(1))
models={}; search_rows=[]
for kind in ['standard','dw_pw']:
    for width in [8,16,32]:
        torch.manual_seed(SEED); model=Candidate(kind,width); opt=torch.optim.Adam(model.parameters(),lr=.003)
        generator=torch.Generator().manual_seed(SEED+10)
        best_loss=float('inf'); best=None
        for epoch in range(18):
            model.train()
            ids=torch.tensor(train_ids)[torch.randperm(len(train_ids),generator=generator)]
            for batch in ids.split(64):
                opt.zero_grad(set_to_none=True); loss=F.cross_entropy(model(X[batch]),y[batch]); loss.backward(); opt.step()
            model.eval()
            with torch.no_grad(): vl=F.cross_entropy(model(X[val_ids]),y[val_ids]).item()
            if vl<best_loss: best_loss=vl; best=deepcopy(model.state_dict())
        model.load_state_dict(best); model.eval(); name=f'{kind}-{width}'; models[name]=model
        with torch.inference_mode(): va=(model(X[val_ids]).argmax(1)==y[val_ids]).float().mean().item()
        params=sum(p.numel() for p in model.parameters())
        mac=64*9*8 + (16*9*8*width if kind=='standard' else 16*(9*8+8*width)) + width*4*10
        search_rows.append({'name':name,'validation_acc':va,'params':params,'mac':mac})
search_samples=benchmark_group(models,X[test_ids[:1]])  # shape-only input; labels unused
for row in search_rows:
    values=[s['latency_us'] for s in search_samples[row['name']]]
    row.update({'median_us':float(np.median(values)),'p10_us':float(np.percentile(values,10)),
                'p90_us':float(np.percentile(values,90))})
save_csv('candidate-latency-samples.csv',[{'name':name,**sample} for name,ss in search_samples.items() for sample in ss])
base=next(r for r in search_rows if r['name']=='standard-16')
feasible=lambda r:r['validation_acc']>=base['validation_acc']-1e-9 and r['median_us']<=base['median_us']
accuracy_best=max(search_rows,key=lambda r:(r['validation_acc'],-r['median_us']))
cost_best=min([r for r in search_rows if feasible(r)],key=lambda r:r['median_us'])
random_names=np.random.default_rng(SEED+33).choice([r['name'] for r in search_rows],size=2,replace=False).tolist()
random_pool=[r for r in search_rows if r['name'] in random_names and feasible(r)]
random_best=min(random_pool,key=lambda r:r['median_us']) if random_pool else base
selections={'manual_baseline':base['name'],'validation_accuracy':accuracy_best['name'],
            'latency_under_baseline_quality_and_cost':cost_best['name'],
            'random_two_under_same_constraints':random_best['name'],'random_two_names':random_names,
            'random_fallback_to_baseline':not bool(random_pool)}
# Selection has finished. Only now calculate held-out test accuracy for all candidates, including nonwinners.
for row in search_rows:
    with torch.inference_mode(): row['test_acc']=(models[row['name']](X[test_ids]).argmax(1)==y[test_ids]).float().mean().item()
    row['pareto_validation_latency']=not any(other['validation_acc']>=row['validation_acc'] and other['median_us']<=row['median_us']
        and (other['validation_acc']>row['validation_acc'] or other['median_us']<row['median_us']) for other in search_rows)
save_csv('candidate-results.csv',search_rows); save_json('selection.json',selections)
print_table(search_rows); print(json.dumps(selections,indent=2))

fig,ax=plt.subplots(figsize=(8,4))
for r in search_rows:
    ax.scatter(r['median_us'],100*r['validation_acc'],s=65,marker='o' if r['name'].startswith('standard') else 's',
               facecolors='#255b87' if r['pareto_validation_latency'] else 'none',edgecolors='#255b87')
    ax.annotate(r['name'],(r['median_us'],100*r['validation_acc']),xytext=(5,5),textcoords='offset points',fontsize=9)
ax.set(xlabel='Measured batch-1 CPU median latency (microseconds)',ylabel='Validation accuracy (%)',
       title='Six trained candidates; filled points are Pareto choices')
ax.margins(x=.23,y=.25); fig.tight_layout(); fig.savefig(OUTPUT/'candidate-pareto.png'); plt.show()