| 49 | super().to(device) |
| 50 | |
| 51 | def training_step(self, dummy): |
| 52 | # Compute depths, poses, and intrinsics using the model. |
| 53 | model_output = self.model(self.batch, self.flows, self.global_step) |
| 54 | |
| 55 | # Compute and log the loss. |
| 56 | total_loss = 0 |
| 57 | for loss_fn in self.losses: |
| 58 | loss = loss_fn.forward( |
| 59 | self.batch, self.flows, self.tracks, model_output, self.global_step |
| 60 | ) |
| 61 | self.log(f"train/loss/{loss_fn.cfg.name}", loss) |
| 62 | total_loss = total_loss + loss |
| 63 | |
| 64 | # Log intrinsics error. |
| 65 | if self.batch.intrinsics is not None: |
| 66 | fx_hat = reduce(model_output.intrinsics[..., 0, 0], "b f ->", "mean") |
| 67 | fy_hat = reduce(model_output.intrinsics[..., 1, 1], "b f ->", "mean") |
| 68 | fx_gt = reduce(self.batch.intrinsics[..., 0, 0], "b f ->", "mean") |
| 69 | fy_gt = reduce(self.batch.intrinsics[..., 1, 1], "b f ->", "mean") |
| 70 | self.log("train/intrinsics/fx_error", (fx_gt - fx_hat).abs()) |
| 71 | self.log("train/intrinsics/fy_error", (fy_gt - fy_hat).abs()) |
| 72 | |
| 73 | return total_loss |
| 74 | |
| 75 | def validation_step(self, dummy): |
| 76 | # Compute depths, poses, and intrinsics using the model. |