3. 视觉特征怎样真正进入语言模型 · 实现代码

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