| 26 | |
| 27 | |
| 28 | class ModelWrapperPretrain(LightningModule): |
| 29 | def __init__( |
| 30 | self, |
| 31 | cfg: ModelWrapperPretrainCfg, |
| 32 | cfg_cropping: CroppingCfg, |
| 33 | cfg_flow: FlowPredictorCfg, |
| 34 | model: Model, |
| 35 | losses: list[Loss], |
| 36 | visualizers: list[Visualizer], |
| 37 | ) -> None: |
| 38 | super().__init__() |
| 39 | self.cfg = cfg |
| 40 | self.cfg_cropping = cfg_cropping |
| 41 | self.flow_predictor = get_flow_predictor(cfg_flow) |
| 42 | self.model = model |
| 43 | self.losses = losses |
| 44 | self.visualizers = visualizers |
| 45 | |
| 46 | @torch.no_grad() |
| 47 | def preprocess_batch(self, batch_dict: dict) -> tuple[Batch, Flows]: |
| 48 | # Convert the batch from an untyped dict to a typed dataclass. |
| 49 | batch_dict.pop("frame_paths", None) |
| 50 | batch = Batch(**batch_dict) |
| 51 | |
| 52 | # Compute optical flow and tracks. |
| 53 | batch_for_model, _ = crop_and_resize_batch_for_model(batch, self.cfg_cropping) |
| 54 | batch_for_flow = crop_and_resize_batch_for_flow(batch, self.cfg_cropping) |
| 55 | _, _, _, h, w = batch_for_model.videos.shape |
| 56 | flows = self.flow_predictor.compute_bidirectional_flow(batch_for_flow, (h, w)) |
| 57 | |
| 58 | return batch_for_model, flows |
| 59 | |
| 60 | def training_step(self, batch): |
| 61 | batch, flows = self.preprocess_batch(batch) |
| 62 | |
| 63 | # Compute depths, poses, and intrinsics using the model. |
| 64 | model_output = self.model(batch, flows, self.global_step) |
| 65 | |
| 66 | # Compute and log the loss. |
| 67 | total_loss = 0 |
| 68 | for loss_fn in self.losses: |
| 69 | loss = loss_fn.forward(batch, flows, None, model_output, self.global_step) |
| 70 | self.log(f"train/loss/{loss_fn.cfg.name}", loss) |
| 71 | total_loss = total_loss + loss |
| 72 | |
| 73 | # Log intrinsics error. |
| 74 | if batch.intrinsics is not None: |
| 75 | fx_hat = reduce(model_output.intrinsics[..., 0, 0], "b f ->", "mean") |
| 76 | fy_hat = reduce(model_output.intrinsics[..., 1, 1], "b f ->", "mean") |
| 77 | fx_gt = reduce(batch.intrinsics[..., 0, 0], "b f ->", "mean") |
| 78 | fy_gt = reduce(batch.intrinsics[..., 1, 1], "b f ->", "mean") |
| 79 | self.log("train/intrinsics/fx_error", (fx_gt - fx_hat).abs()) |
| 80 | self.log("train/intrinsics/fy_error", (fy_gt - fy_hat).abs()) |
| 81 | |
| 82 | return total_loss |
| 83 | |
| 84 | def validation_step(self, batch): |
| 85 | batch, flows = self.preprocess_batch(batch) |