LAB 13 · 看图生成一句话：训练、解码与消融

# Detached readonly runs write only to an isolated ephemeral directory.
import os, tempfile
_LAB_WORKDIR = tempfile.mkdtemp(prefix="lecture03-lab13-")
os.chdir(_LAB_WORKDIR)
print("ephemeral workdir:", _LAB_WORKDIR)

seqs=[[1,4,5,2],[1,6,2,0]]
print('warm-up input',[s[:-1] for s in seqs])
print('warm-up target',[s[1:] for s in seqs])
print('warm-up loss mask',[[int(t!=0) for t in s[1:]] for s in seqs])
# Every outgoing distribution sums to one. Both paths terminate after two tokens.
tree={():{'A':.55,'B':.45},('A',):{'X':.45,'Z':.55},('B',):{'Y':.80,'Z':.20}}
assert all(abs(sum(p.values())-1)<1e-9 for p in tree.values())
greedy_path=();greedy_probability=1.
for _ in range(2):
    token=max(tree[greedy_path],key=tree[greedy_path].get)
    greedy_probability*=tree[greedy_path][token];greedy_path+= (token,)
beams=[((),1.)]
for _ in range(2):
    beams=sorted([(seq+(word,),p*conditional) for seq,p in beams for word,conditional in tree[seq].items()],key=lambda b:b[1],reverse=True)[:2]
print('actual tree greedy:',greedy_path,greedy_probability,'beam2:',beams[0])

from pathlib import Path
import json, hashlib, platform, sys, time
from datetime import datetime, timezone
import numpy as np
from PIL import Image, ImageDraw
import matplotlib
if 'get_ipython' not in globals(): matplotlib.use('Agg')
import matplotlib.pyplot as plt
import torch
from torch import nn
import torch.nn.functional as F
OUT = Path('outputs/lab13'); OUT.mkdir(parents=True, exist_ok=True)
torch.set_num_threads(1)
torch.manual_seed(1303)
np.random.seed(1303)
START = time.perf_counter()
VOCAB = ['<pad>', '<bos>', '<eos>', 'a', 'red', 'blue', 'green', 'square', 'circle', 'on', 'left', 'right']
W = {s:i for i,s in enumerate(VOCAB)}
PAD, BOS, EOS = 0, 1, 2
COLORS = {'red':(220,55,55), 'blue':(45,95,225), 'green':(35,170,75)}
print('CPU', torch.__version__, 'seed 1303; no downloaded data or weights')

def render_object(color, shape, position, rng):
    canvas = Image.new('RGB', (32,32), (244,242,236))
    painter = ImageDraw.Draw(canvas)
    cx = {'left':8, 'center':16, 'right':24}[position] + int(rng.integers(-1,2))
    cy = 16 + int(rng.integers(-4,5)); radius = int(rng.integers(4,7))
    box = (cx-radius, cy-radius, cx+radius, cy+radius)
    if shape == 'square': painter.rectangle(box, fill=COLORS[color])
    else: painter.ellipse(box, fill=COLORS[color])
    arr = np.array(canvas).astype(np.float32) / 255
    arr = np.clip(arr + rng.normal(0,.018,arr.shape),0,1).astype(np.float32)
    return arr

def make_caption_data(n, seed):
    rng = np.random.default_rng(seed); images=[]; sequences=[]; descriptions=[]
    conditions = [(c,s,p) for c in COLORS for s in ['square','circle'] for p in ['left','center','right']]
    for i in range(n):
        color, shape, pos = conditions[i % len(conditions)]
        words = ['a',color,shape] + ([] if pos=='center' else ['on',pos])
        sequence = [BOS] + [W[w] for w in words] + [EOS]
        sequences.append(sequence+[PAD]*(7-len(sequence)))
        images.append(render_object(color,shape,pos,rng)); descriptions.append(' '.join(words))
    return torch.tensor(np.stack(images)).permute(0,3,1,2), torch.tensor(sequences), descriptions

Xtrain, Ytrain, train_text = make_caption_data(720,13031)
Xtest, Ytest, test_text = make_caption_data(180,13033)
# Independent random pixels/geometry; no duplicate image crosses the split.
train_hashes = {hashlib.sha256(x.numpy().tobytes()).hexdigest() for x in Xtrain}
assert not train_hashes.intersection(hashlib.sha256(x.numpy().tobytes()).hexdigest() for x in Xtest)
np.savez_compressed(OUT/'lab13-data.npz', train_images=Xtrain.numpy(), train_tokens=Ytrain.numpy(), test_images=Xtest.numpy(), test_tokens=Ytest.numpy())
print('images:', tuple(Xtrain.shape), tuple(Xtest.shape), 'tokens:', tuple(Ytrain.shape))
print('teacher input:', Ytrain[:3,:-1].tolist())
print('shifted target:',Ytrain[:3,1:].tolist())
print('loss mask:',Ytrain[:3,1:].ne(PAD).int().tolist())

