x1: [bxa, C, H, W] x_next: [bxa, ts-1, C, H, W] beq_next_frames
(self, x1, x_next, mask_ratio)
| 315 | return x, mask, ids_restore |
| 316 | |
| 317 | def forward_encoder_partial(self, x1, x_next, mask_ratio): |
| 318 | """ |
| 319 | x1: [bxa, C, H, W] |
| 320 | x_next: [bxa, ts-1, C, H, W] beq_next_frames |
| 321 | """ |
| 322 | # cat x1 and x_next to encoder independently |
| 323 | BA, C, H, W = x1.size() |
| 324 | x1 = x1.unsqueeze(1) |
| 325 | if self.time_stamp>1: |
| 326 | x_ind = torch.cat((x1, x_next), dim=1) # [bxa, ts, C, H, W] |
| 327 | else: |
| 328 | x_ind = x1 |
| 329 | # print(x_ind.size()) |
| 330 | assert x_ind.size(1) == self.time_stamp |
| 331 | x_ind = x_ind.reshape(BA*self.time_stamp, C, H, W) |
| 332 | # embed patches |
| 333 | x_ind = self.patch_embed(x_ind) |
| 334 | x_ind = x_ind + self.pos_embed[:, 1:, :] |
| 335 | |
| 336 | # mask before transformer encoding |
| 337 | # amortized masking, complement and random |
| 338 | x_masked, mask, ids_restore = self.masking_handle(x_ind, mask_ratio) |
| 339 | # print(x_masked.size()) |
| 340 | # print(mask.size()) |
| 341 | for blk in self.blocks: |
| 342 | x_masked = blk(x_masked) |
| 343 | x = self.norm(x_masked) |
| 344 | x = self.compressor(x) |
| 345 | |
| 346 | # # -------- encoder then mask ---------- |
| 347 | # # apply Transformer blocks |
| 348 | # for blk in self.blocks: |
| 349 | # x_ind = blk(x_ind) |
| 350 | # x = self.norm(x_ind) |
| 351 | # # compress for communication |
| 352 | # # mask ONLY for transmission, encode the complete sequence |
| 353 | # x_masked, mask, ids_restore = self.masking_handle(x, mask_ratio) |
| 354 | # x = self.compressor(x_masked) |
| 355 | # # -------------------------------------- |
| 356 | |
| 357 | return x, mask, ids_restore |
| 358 | |
| 359 | def forward_decoder(self, latent, mask, ids_restore): |
| 360 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected