Modified from the original forward, make fusion happen
(self, imgs1, imgs_next, teacher, trans_matrices, num_agent_tensor, batch_size, mask_ratio=0.75)
| 233 | return loss |
| 234 | |
| 235 | def forward(self, imgs1, imgs_next, teacher, trans_matrices, num_agent_tensor, batch_size, mask_ratio=0.75): |
| 236 | """ |
| 237 | Modified from the original forward, make fusion happen |
| 238 | """ |
| 239 | p = self.patch_embed.patch_size[0] |
| 240 | self.patch_h = self.patch_w = imgs1.shape[2]//p |
| 241 | |
| 242 | latent, mask1, ids_restore, size1 = self.forward_encoder(imgs1, imgs_next, mask_ratio) |
| 243 | # latent: [Bxa, L, D] |
| 244 | # now we only consider reconstruct the first frame |
| 245 | # this can be extended to reconstruct both |
| 246 | latent_to_decode = latent[:, :1+size1, :] # CLS(0) + timesamp t (size1) |
| 247 | fused_latent = self.forward_fusion(latent_to_decode, ids_restore, trans_matrices, num_agent_tensor, batch_size) |
| 248 | pred = self.forward_decoder(fused_latent) # [N, L, p*p*3] |
| 249 | |
| 250 | # --- check communication --- |
| 251 | # fused_latent = self.forward_fusion(imgs1, ids_restore, trans_matrices, num_agent_tensor, batch_size) |
| 252 | # pred = fused_latent # cls token is not included in forward fusion |
| 253 | #---------------------------- |
| 254 | loss = self.forward_loss(teacher, pred, mask1) |
| 255 | # loss = self.forward_loss(imgs1, pred, mask1) # use single view as supervision |
| 256 | # what to do with the masking? now it is fused. So not exactly mask1 |
| 257 | # ONE solution: just calculate the loss over the entire input |
| 258 | return loss, pred, mask1, fused_latent[:,1:,:] # remove cls |
| 259 | |
| 260 | |
| 261 | class IndivMultiAgentMAEViT(MultiAgentMaskedAutoencoderViT): |
nothing calls this directly
no test coverage detected