| 22 | |
| 23 | |
| 24 | class ModelWrapperOverfit(LightningModule): |
| 25 | def __init__( |
| 26 | self, |
| 27 | cfg: ModelWrapperOverfitCfg, |
| 28 | model: Model, |
| 29 | batch: Batch, |
| 30 | flows: Flows, |
| 31 | tracks: list[Tracks] | None, |
| 32 | losses: list[Loss], |
| 33 | visualizers: list[Visualizer], |
| 34 | ) -> None: |
| 35 | super().__init__() |
| 36 | self.cfg = cfg |
| 37 | self.batch = batch |
| 38 | self.flows = flows |
| 39 | self.tracks = tracks |
| 40 | self.model = model |
| 41 | self.losses = losses |
| 42 | self.visualizers = visualizers |
| 43 | |
| 44 | def to(self, device: torch.device) -> None: |
| 45 | self.batch = self.batch.to(device) |
| 46 | self.flows = self.flows.to(device) |
| 47 | if self.tracks is not None: |
| 48 | self.tracks = [tracks.to(device) for tracks in self.tracks] |
| 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. |
| 77 | model_output = self.model(self.batch, self.flows, self.global_step) |
| 78 | |
| 79 | # Generate visualizations. |
| 80 | for visualizer in self.visualizers: |
| 81 | visualizations = visualizer.visualize( |