as the pre process of the decoder, handles: 1) add mask, restore feature maps shape like (b x agent x channel x h x w) 2) communication/fushion/aggregation 3) ready for decoder
(self, x, ids_restore, trans_matrices, num_agent_tensor, batch_size)
| 63 | raise NotImplementedError(mask_method) |
| 64 | |
| 65 | def forward_fusion(self, x, ids_restore, trans_matrices, num_agent_tensor, batch_size): |
| 66 | """ |
| 67 | as the pre process of the decoder, handles: |
| 68 | 1) add mask, restore feature maps shape like (b x agent x channel x h x w) |
| 69 | 2) communication/fushion/aggregation |
| 70 | 3) ready for decoder |
| 71 | """ |
| 72 | device = x.device |
| 73 | # # decompress |
| 74 | x = self.decompressor(x) |
| 75 | # # embed tokens |
| 76 | x = self.decoder_embed(x) |
| 77 | # append mask tokens to sequence |
| 78 | mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1) |
| 79 | x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1) # no cls token |
| 80 | x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])) # unshuffle |
| 81 | # print("x_", x_.size()) |
| 82 | ## --- check fusion (communication) --- |
| 83 | # x_ = self.patchify(x) # here x is input image, directly |
| 84 | |
| 85 | # x_: (B, seq, chns) |
| 86 | ## -- reshape back to 256x256 --- |
| 87 | # feature_maps = self.unpatchify(x_) |
| 88 | ## ------------- |
| 89 | ## --- reshape into B C H W --- |
| 90 | feature_maps = x_.reshape(x_.shape[0], self.patch_h, self.patch_w, x_.shape[-1]) # (B, h, w, chns) |
| 91 | feature_maps = feature_maps.permute(0, 3, 1, 2) # (B, chns, h, w) |
| 92 | # print("feature map", feature_maps.size()) |
| 93 | # ------- |
| 94 | ## --- do fusion --- |
| 95 | size = self.get_feature_maps_size(feature_maps) |
| 96 | # print(size) |
| 97 | assert feature_maps.size(0) % batch_size == 0, (feature_maps.size(), batch_size) |
| 98 | self.num_agent = feature_maps.size(0) // batch_size |
| 99 | feat_list = self.build_feature_list(batch_size, feature_maps) |
| 100 | # [[1,1,256,32,32]x5] NOTE should it be [[B, 1, 256, 32, 32]x5]? |
| 101 | # print(feat_list) |
| 102 | local_com_mat = self.build_local_communication_matrix( |
| 103 | feat_list) # [2 5 512 32 32] [batch, agent, channel, height, width] |
| 104 | # # FIXME: check size |
| 105 | # print("local com mat size", local_com_mat.shape) #[2,5,256,32,32] |
| 106 | local_com_mat_update = self.build_local_communication_matrix(feat_list) # to avoid the inplace operation |
| 107 | |
| 108 | for b in range(batch_size): |
| 109 | self.num_agent = num_agent_tensor[b, 0] |
| 110 | for i in range(self.num_agent): |
| 111 | self.tg_agent = local_com_mat[b, i] |
| 112 | # print("tg agent shape", self.tg_agent.shape) #[256,32,32] |
| 113 | self.neighbor_feat_list = [] |
| 114 | self.neighbor_feat_list.append(self.tg_agent) |
| 115 | all_warp = trans_matrices[b, i] # transformation [2 5 5 4 4] |
| 116 | # print(all_warp.shape)[5,4,4] |
| 117 | self.build_neighbors_feature_list(b, i, all_warp, self.num_agent, local_com_mat, |
| 118 | device, size) |
| 119 | |
| 120 | # feature update |
| 121 | # torch.save(torch.stack(self.neighbor_feat_list).detach().cpu(), "/mnt/NAS/home/zjx/Masked-Multiagent-Autoencoder/debug/nbf-{}-{}.pt".format(b, i)) |
| 122 | local_com_mat_update[b, i] = self.fusion() |
no test coverage detected