| 644 | return feat_mat |
| 645 | |
| 646 | def forward_encoder(self, x1, x_next, mask_ratio): |
| 647 | # embed patches |
| 648 | x1 = self.patch_embed(x1) |
| 649 | # x2 = self.patch_embed(x2) |
| 650 | |
| 651 | |
| 652 | # add pos embed w/o cls token |
| 653 | x1 = x1 + self.pos_embed[:, 1:, :] + self.temp_embed[:, 0, :] |
| 654 | |
| 655 | # handles other time stamps |
| 656 | xs = [] |
| 657 | for ts in range(self.time_stamp-1): |
| 658 | xt = x_next[:, ts, :, :, :] # [Bxa, C, H, W] |
| 659 | # print("xt size", xt.size()) |
| 660 | xt = self.patch_embed(xt) + self.pos_embed[:, 1:, :] + self.temp_embed[:, ts+1, :] |
| 661 | xs.append(xt) |
| 662 | # x2 = x2 + self.pos_embed[:, 1:, :] + self.temp_embed[:, 1, :] |
| 663 | |
| 664 | # masking: length -> length * mask_ratio |
| 665 | # x1, mask1, ids_restore1 = self.random_masking(x1, mask_ratio) |
| 666 | # x2, mask2, ids_restore2 = self.random_masking(x2, mask_ratio) |
| 667 | # x_masked, x1len, mask1, ids_restore1 = self.more_random_masking(x1, xs, mask_ratio) |
| 668 | # complement masking |
| 669 | # x_masked, x1len, mask1, ids_restore1 = self.complement_masking(x1, xs, mask_ratio) |
| 670 | x_masked, x1len, mask1, ids_restore1 = self.masking_handle(x1, xs, mask_ratio) |
| 671 | # print(x_masked.size()) |
| 672 | # print(x1len) |
| 673 | |
| 674 | # append cls token |
| 675 | cls_token = self.cls_token + self.pos_embed[:, :1, :] |
| 676 | cls_tokens = cls_token.expand(x_masked.shape[0], -1, -1) |
| 677 | x = torch.cat((cls_tokens, x_masked), dim=1) |
| 678 | |
| 679 | # apply Transformer blocks |
| 680 | for blk in self.blocks: |
| 681 | x = blk(x) |
| 682 | x = self.norm(x) |
| 683 | # compress for communication |
| 684 | x = self.compressor(x) |
| 685 | |
| 686 | return x, mask1, ids_restore1, x1len |
| 687 | |
| 688 | def forward_decoder(self, x, ids_restore): |
| 689 | # decompress |