class Captioner(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = nn.Sequential(nn.Conv2d(3,12,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),nn.Conv2d(12,24,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(24*8*8,48),nn.LayerNorm(48),nn.ReLU())
        self.word = nn.Embedding(len(VOCAB),24,padding_idx=PAD)
        self.initial = nn.Linear(48,64)
        self.rnn = nn.GRU(24+48,64,batch_first=True)
        self.output = nn.Linear(64,len(VOCAB))
    def encode(self, pixels):
        return self.cnn(pixels - 0.5)
    def forward(self, pixels, prefix):
        visual = self.encode(pixels)
        features = visual[:,None,:].expand(-1,prefix.shape[1],-1)
        hidden,_ = self.rnn(torch.cat((self.word(prefix),features),-1),self.initial(visual)[None])
        return self.output(hidden)

model = Captioner()
optimizer = torch.optim.Adam(model.parameters(),lr=.001)
curve=[]
for epoch in range(45):
    model.train(); order=torch.randperm(len(Xtrain)); total=0.; tokens=0
    for ix in order.split(60):
        logits=model(Xtrain[ix],Ytrain[ix,:-1])
        targets=Ytrain[ix,1:]
        # Sum over real targets then divide by their count, not padded positions.
        loss=F.cross_entropy(logits.reshape(-1,len(VOCAB)),targets.reshape(-1),ignore_index=PAD,reduction='sum')/targets.ne(PAD).sum()
        optimizer.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(),5); optimizer.step()
        count=int(targets.ne(PAD).sum()); total+=float(loss.detach())*count; tokens+=count
    model.eval()
    with torch.no_grad():
        logits=model(Xtest,Ytest[:,:-1]); gold=Ytest[:,1:]
        test_loss=float(F.cross_entropy(logits.reshape(-1,len(VOCAB)),gold.reshape(-1),ignore_index=PAD))
    curve.append({'epoch':epoch+1,'train_token_ce':total/tokens,'heldout_teacher_forced_ce':test_loss})
    if epoch in [0,5,11,23,44]: print(curve[-1])
torch.save({'state_dict':model.state_dict(),'vocab':VOCAB,'seed':1303},OUT/'lab13-caption-weights.pt')

@torch.no_grad()
def greedy(pixels, prefix=(BOS,), max_new=8):
    model.eval(); seq=torch.tensor(prefix).repeat(len(pixels),1); finished=torch.zeros(len(pixels),dtype=torch.bool)
    for _ in range(max_new):
        nxt=model(pixels,seq)[:,-1].argmax(-1)
        nxt=torch.where(finished,torch.tensor(EOS),nxt)
        seq=torch.cat((seq,nxt[:,None]),1); finished|=nxt.eq(EOS)
        if bool(finished.all()): break
    return seq.tolist()

@torch.no_grad()
def beam_one(pixel, width=3, max_new=8):
    beams=[([BOS],0.0,False)]
    for _ in range(max_new):
        candidates=[]
        for seq,score,done in beams:
            if done: candidates.append((seq,score,True)); continue
            lp=F.log_softmax(model(pixel[None],torch.tensor([seq]))[0,-1],dim=-1)
            values,indices=lp.topk(width)
            for value,idx in zip(values.tolist(),indices.tolist()):
                candidates.append((seq+[idx],score+value,idx==EOS))
        beams=sorted(candidates,key=lambda b:b[1],reverse=True)[:width]
        if all(b[2] for b in beams): break
    # Ended hypotheses only if at least one ended; otherwise max-length fallback.
    ended=[b for b in beams if b[2]]
    return max(ended or beams,key=lambda b:b[1])

def words(ids):
    out=[]
    for idx in ids:
        if idx==EOS: break
        if idx not in [BOS,PAD]: out.append(VOCAB[idx])
    return ' '.join(out)

pred=[words(s) for s in greedy(Xtest)]
beam_pred=[words(beam_one(x)[0]) for x in Xtest]
zero_pred=[words(s) for s in greedy(torch.zeros_like(Xtest))]
# A deterministic derangement shifts test images; target captions stay fixed.
shuffle_index=torch.roll(torch.arange(len(Xtest)),1)
shuffle_pred=[words(s) for s in greedy(Xtest[shuffle_index])]
exact=lambda p: float(np.mean([a==b for a,b in zip(p,test_text)]))
metrics={'heldout_images':len(Xtest),'greedy_exact_match':exact(pred),'beam3_exact_match':exact(beam_pred),'zero_pixels_exact_match':exact(zero_pred),'shuffled_pixels_exact_match':exact(shuffle_pred),'final_train_token_ce':curve[-1]['train_token_ce'],'final_heldout_teacher_forced_ce':curve[-1]['heldout_teacher_forced_ce']}
print(json.dumps(metrics,indent=2))
# Hold the language prefix fixed; change only the input image.
rng=np.random.default_rng(999)
control_specs=[('red','square','left'),('blue','circle','right'),('green','square','center')]
control_pixels=torch.tensor(np.stack([render_object(*spec,rng) for spec in control_specs])).permute(0,3,1,2)
control_pred=[words(s) for s in greedy(control_pixels,prefix=(BOS,W['a']))]
controls=[{'condition':list(spec),'fixed_prefix':'<bos> a','generated':p} for spec,p in zip(control_specs,control_pred)]
print('same prefix, changed image:',json.dumps(controls,indent=2))

fig,axes=plt.subplots(2,3,figsize=(10,6))
chosen=[0,3,7,10,14,17]
for ax,i in zip(axes.flat,chosen):
    ax.imshow(Xtest[i].permute(1,2,0)); ax.axis('off')
    ax.set_title('truth: '+test_text[i]+'\ngreedy: '+pred[i]+'\nbeam: '+beam_pred[i],fontsize=9)
fig.suptitle('LAB13: held-out images; model-generated captions',fontsize=13); fig.tight_layout()
fig.savefig(OUT/'lab13-caption-predictions.png',dpi=160); plt.show(); plt.close(fig)
fig,axes=plt.subplots(1,2,figsize=(10,3.4))
axes[0].plot([r['epoch'] for r in curve],[r['train_token_ce'] for r in curve],label='train'); axes[0].plot([r['epoch'] for r in curve],[r['heldout_teacher_forced_ce'] for r in curve],label='held-out'); axes[0].set(xlabel='epoch',ylabel='mean non-pad token CE');axes[0].legend()
labels=['greedy','beam 3','zero image','shuffled image'];values=[metrics[k] for k in ['greedy_exact_match','beam3_exact_match','zero_pixels_exact_match','shuffled_pixels_exact_match']]
axes[1].bar(labels,values,color=['#356b91','#40868e','#a9adb0','#a9adb0']);axes[1].set(ylim=(0,1.12),ylabel='exact caption accuracy (180 images)')
for i,v in enumerate(values):axes[1].text(i,v+.025,f'{v:.1%}',ha='center')
fig.tight_layout();fig.savefig(OUT/'lab13-caption-controls.png',dpi=160);plt.show();plt.close(fig)
fig,axes=plt.subplots(1,3,figsize=(9,3))
for ax,x,row in zip(axes,control_pixels,controls):
    ax.imshow(x.permute(1,2,0));ax.axis('off');ax.set_title('prefix: <bos> a\n'+row['generated'],fontsize=10)
fig.tight_layout();fig.savefig(OUT/'lab13-same-prefix.png',dpi=160);plt.show();plt.close(fig)
result={'notebook_id':'3da47f3ee3554ee0a36e0bb6a2fdc61f','execution':'local CPU, fresh process; not a Deepnote cloud run','created_utc':datetime.now(timezone.utc).isoformat(),'elapsed_seconds':time.perf_counter()-START,'environment':{'python':sys.version,'torch':torch.__version__,'numpy':np.__version__,'platform':platform.platform()},'split':{'train_seed':13031,'test_seed':13033,'train_images':720,'test_images':180,'description':'independent noisy renders; same 18 semantic combinations on both sides; no compositional/general-photo claim'},'metrics':metrics,'curve':curve,'same_prefix_controls':controls,'predictions':[{'truth':t,'greedy':p,'beam3':b,'zero':z,'shuffled':s} for t,p,b,z,s in zip(test_text,pred,beam_pred,zero_pred,shuffle_pred)],'weight_sha256':hashlib.sha256((OUT/'lab13-caption-weights.pt').read_bytes()).hexdigest(),'data_sha256':hashlib.sha256((OUT/'lab13-data.npz').read_bytes()).hexdigest()}
(OUT/'lab13-results.json').write_text(json.dumps(result,indent=2))
print('saved',OUT,'seconds',round(result['elapsed_seconds'],2))