| 132 | |
| 133 | # inference mae, ego use groundtruth |
| 134 | def infer_mae_completion(self, data, batch_size, mask_ratio, loss_fn='mse'): |
| 135 | # figure out the dimensions: squeeze, permute etc |
| 136 | bev_seq = data['bev_seq'].squeeze(1).permute(0,3,1,2) |
| 137 | bev_seq_next = data['bev_seq_next'] # [bxa, ts-1, H, W, C] |
| 138 | if bev_seq_next.dim()>2: |
| 139 | bev_seq_next = bev_seq_next.permute(0,1,4,2,3) # [bxa, ts-1, C, H, W] |
| 140 | # print(data['bev_seq_teacher'].size()) |
| 141 | bev_teacher = data['bev_seq_teacher'].squeeze(1).permute(0,3,1,2) |
| 142 | # print("teacher size", bev_teacher.size()) |
| 143 | num_agent_tensor = data['num_agent'] |
| 144 | trans_matrices = data['trans_matrices'] |
| 145 | |
| 146 | with torch.cuda.amp.autocast(enabled=False): |
| 147 | loss, result, latent, ind_pred = self.model.inference(bev_seq, |
| 148 | bev_seq_next, |
| 149 | bev_teacher, |
| 150 | trans_matrices, |
| 151 | num_agent_tensor, |
| 152 | batch_size, |
| 153 | mask_ratio=mask_ratio) |
| 154 | |
| 155 | loss_value = loss.item() |
| 156 | return loss_value, result, ind_pred.detach(), latent.detach() |
| 157 | |
| 158 | def step_vae_completion(self, data, batch_size, loss_fn='ce', trainable=False): |
| 159 | bev_seq = data['bev_seq'] |