| 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'] |
| 160 | trans_matrices = data['trans_matrices'] |
| 161 | num_agent = data['num_agent'] |
| 162 | |
| 163 | vq_loss, result, ind_recon, perplexity = self.model(bev_seq, trans_matrices, num_agent, batch_size=batch_size) |
| 164 | target = bev_seq.squeeze(1).permute(0,3,1,2).detach() |
| 165 | # print("target", target.size()) |
| 166 | # recon_error = F.mse_loss(ind_recon, target) |
| 167 | ## classification task ## |
| 168 | # target = bev_seq.permute(0, 1, 4, 2, 3).squeeze(1) |
| 169 | target = target.type(torch.LongTensor).to(ind_recon.device) |
| 170 | # print("target", target.size()) |
| 171 | # print("ind_recon", ind_recon.size(), ind_recon.type()) |
| 172 | loss_fn = nn.CrossEntropyLoss() |
| 173 | recon_error = loss_fn(ind_recon, target) |
| 174 | loss = recon_error + vq_loss |
| 175 | |
| 176 | if trainable: |
| 177 | self.optimizer.zero_grad() |
| 178 | loss.backward() |
| 179 | self.optimizer.step() |
| 180 | |
| 181 | return loss.item(), result, ind_recon, perplexity.detach() |
| 182 | |
| 183 | def infer_vae_completion(self, data, batch_size, loss_fn='ce', trainable=False): |
| 184 | bev_seq = data['bev_seq'] |