seq2seq.py · 可读全文
这里保留原文的中文、公式标记、代码和链接。可直接交给 Agent 读取，也可选取下方原文复制。此入口返回 HTML；原始 .md / .py 文件仍保留，完整候选包同时包含 UTF-8 原件。
"""DN1 controlled CPU experiment. Natural source order, reversed canonical target.
Both systems use the same BiLSTM encoder, zero decoder initial state and context injection.
Only the read interface differs; dynamic reading adds a trainable query projection.
"""
from pathlib import Path
import json, time, math, hashlib, platform, copy
import numpy as np
import torch
import experiment_io as runs
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
torch.set_num_threads(1)
PAD, BOS, EOS = 0, 1, 2
VOCAB = ['<PAD>','<BOS>','<EOS>'] + list('abcdef')
CONFIG = dict(seed=20261018, train=3000, validation=300, test=300,
source_lengths=[3,7], embed=16, encoder_hidden=16, decoder_hidden=32,
batch=32, updates=500, checkpoint_interval=100, lr=.003,
max_seconds_per_model=600, source_order='natural', target='reverse(source)+EOS',
initialization='copy all same-shaped common tensors; separate optimizers',
summary='masked mean of BiLSTM outputs', decoder_initial_state='zeros',
output_score='sum log p including EOS; PAD/BOS forbidden; no length penalty')
def validate_config(config, training=False):
"""Reject descriptive switches that have no implementation, before writing."""
if not isinstance(config, dict):
raise ValueError('配置必须是字典；请从 dict(seq.CONFIG) 开始修改。')
missing, unknown = set(CONFIG) - set(config), set(config) - set(CONFIG)
if missing or unknown:
raise ValueError(f'配置字段不匹配：缺少{sorted(missing)}；不支持{sorted(unknown)}。')
config = copy.deepcopy(config)
for key in ('source_order', 'target', 'initialization', 'summary', 'decoder_initial_state', 'output_score'):
if config[key] != CONFIG[key]:
raise ValueError(f'{key} 当前只支持 {CONFIG[key]!r}；该模式尚未实现，不能只改报告文字。')
lengths = config['source_lengths']
if not isinstance(lengths, (list, tuple)) or len(lengths) != 2 or any(type(n) is not int for n in lengths):
raise ValueError('source_lengths 应为 [最短长度, 最长长度] 两个整数。')
lo, hi = lengths
if not 1 <= lo <= hi <= 64:
raise ValueError('source_lengths 需满足 1 <= 最短 <= 最长 <= 64；本CPU教学实现不支持更长训练。')
config['source_lengths'] = list(lengths)
for key in ('seed', 'train', 'validation', 'test', 'embed', 'encoder_hidden', 'decoder_hidden', 'batch', 'updates', 'checkpoint_interval'):
if type(config[key]) is not int or config[key] < (0 if key in ('seed', 'train', 'validation', 'test') else 1):
raise ValueError(f'{key} 必须是合法的整数，不能为负数或布尔值。')
if config['seed'] > 2**32 - 1:
raise ValueError('seed 需在0至2**32-1之间。')
total = sum(config[k] for k in ('train', 'validation', 'test'))
capacity = sum((len(VOCAB) - 3)**n for n in range(lo, hi + 1))
if total < 1 or total > capacity:
raise ValueError(f'请求{total}条唯一序列，但长度{lo}–{hi}的六字符词表容量只有{capacity}；请减少样本数或扩大长度范围。')
if total > 100000:
raise ValueError('本CPU教学生成器最多100000条；请使用小规模配置。')
for key in ('lr', 'max_seconds_per_model'):
if isinstance(config[key], bool) or not isinstance(config[key], (int, float)) or not math.isfinite(config[key]) or config[key] <= 0:
raise ValueError(f'{key} 必须为有限正数。')
if training:
if any(config[k] == 0 for k in ('train', 'validation', 'test')) or config['train'] < config['batch']:
raise ValueError('训练需要非空train/validation/test，且train样本数至少等于batch。')
if config['updates'] > 5000 or config['max_seconds_per_model'] > 600:
raise ValueError('本教学入口每模型最多5000更新且600秒；请先用小规模配置。')
return config
def data_generation_method(config):
lo, hi = config['source_lengths']
capacity = sum(6**n for n in range(lo, hi + 1))
total = sum(config[k] for k in ('train', 'validation', 'test'))
return 'bounded universe permutation' if capacity <= 100000 and total * 2 >= capacity else 'uniform length and symbols; bounded rejection'
def make_data(config=CONFIG):
config = validate_config(config)
rng=np.random.default_rng(config['seed']); seen=set(); rows=[]
lo, hi = config['source_lengths']
total = sum(config[k] for k in ('train','validation','test'))
if data_generation_method(config) == 'bounded universe permutation':
# Dense small spaces: sample integer ranks without a coupon-collector loop.
capacity = sum(6**n for n in range(lo, hi + 1))
for rank in rng.permutation(capacity)[:total]:
rank = int(rank)
for n in range(lo, hi + 1):
if rank < 6**n: break
rank -= 6**n
x = [3] * n
for i in range(n - 1, -1, -1):
rank, digit = divmod(rank, 6); x[i] += digit
rows.append({'source': x, 'target': x[::-1] + [EOS]})
else:
# DN1 defaults take exactly the original random draws, preserving its split.
for _ in range(max(1000, 50 * total)):
n=int(rng.integers(lo,hi+1)); x=tuple(rng.integers(3,len(VOCAB),n).tolist())
if x not in seen:
seen.add(x); rows.append({'source':list(x),'target':list(x[::-1])+[EOS]})
if len(rows) == total: break
if len(rows) != total:
raise ValueError('已达到有界去重采样上限；请减少样本数或扩大source_lengths，不会无限循环。')
result={}; start=0
for k in ['train','validation','test']:
result[k]=rows[start:start+config[k]]; start+=config[k]
groups=[{tuple(r['source']) for r in result[k]} for k in result]
assert all(not groups[i]&groups[j] for i in range(3) for j in range(i))
return result
def batch(rows):
lengths=torch.tensor([len(r['source']) for r in rows])
n,m=int(lengths.max()),max(len(r['target']) for r in rows)
src=torch.full((len(rows),n),PAD,dtype=torch.long)
target=torch.full((len(rows),m),PAD,dtype=torch.long)
for i,r in enumerate(rows):
src[i,:len(r['source'])]=torch.tensor(r['source'])
target[i,:len(r['target'])]=torch.tensor(r['target'])
decoder_input=torch.cat([torch.full((len(rows),1),BOS),target[:,:-1]],1)
return src,lengths,decoder_input,target
class Seq2Seq(nn.Module):
def __init__(self,reading='fixed',config=CONFIG):
if reading not in ('fixed', 'dynamic'):
raise ValueError('reading 只支持 fixed 或 dynamic。')
super().__init__(); self.reading=reading; self.config=validate_config(config)
e,h,d=config['embed'],config['encoder_hidden'],config['decoder_hidden']
self.embedding=nn.Embedding(len(VOCAB),e,padding_idx=PAD)
self.encoder=nn.LSTM(e,h,batch_first=True,bidirectional=True)
self.decoder=nn.LSTMCell(e+2*h,d)
self.output=nn.Linear(d+2*h,len(VOCAB))
if reading=='dynamic': self.query=nn.Linear(d,2*h,bias=False)
def encode(self,src,lengths):
p=pack_padded_sequence(self.embedding(src),lengths.cpu(),batch_first=True,enforce_sorted=False)
z,_=self.encoder(p)
z,_=pad_packed_sequence(z,batch_first=True,total_length=src.shape[1])
mask=torch.arange(src.shape[1])[None,:]<lengths[:,None]
return z,mask
def initial(self,n):
return (torch.zeros(n,self.config['decoder_hidden']),torch.zeros(n,self.config['decoder_hidden']))
def step(self,token,state,z,mask):
h,c=state
if self.reading=='dynamic':
scores=torch.einsum('bd,bnd->bn',self.query(h),z)/math.sqrt(z.shape[-1])
weights=torch.softmax(scores.masked_fill(~mask,-torch.inf),-1)
else: weights=mask.float()/mask.sum(-1,keepdim=True)
context=torch.einsum('bn,bnd->bd',weights,z)
h,c=self.decoder(torch.cat([self.embedding(token),context],-1),(h,c))
logits=self.output(torch.cat([h,context],-1))
logits=logits.clone(); logits[:,:2]=-torch.inf
return logits,(h,c),weights
def forward(self,src,lengths,decoder_input):
z,mask=self.encode(src,lengths); state=self.initial(len(src)); outputs=[]; weights=[]
for t in range(decoder_input.shape[1]):
logit,state,a=self.step(decoder_input[:,t],state,z,mask)
outputs.append(logit);weights.append(a)
return torch.stack(outputs,1),torch.stack(weights,1)
def loss_on(model,rows):
src,l,di,y=batch(rows); logits,_=model(src,l,di)
# PAD positions are excluded; EOS is a target; denominator is predicted tokens.
return nn.functional.cross_entropy(logits.reshape(-1,len(VOCAB)),y.flatten(),ignore_index=PAD)
@torch.no_grad()
def greedy(model,rows,max_steps=None):
if max_steps is None: max_steps=max(12, max(len(r['source']) for r in rows)+2)
src,l,_,_=batch(rows);z,mask=model.encode(src,l);state=model.initial(len(src))
token=torch.full((len(src),),BOS); done=torch.zeros(len(src),dtype=torch.bool)
predictions=[[] for _ in rows]; total=torch.zeros(len(rows));trace=[]
for t in range(max_steps):
logits,state,a=model.step(token,state,z,mask); lp=logits.log_softmax(-1);token=lp.argmax(-1)
for i in range(len(rows)):
if not done[i]:predictions[i].append(int(token[i]));total[i]+=lp[i,token[i]]
trace.append(a.numpy());done|=token==EOS
if done.all():break
return predictions,total.tolist(),trace
@torch.no_grad()
def evaluate(model,rows):
model.eval(); loss_sum=0.;tokens=0;predictions=[];scores=[]
for start in range(0,len(rows),64):
r=rows[start:start+64]; nt=sum(len(x['target']) for x in r)
loss_sum+=float(loss_on(model,r))*nt;tokens+=nt
p,s,_=greedy(model,r);predictions+=p;scores+=s
correct=[p==r['target'] for p,r in zip(predictions,rows)]
lengths=sorted({len(r['source']) for r in rows})
by_length={str(n):{'n':sum(len(r['source'])==n for r in rows),
'exact_match':float(np.mean([c for c,r in zip(correct,rows) if len(r['source'])==n]))} for n in lengths}
examples=[dict(source=''.join(VOCAB[t] for t in r['source']),
target=' '.join(VOCAB[t] for t in r['target']),prediction=' '.join(VOCAB[t] for t in p),
correct=c,log_probability=s,ended=bool(p and p[-1]==EOS))
for r,p,c,s in zip(rows,predictions,correct,scores)]
return dict(nll=loss_sum/tokens,token_count=tokens,exact_match=float(np.mean(correct)),
by_length=by_length,first_examples=examples[:8],failures=[e for e in examples if not e['correct']][:8])
@torch.no_grad()
def beam(model,row,width=2,max_steps=None):
if max_steps is None: max_steps=max(12, len(row['source'])+2)
src,l,_,_=batch([row]);z,mask=model.encode(src,l)
beams=[([],0.,model.initial(1))]; traces=[]
for t in range(max_steps):
candidates=[]
for seq,score,state in beams:
if seq and seq[-1]==EOS:candidates.append((seq,score,state));continue
token=torch.tensor([seq[-1] if seq else BOS])
logits,newstate,_=model.step(token,state,z,mask)
lp=logits.log_softmax(-1)[0]
for k in range(2,len(VOCAB)):candidates.append((seq+[k],score+float(lp[k]),newstate))
beams=sorted(candidates,key=lambda x:(-x[1],x[0]))[:width]
traces.append([{'tokens':[VOCAB[k] for k in s],'logp':p} for s,p,_ in beams])
if all(s[-1]==EOS for s,_,_ in beams):break
finished=[b for b in beams if b[0][-1]==EOS]
seq,score,_=max(finished or beams,key=lambda b:b[1])
return dict(tokens=[VOCAB[k] for k in seq],log_probability=score,ended=seq[-1]==EOS,
correct=seq==row['target'],trace=traces)
def initialized_pair(config=CONFIG):
torch.manual_seed(config['seed']); fixed=Seq2Seq('fixed',config)
torch.manual_seed(config['seed']+1);dynamic=Seq2Seq('dynamic',config)
state=dynamic.state_dict()
for k,v in fixed.state_dict().items():state[k]=v.clone()
dynamic.load_state_dict(state)
assert all(torch.equal(v,dynamic.state_dict()[k]) for k,v in fixed.state_dict().items())
return {'fixed':fixed,'dynamic':dynamic}
def overfit_check(config=CONFIG):
# A fixed short-task implementation check, independent of the student's range.
# Under DN1 defaults these are the exact original 16 examples.
sanity_config={**config, 'source_lengths':[3,7], 'train':3000, 'validation':300, 'test':300}
rows=[r for r in make_data(sanity_config)['train'] if len(r['source'])<=4][:16]
result={}
for kind,model in initialized_pair(config).items():
opt=torch.optim.Adam(model.parameters(),lr=.01); started=time.perf_counter()
for step in range(300):
opt.zero_grad();loss=loss_on(model,rows);loss.backward();nn.utils.clip_grad_norm_(model.parameters(),1.);opt.step()
evaluation=evaluate(model,rows)
result[kind]={'updates':300,'examples':16,'source_lengths':[3,4],
'role':'independent short-task implementation check; formal training reinitializes',
'seconds':time.perf_counter()-started,**evaluation}
assert evaluation['exact_match']==1.,f'{kind} short overfit failed'
return result
def train_pair(root,scope='local',config=None):
config=validate_config(dict(CONFIG if config is None else config),training=True)
root=Path(root);root.mkdir(parents=True,exist_ok=True)
data=make_data(config); (root/'data.json').write_text(json.dumps(data))
(root/'config.json').write_text(json.dumps(config,indent=2))
started=time.perf_counter(); sanity=overfit_check(config); models=initialized_pair(config)
optimizers={k:torch.optim.Adam(m.parameters(),lr=config['lr']) for k,m in models.items()}
curves={k:[] for k in models};elapsed={k:0. for k in models};checkpoints=[]
rng=np.random.default_rng(config['seed']+2);token_budget=0;sample_budget=0
for step in range(1,config['updates']+1):
idx=rng.choice(len(data['train']),config['batch'],replace=False); rows=[data['train'][int(i)] for i in idx]
token_budget+=sum(len(r['target']) for r in rows);sample_budget+=len(rows)
for kind,model in models.items():
t=time.perf_counter();model.train();optimizers[kind].zero_grad();loss=loss_on(model,rows)
loss.backward();grad=nn.utils.clip_grad_norm_(model.parameters(),1.);optimizers[kind].step();elapsed[kind]+=time.perf_counter()-t
curves[kind].append({'step':step,'batch_nll':float(loss.detach()),'gradient_norm_before_clip':float(grad)})
if step%config['checkpoint_interval']==0 or step==config['updates']:
for kind,model in models.items():
validation=evaluate(model,data['validation']);curves[kind][-1]['validation']=validation
torch.save({'state':model.state_dict(),'config':config,'reading':kind,'step':step,
'scope':scope,'token_budget':token_budget},root/f'{kind}-{step}.pt')
print(scope,kind,'step',step,'validation NLL',round(validation['nll'],4),'EM',validation['exact_match'],flush=True)
checkpoints.append(step)
if max(elapsed.values())>=config['max_seconds_per_model']:break
if step==30:
print('30-update pilot seconds',elapsed,flush=True)
# Primary comparison: same last common checkpoint, never choose with test scores.
final=checkpoints[-1];evaluation={k:evaluate(m,data['test']) for k,m in models.items()}
report={'scope':scope,'created_at_utc':__import__('datetime').datetime.now(__import__('datetime').timezone.utc).isoformat(),
'config':config,'data_generation_method':data_generation_method(config),
'common_step':final,'training_tokens_each':token_budget,'training_samples_each':sample_budget,
'parameter_counts':{k:sum(p.numel() for p in m.parameters()) for k,m in models.items()},
'training_seconds':elapsed,'total_seconds':time.perf_counter()-started,'sanity':sanity,'curves':curves,
'evaluation':evaluation,'data_sha256':hashlib.sha256((root/'data.json').read_bytes()).hexdigest(),
'environment':{'python':platform.python_version(),'torch':torch.__version__,'device':'cpu','threads':torch.get_num_threads()},
'search_examples':{k:{str(w):beam(m,data['test'][0],w) for w in [1,2,4]} for k,m in models.items()},
'limit':'one seed, synthetic reversal, unequal active parameter counts; no general architecture ranking'}
(root/'report.json').write_text(json.dumps(report,indent=2))
return models,data,report
def train_student(root,experiment_name='seq2seq-trial',config=None):
config=validate_config(dict(CONFIG if config is None else config),training=True)
before=runs.reference_hashes(root)
destination=runs.new_student_run(root,'seq2seq',experiment_name)
try:
models,data,report=train_pair(destination,'student',config)
finally:
runs.assert_references_unchanged(root,before)
return models,data,report,destination
def load_pair(root):
root=Path(root);report=json.loads((root/'report.json').read_text());models={}
for kind in ['fixed','dynamic']:
saved=torch.load(root/f"{kind}-{report['common_step']}.pt",map_location='cpu',weights_only=True)
model=Seq2Seq(kind,saved['config']);model.load_state_dict(saved['state']);model.eval();models[kind]=model
return models,json.loads((root/'data.json').read_text()),report
if __name__=='__main__':
import argparse
parser=argparse.ArgumentParser();parser.add_argument('--output',required=True);parser.add_argument('--scope',default='local')
a=parser.parse_args();train_pair(a.output,a.scope)