| 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 | """ |