| 42 | |
| 43 | # used by scene completion task |
| 44 | def step_completion(self, data, batch_size, loss_fn='ce', trainable=False): |
| 45 | bev_seq = data['bev_seq'] |
| 46 | trans_matrices = data['trans_matrices'] |
| 47 | num_agent = data['num_agent'] |
| 48 | |
| 49 | result, ind_pred = self.model(bev_seq, trans_matrices, num_agent, batch_size=batch_size) |
| 50 | |
| 51 | loss_fn_dict = { |
| 52 | 'mse': nn.MSELoss(), |
| 53 | 'bce': nn.BCELoss(), |
| 54 | 'ce': nn.CrossEntropyLoss(), |
| 55 | 'l1': nn.L1Loss(), |
| 56 | 'smooth_l1': nn.SmoothL1Loss(), |
| 57 | } |
| 58 | |
| 59 | loss = -1 |
| 60 | if trainable: |
| 61 | # labels = data['bev_seq_teacher'] |
| 62 | # labels = labels.permute(0, 1, 4, 2, 3).squeeze() # (Batch, seq, z, h, w) |
| 63 | # loss = 10000 * loss_fn_dict[loss_fn](result, labels) |
| 64 | target = bev_seq.permute(0, 1, 4, 2, 3).squeeze(1) |
| 65 | target = target.type(torch.LongTensor).to(ind_pred.device) |
| 66 | loss = loss_fn_dict[loss_fn](ind_pred, target) |
| 67 | |
| 68 | if self.MGDA: |
| 69 | self.optimizer_encoder.zero_grad() |
| 70 | self.optimizer_head.zero_grad() |
| 71 | loss.backward() |
| 72 | self.optimizer_encoder.step() |
| 73 | self.optimizer_head.step() |
| 74 | else: |
| 75 | self.optimizer.zero_grad() |
| 76 | loss.backward() |
| 77 | self.optimizer.step() |
| 78 | |
| 79 | return loss, result |
| 80 | |
| 81 | def infer_completion(self, data, batch_size): |
| 82 | bev_seq = data['bev_seq'] |