joint reconstruction Serves as the DetModel, handles fusion/communication between encoder and decoder
| 22 | |
| 23 | |
| 24 | class FusionMultiAgentMAEViT(MultiAgentMaskedAutoencoderViT): |
| 25 | """ |
| 26 | joint reconstruction |
| 27 | Serves as the DetModel, handles fusion/communication between encoder and decoder |
| 28 | """ |
| 29 | def __init__(self, img_size=224, patch_size=16, in_chans=3, |
| 30 | embed_dim=1024, depth=24, num_heads=16, |
| 31 | decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, |
| 32 | mlp_ratio=4., norm_layer=nn.LayerNorm, decoder_head="conv3", norm_pix_loss=False, time_stamp=1, mask_method="random"): |
| 33 | super(FusionMultiAgentMAEViT, self).__init__( |
| 34 | img_size=img_size, |
| 35 | patch_size=patch_size, |
| 36 | in_chans=in_chans, |
| 37 | embed_dim=embed_dim, |
| 38 | depth=depth, |
| 39 | num_heads=num_heads, |
| 40 | decoder_embed_dim=decoder_embed_dim, |
| 41 | decoder_depth=decoder_depth, |
| 42 | decoder_num_heads=decoder_num_heads, |
| 43 | mlp_ratio=mlp_ratio, |
| 44 | norm_layer=norm_layer, |
| 45 | decoder_head = decoder_head, |
| 46 | norm_pix_loss=norm_pix_loss, |
| 47 | time_stamp = time_stamp, |
| 48 | mask_method = mask_method |
| 49 | ) |
| 50 | # for neighbor agents' features |
| 51 | self.patch_h = 0 |
| 52 | self.patch_w = 0 |
| 53 | # self.num_agent = 0 |
| 54 | # self.neighbor_feat_list = [] |
| 55 | # self.tg_agent = None |
| 56 | if mask_method == "random": |
| 57 | print("do random masking") |
| 58 | self.masking_handle = self.more_random_masking |
| 59 | elif mask_method == "complement": |
| 60 | print("do complement masking") |
| 61 | self.masking_handle = self.more_complement_masking |
| 62 | else: |
| 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()) |
no outgoing calls
no test coverage detected