Spaces:
Running
Running
| import pytorch_lightning as pl | |
| import torch | |
| import torch.nn as nn | |
| import timm | |
| from torchmetrics import Accuracy | |
| from utils.focal_loss import FocalLoss | |
| class MorphDetectorPL(pl.LightningModule): | |
| """ | |
| PyTorch Lightning module for Morph Detection. | |
| Dual-stream: Vision Transformer + EfficientNet backbone, fused with shallow classifier head. | |
| Uses FocalLoss for hard example mining. | |
| """ | |
| def __init__( | |
| self, | |
| vit_model_name: str = 'vit_base_patch16_224', | |
| effnet_model_name: str = 'efficientnet_b0', | |
| pretrained: bool = True, | |
| lr: float = 1e-4, | |
| dropout: float = 0.2, | |
| focal_alpha: float = 0.25, | |
| focal_gamma: float = 2.0, | |
| ): | |
| super().__init__() | |
| # save hyperparameters | |
| self.save_hyperparameters() | |
| # Global feature extractor | |
| self.vit = timm.create_model( | |
| vit_model_name, | |
| pretrained=pretrained, | |
| num_classes=0, | |
| ) | |
| # Local feature extractor | |
| self.effnet = timm.create_model( | |
| effnet_model_name, | |
| pretrained=pretrained, | |
| num_classes=0, | |
| ) | |
| # Fusion classifier head | |
| fused_dim = self.vit.num_features + self.effnet.num_features | |
| self.classifier = nn.Sequential( | |
| nn.Linear(fused_dim, 256), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(256, 1), | |
| nn.Sigmoid(), | |
| ) | |
| # Loss and metrics | |
| self.criterion = FocalLoss(alpha=focal_alpha, gamma=focal_gamma) | |
| self.train_acc = Accuracy() | |
| self.val_acc = Accuracy() | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| # x: (B, C, H, W) | |
| feat1 = self.vit(x) | |
| feat2 = self.effnet(x) | |
| fused = torch.cat([feat1, feat2], dim=1) | |
| return self.classifier(fused) | |
| def training_step(self, batch, batch_idx): | |
| imgs, labels = batch | |
| labels = labels.float().unsqueeze(1) | |
| logits = self(imgs) | |
| loss = self.criterion(logits, labels) | |
| preds = (logits > 0.5).int() | |
| self.train_acc.update(preds, labels.int()) | |
| self.log('train_loss', loss, on_step=True, on_epoch=True) | |
| self.log('train_acc', self.train_acc, on_step=False, on_epoch=True) | |
| return loss | |
| def validation_step(self, batch, batch_idx): | |
| imgs, labels = batch | |
| labels = labels.float().unsqueeze(1) | |
| logits = self(imgs) | |
| loss = self.criterion(logits, labels) | |
| preds = (logits > 0.5).int() | |
| self.val_acc.update(preds, labels.int()) | |
| self.log('val_loss', loss, prog_bar=True) | |
| self.log('val_acc', self.val_acc, prog_bar=True) | |
| def configure_optimizers(self): | |
| optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr) | |
| scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( | |
| optimizer, mode='min', factor=0.5, patience=3 | |
| ) | |
| return { | |
| 'optimizer': optimizer, | |
| 'lr_scheduler': { | |
| 'scheduler': scheduler, | |
| 'monitor': 'val_loss', | |
| } | |
| } |