| 508 | ] |
| 509 | |
| 510 | def forward_test(self, x): |
| 511 | self.attn_maps = [] |
| 512 | if not self.ds: |
| 513 | visual_f = x + self.vis_pos_embed |
| 514 | elif self.pos2d: |
| 515 | visual_f = x + self.vis_pos_embed[:, :, :x.shape[2], :x.shape[3]] |
| 516 | visual_f = x.flatten(2).transpose(1, 2) |
| 517 | else: |
| 518 | visual_f = x |
| 519 | bs = x.shape[0] |
| 520 | |
| 521 | if self.next_mode: |
| 522 | ques_next = self.next_token.tile([bs, 1, 1, 1]).squeeze(1) |
| 523 | prompt_next_embed = self.prompt_next_embed.tile([bs, 1, 1, |
| 524 | 1]).squeeze(1) |
| 525 | prompt_id = torch.full([bs, self.sub_str_len], |
| 526 | self.bos_next, |
| 527 | dtype=torch.long, |
| 528 | device=x.device) |
| 529 | pred_id_list = torch.full([bs, self.max_len], |
| 530 | self.ignore_index, |
| 531 | dtype=torch.long, |
| 532 | device=x.device) |
| 533 | logits_all = [] |
| 534 | mask_pad = torch.zeros([bs, 1], |
| 535 | dtype=torch.float32, |
| 536 | device=x.device) |
| 537 | for j in range(0, self.max_len - 1): |
| 538 | |
| 539 | prompt_next = torch.concat([ |
| 540 | prompt_next_embed[:, :1, :], |
| 541 | prompt_next_embed[:, 1:, :] + self.char_embed(prompt_id) |
| 542 | ], 1) # b, sub_l, dim |
| 543 | mask_next = torch.where(prompt_id == self.bos_next, |
| 544 | float('-inf'), 0) # b, subs_l |
| 545 | mask = torch.concat([mask_pad, mask_next], 1) |
| 546 | ques_next_i = ques_next |
| 547 | visual_f_i = visual_f |
| 548 | for layer in self.cmff_decoder: |
| 549 | ques_next_i = layer(ques_next_i, prompt_next, visual_f_i, |
| 550 | mask.unsqueeze(1)) |
| 551 | self.attn_maps.append( |
| 552 | self.cmff_decoder[-1].question_to_images_cross_attn. |
| 553 | attn_map[0]) |
| 554 | logits_next_i = self.ques1_head(self.norm_pred(ques_next_i)) |
| 555 | logits = F.softmax(logits_next_i, -1) |
| 556 | pred_id_i = logits.argmax(-1) |
| 557 | logits_all.append(logits) |
| 558 | pred_id_list[:, j:j + 1] = pred_id_i |
| 559 | if (pred_id_list == self.eos).any(dim=-1).all(): |
| 560 | break |
| 561 | prompt_id = torch.concat( |
| 562 | [ |
| 563 | prompt_id[:, 1:, ], |
| 564 | pred_id_i, |
| 565 | ], |
| 566 | 1, |
| 567 | ) |