from typing import Any, Optional import numpy as np import torch from torch import nn, optim import lightning.pytorch as pl import torchvision.models.video as tvmv import sklearn.metrics as skm class SyntaxLightningModule(pl.LightningModule): """ LightningModule for training the 3D backbone on SYNTAX score. Architecture: - backbone: ResNet3D (r3d_18) from torchvision - output linear layer with two units [0] binary classification logit for significant disease [1] regression output for SYNTAX score (log1p) Training modes: - pretrain (weight_path is None): freeze the backbone and train only the final fc layer - finetune (weight_path is set): load checkpoint weights and fine-tune the full network """ def __init__( self, num_classes: int, lr: float, weight_decay: float = 0.0, max_epochs: Optional[int] = None, weight_path: Optional[str] = None, sigma_a: float = 0.0, sigma_b: float = 1.0, **kwargs, ): super().__init__() self.save_hyperparameters() self.num_classes = int(num_classes) self.lr = float(lr) self.weight_decay = float(weight_decay) self.max_epochs = max_epochs self.weight_path = weight_path self.sigma_a = float(sigma_a) self.sigma_b = float(sigma_b) self.model = tvmv.r3d_18(weights=tvmv.R3D_18_Weights.DEFAULT) in_features = self.model.fc.in_features self.model.fc = nn.Linear(in_features=in_features, out_features=self.num_classes, bias=True) if self.weight_path is not None: self._load_backbone_weights(self.weight_path) self.loss_clf = nn.BCEWithLogitsLoss(reduction="none") self.loss_reg = nn.MSELoss(reduction="none") self._y_true = [] self._y_prob = [] self._y_pred = [] self._t_true = [] self._t_pred = [] def _load_backbone_weights(self, weight_path: str) -> None: """ Load backbone weights from either: - a Lightning checkpoint (dict with a 'state_dict' key) - or a raw state_dict saved via model.state_dict() (.pt/.pth) Prints the source type and key statistics. """ obj = torch.load(weight_path, map_location="cpu", weights_only=False) if isinstance(obj, dict) and "state_dict" in obj: state_dict = obj["state_dict"] state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()} src_type = "lightning_checkpoint" else: state_dict = obj src_type = "raw_state_dict" incompatible = self.model.load_state_dict(state_dict, strict=False) loaded_keys = [k for k in state_dict.keys() if k not in incompatible.missing_keys] print( f"[Backbone] Loaded weights from '{weight_path}' " f"(type={src_type}): {len(loaded_keys)} params, " f"missing={len(incompatible.missing_keys)}, " f"unexpected={len(incompatible.unexpected_keys)}" ) if incompatible.missing_keys: print(f"[Backbone] Missing keys example: {incompatible.missing_keys[:5]}") if incompatible.unexpected_keys: print(f"[Backbone] Unexpected keys example: {incompatible.unexpected_keys[:5]}") def forward(self, x: torch.Tensor) -> torch.Tensor: """ Input: x: (B, C, T, H, W) Output: y_hat: (B, 2) — [clf_logit, reg_output] """ return self.model(x) def training_step(self, batch: Any, batch_idx: int) -> torch.Tensor: """ One backbone training step. """ x, y, target, sample_weight, path, original_label = batch y_hat = self(x) yp_clf = y_hat[:, 0:1] yp_reg = y_hat[:, 1:2] weights_clf = torch.where(y > 0, 1.0, 0.45).to(y.dtype) clf_loss = (self.loss_clf(yp_clf, y) * weights_clf).mean() reg_loss_raw = self.loss_reg(yp_reg, target) sigma = self.sigma_a * target + self.sigma_b reg_loss = (reg_loss_raw / (sigma ** 2)).mean() loss = clf_loss + 0.5 * reg_loss y_prob = torch.sigmoid(yp_clf).detach() y_pred = (y_prob > 0.5).int().cpu().numpy() y_true = y.detach().int().cpu().numpy() self.log("train_clf_loss", clf_loss, prog_bar=True, sync_dist=True) self.log("train_reg_loss", reg_loss, prog_bar=True, sync_dist=True) self.log("train_loss", loss, prog_bar=True, sync_dist=True) self.log("train_f1", skm.f1_score(y_true, y_pred, zero_division=0), prog_bar=True, sync_dist=True) self.log("train_acc", skm.accuracy_score(y_true, y_pred), prog_bar=True, sync_dist=True) return loss def validation_step(self, batch: Any, batch_idx: int) -> torch.Tensor: """ One backbone validation step. """ x, y, target, sample_weight, path, original_label = batch y_hat = self(x) yp_clf = y_hat[:, 0:1] yp_reg = y_hat[:, 1:2] clf_loss = self.loss_clf(yp_clf, y).mean() reg_loss_raw = self.loss_reg(yp_reg, target) sigma = self.sigma_a * target + self.sigma_b reg_loss = (reg_loss_raw / (sigma ** 2)).mean() loss = clf_loss + 0.5 * reg_loss y_prob = torch.sigmoid(yp_clf).float() self._y_true.append(float(y[..., 0].float().cpu())) self._y_prob.append(float(y_prob[..., 0].cpu())) self._y_pred.append(int((y_prob[..., 0] > 0.5).cpu())) self._t_true.append(float(target[..., 0].float().cpu())) self._t_pred.append(float(yp_reg[..., 0].cpu())) self.log("val_loss", loss, prog_bar=True, sync_dist=True) self.log("val_clf_loss", clf_loss, prog_bar=False, sync_dist=True) self.log("val_reg_loss", reg_loss, prog_bar=False, sync_dist=True) return loss def on_validation_epoch_end(self) -> None: """ Compute and log metrics at the end of validation. """ if len(self._t_true) > 0: rmse = skm.root_mean_squared_error(self._t_true, self._t_pred) mae = skm.mean_absolute_error(self._t_true, self._t_pred) self.log("val_rmse", rmse, prog_bar=True, sync_dist=True) self.log("val_reg_mae", mae, prog_bar=True, sync_dist=True) if len(set(self._y_true)) > 1: auc = skm.roc_auc_score(self._y_true, self._y_prob) f1 = skm.f1_score(self._y_true, self._y_pred, zero_division=0) acc = skm.accuracy_score(self._y_true, self._y_pred) self.log("val_auc", auc, prog_bar=True, sync_dist=True) self.log("val_f1", f1, prog_bar=True, sync_dist=True) self.log("val_acc", acc, prog_bar=True, sync_dist=True) self._y_true.clear() self._y_prob.clear() self._y_pred.clear() self._t_true.clear() self._t_pred.clear() def on_train_epoch_end(self) -> None: """ Log the current learning rate. """ opt = self.optimizers() self.log( "lr", opt.optimizer.param_groups[0]["lr"], on_step=False, on_epoch=True, sync_dist=True, ) def configure_optimizers(self): """ Configure the optimizer and OneCycleLR. If weight_path is None: train only self.model.fc (pretrain). Otherwise: train the full model (fine-tuning). """ if self.weight_path is None: for p in self.parameters(): p.requires_grad = False for p in self.model.fc.parameters(): p.requires_grad = True params = self.model.fc.parameters() else: for p in self.parameters(): p.requires_grad = True params = self.parameters() optimizer = optim.AdamW(params, lr=self.lr, weight_decay=self.weight_decay) if self.max_epochs is not None and getattr(self, "trainer", None) is not None: total_steps = self.trainer.estimated_stepping_batches scheduler = optim.lr_scheduler.OneCycleLR( optimizer=optimizer, max_lr=self.lr, total_steps=total_steps, ) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", }, } return optimizer def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any: """ Backbone inference step. """ x, y, target, sample_weight, path, original_label = batch y_hat = self(x) yp_clf = y_hat[:, 0:1] yp_reg = y_hat[:, 1:2] y_prob = torch.sigmoid(yp_clf) return { "y": y, "y_pred": (y_prob > 0.5).int(), "y_prob": y_prob, "y_reg": yp_reg, "target": target, "original_label": original_label, "path": path, }