| 89 | |
| 90 | # Used my MAE model in scene completion, added by Juexiao |
| 91 | def step_mae_completion(self, data, batch_size, mask_ratio, loss_fn='mse', trainable = False): |
| 92 | # figure out the dimensions: squeeze, permute etc |
| 93 | bev_seq = data['bev_seq'].squeeze(1).permute(0,3,1,2) |
| 94 | bev_seq_next = data['bev_seq_next'] # [bxa, ts-1, H, W, C] |
| 95 | if bev_seq_next.dim()>2: |
| 96 | bev_seq_next = bev_seq_next.permute(0,1,4,2,3) # [bxa, ts-1, C, H, W] |
| 97 | # print(data['bev_seq_teacher'].size()) |
| 98 | bev_teacher = data['bev_seq_teacher'].squeeze(1).permute(0,3,1,2) |
| 99 | # print("teacher size", bev_teacher.size()) |
| 100 | num_agent_tensor = data['num_agent'] |
| 101 | trans_matrices = data['trans_matrices'] |
| 102 | |
| 103 | with torch.cuda.amp.autocast(enabled=False): |
| 104 | loss, result, _, ind_pred = self.model(bev_seq, |
| 105 | bev_seq_next, |
| 106 | bev_teacher, |
| 107 | trans_matrices, |
| 108 | num_agent_tensor, |
| 109 | batch_size, |
| 110 | mask_ratio=mask_ratio) |
| 111 | # print("result size", result.size()) |
| 112 | # exit(1) |
| 113 | |
| 114 | loss_value = loss.item() |
| 115 | |
| 116 | if not math.isfinite(loss_value): |
| 117 | print("Loss is {}, stopping training".format(loss_value)) |
| 118 | # torch.save(bev_seq.cpu(), "lossnan-bev.pt") |
| 119 | # torch.save(self.model.cpu(), "lossnan-mae.pt") |
| 120 | # torch.save(result.cpu(), "lossnan-pred.pt") |
| 121 | # torch.save(mask.cpu(), "lossnan-mask.pt") |
| 122 | sys.exit(1) |
| 123 | |
| 124 | if trainable: |
| 125 | self.mae_loss_scaler(loss, self.optimizer, parameters=self.model.parameters(), |
| 126 | update_grad=True) |
| 127 | self.optimizer.zero_grad() |
| 128 | # loss.backward() |
| 129 | # self.optimizer.step() |
| 130 | |
| 131 | return loss_value, result, ind_pred.detach() |
| 132 | |
| 133 | # inference mae, ego use groundtruth |
| 134 | def infer_mae_completion(self, data, batch_size, mask_ratio, loss_fn='mse'): |