| 191 | |
| 192 | # Used by VQ STAR model in scene completion, added by Juexiao |
| 193 | def step_vqstar_completion(self, data, batch_size, mask_ratio, loss_fn='mse', trainable = False): |
| 194 | # figure out the dimensions: squeeze, permute etc |
| 195 | bev_seq = data['bev_seq'].squeeze(1).permute(0,3,1,2) |
| 196 | bev_seq_next = data['bev_seq_next'] # [bxa, ts-1, H, W, C] |
| 197 | if bev_seq_next.dim()>2: |
| 198 | bev_seq_next = bev_seq_next.permute(0,1,4,2,3) # [bxa, ts-1, C, H, W] |
| 199 | # print(data['bev_seq_teacher'].size()) |
| 200 | bev_teacher = data['bev_seq_teacher'].squeeze(1).permute(0,3,1,2) |
| 201 | # print("teacher size", bev_teacher.size()) |
| 202 | num_agent_tensor = data['num_agent'] |
| 203 | trans_matrices = data['trans_matrices'] |
| 204 | |
| 205 | with torch.cuda.amp.autocast(enabled=False): |
| 206 | loss, result, _, ind_pred, perplexity = self.model(bev_seq, |
| 207 | bev_seq_next, |
| 208 | bev_teacher, |
| 209 | trans_matrices, |
| 210 | num_agent_tensor, |
| 211 | batch_size, |
| 212 | mask_ratio=mask_ratio) |
| 213 | # print("result size", result.size()) |
| 214 | # exit(1) |
| 215 | |
| 216 | loss_value = loss.item() |
| 217 | |
| 218 | if not math.isfinite(loss_value): |
| 219 | print("Loss is {}, stopping training".format(loss_value)) |
| 220 | sys.exit(1) |
| 221 | |
| 222 | if trainable: |
| 223 | self.mae_loss_scaler(loss, self.optimizer, parameters=self.model.parameters(), |
| 224 | update_grad=True) |
| 225 | self.optimizer.zero_grad() |
| 226 | # loss.backward() |
| 227 | # self.optimizer.step() |
| 228 | |
| 229 | return loss_value, result, ind_pred.detach(), perplexity.detach() |
| 230 | |
| 231 | def infer_vqstar_completion(self, data, batch_size, mask_ratio, loss_fn='mse'): |
| 232 | # figure out the dimensions: squeeze, permute etc |