x_next: [bxa, ts-1, C, H, W] beq_next_frames
(self, x1, x_next, mask_ratio)
| 140 | return fused_latent |
| 141 | |
| 142 | def forward_encoder(self, x1, x_next, mask_ratio): |
| 143 | """ |
| 144 | x_next: [bxa, ts-1, C, H, W] beq_next_frames |
| 145 | """ |
| 146 | # embed patches |
| 147 | x1 = self.patch_embed(x1) |
| 148 | # add pos embed w/o cls token |
| 149 | x1 = x1 + self.pos_embed[:, 1:, :] + self.temp_embed[:, 0, :] |
| 150 | |
| 151 | # handles other time stamps |
| 152 | xs = [] |
| 153 | for ts in range(self.time_stamp-1): |
| 154 | xt = x_next[:, ts, :, :, :] # [Bxa, C, H, W] |
| 155 | # print("xt size", xt.size()) |
| 156 | xt = self.patch_embed(xt) + self.pos_embed[:, 1:, :] + self.temp_embed[:, ts+1, :] |
| 157 | xs.append(xt) |
| 158 | |
| 159 | # masking: length -> length * mask_ratio |
| 160 | # x_masked, x1len, mask1, ids_restore1 = self.more_random_masking(x1, xs, mask_ratio) |
| 161 | # complement masking |
| 162 | # x_masked, x1len, mask1, ids_restore1 = self.complement_masking(x1, xs, mask_ratio) |
| 163 | x_masked, x1len, mask1, ids_restore1 = self.masking_handle(x1, xs, mask_ratio) |
| 164 | # print(x_masked.size()) |
| 165 | |
| 166 | # append cls token |
| 167 | cls_token = self.cls_token + self.pos_embed[:, :1, :] |
| 168 | cls_tokens = cls_token.expand(x_masked.shape[0], -1, -1) |
| 169 | x = torch.cat((cls_tokens, x_masked), dim=1) |
| 170 | |
| 171 | # apply Transformer blocks |
| 172 | for blk in self.blocks: |
| 173 | x = blk(x) |
| 174 | x = self.norm(x) |
| 175 | # compress for communication |
| 176 | x = self.compressor(x) |
| 177 | |
| 178 | return x, mask1, ids_restore1, x1len |
| 179 | |
| 180 | def forward_decoder(self, x): |
| 181 | """ |