| 288 | return logit |
| 289 | |
| 290 | def forward_test(self, vis_feat): |
| 291 | bs = vis_feat.size(0) |
| 292 | |
| 293 | dec_seq = torch.full( |
| 294 | (bs, self.max_len + 1), |
| 295 | self.ignore_index, |
| 296 | dtype=torch.int64, |
| 297 | device=vis_feat.device, |
| 298 | ) |
| 299 | dec_seq[:, 0] = self.bos |
| 300 | logits = [] |
| 301 | for len_dec_seq in range(0, self.max_len): |
| 302 | sem_feat, sem_mask = self.semantic_branch(dec_seq[:, :len_dec_seq + |
| 303 | 1]) |
| 304 | pos_feat = self.positional_branch(sem_feat) |
| 305 | output = self.mdcdp( |
| 306 | sem_feat, |
| 307 | vis_feat, |
| 308 | pos_feat, |
| 309 | tgt_mask=sem_mask, |
| 310 | memory_mask=None, |
| 311 | ) |
| 312 | |
| 313 | dec_output = output[:, -1:, :] |
| 314 | |
| 315 | word_prob = F.softmax(self.tgt_word_prj(dec_output), dim=-1) |
| 316 | logits.append(word_prob) |
| 317 | if len_dec_seq < self.max_len: |
| 318 | # greedy decode. add the next token index to the target input |
| 319 | dec_seq[:, len_dec_seq + 1] = word_prob.squeeze(1).argmax(-1) |
| 320 | # Efficient batch decoding: If all output words have at least one EOS token, end decoding. |
| 321 | if (dec_seq == self.eos).any(dim=-1).all(): |
| 322 | break |
| 323 | logits = torch.cat(logits, dim=1) |
| 324 | return logits |
| 325 | |
| 326 | def forward_beam(self, x): |
| 327 | """Translation work in one batch.""" |