""" model.py -- Part B detector and the loss-function ablation. Architecture ------------ Faster R-CNN with a ResNet-50 FPN backbone (torchvision), head replaced by a 2-class predictor (background + polyp). Faster R-CNN is a reasonable default for Kvasir-SEG: polyps are large relative to typical COCO objects, there is only one class, and the two-stage design gives cleanly separable classification and regression losses -- which is exactly what the loss ablation needs to modify. Loss ablation (assignment 3.3.3) -------------------------------- `--loss ce` : the stock torchvision loss L = CE(proposal classification) + smooth-L1(box regression) Proposals are sampled 512 per image with a target 25% positive fraction, so the classification term is dominated by easy background proposals. `--loss focal` : replaces ONLY the ROI-head classification term with focal loss FL(p_t) = -alpha_t (1 - p_t)^gamma log(p_t) (Lin et al., ICCV 2017). Down-weights the easy negatives that dominate the proposal set so the gradient signal concentrates on hard/ambiguous proposals -- the standard remedy for the foreground/background imbalance quantified in pre_process.py. `--loss weighted_ce` : keeps cross-entropy but applies a class weight to the foreground logit, a cheaper alternative to focal loss. The swap is done by rebinding `torchvision.models.detection.roi_heads.fastrcnn_loss`, because `RoIHeads.forward` resolves that name at call time. Everything else (anchors, sampling, NMS, box loss) is byte-for-byte identical between the two arms, so the comparison isolates the loss function. """ from __future__ import annotations from typing import List, Optional, Tuple import torch import torch.nn.functional as F from common import NUM_CLASSES # --------------------------------------------------------------------------- # Loss variants # --------------------------------------------------------------------------- _FOCAL_ALPHA = 0.25 _FOCAL_GAMMA = 2.0 _CLS_WEIGHT: Optional[torch.Tensor] = None _ORIGINAL_FASTRCNN_LOSS = None def _box_loss(box_regression, labels, regression_targets): """Identical to torchvision's box term -- kept separate so the ablation only ever changes the classification term.""" labels = torch.cat(labels, dim=0) regression_targets = torch.cat(regression_targets, dim=0) sampled_pos_inds = torch.where(labels > 0)[0] labels_pos = labels[sampled_pos_inds] N, num_classes = box_regression.shape[0], box_regression.shape[1] // 4 box_regression = box_regression.reshape(N, num_classes, 4) box_loss = F.smooth_l1_loss( box_regression[sampled_pos_inds, labels_pos], regression_targets[sampled_pos_inds], beta=1 / 9, reduction="sum") / max(1, labels.numel()) return box_loss def fastrcnn_loss_focal(class_logits, box_regression, labels, regression_targets): labels_cat = torch.cat(labels, dim=0) ce = F.cross_entropy(class_logits, labels_cat, reduction="none") pt = torch.exp(-ce) # = p_t alpha_t = torch.where(labels_cat > 0, torch.full_like(pt, _FOCAL_ALPHA), torch.full_like(pt, 1.0 - _FOCAL_ALPHA)) classification_loss = (alpha_t * (1.0 - pt) ** _FOCAL_GAMMA * ce).mean() return classification_loss, _box_loss(box_regression, labels, regression_targets) def fastrcnn_loss_weighted(class_logits, box_regression, labels, regression_targets): labels_cat = torch.cat(labels, dim=0) w = _CLS_WEIGHT.to(class_logits.device) if _CLS_WEIGHT is not None else None classification_loss = F.cross_entropy(class_logits, labels_cat, weight=w) return classification_loss, _box_loss(box_regression, labels, regression_targets) def set_roi_loss(kind: str = "ce", alpha: float = 0.25, gamma: float = 2.0, fg_weight: float = 4.0) -> str: """Install the requested ROI-head classification loss. Returns a description string that is stored in the checkpoint and printed in the report.""" global _FOCAL_ALPHA, _FOCAL_GAMMA, _CLS_WEIGHT, _ORIGINAL_FASTRCNN_LOSS from torchvision.models.detection import roi_heads if _ORIGINAL_FASTRCNN_LOSS is None: _ORIGINAL_FASTRCNN_LOSS = roi_heads.fastrcnn_loss kind = kind.lower() if kind == "ce": roi_heads.fastrcnn_loss = _ORIGINAL_FASTRCNN_LOSS return "torchvision default: cross-entropy (ROI cls) + smooth-L1 (ROI box)" if kind == "focal": _FOCAL_ALPHA, _FOCAL_GAMMA = alpha, gamma roi_heads.fastrcnn_loss = fastrcnn_loss_focal return (f"focal loss on ROI classification (alpha={alpha}, gamma={gamma}) " f"+ unchanged smooth-L1 box loss") if kind == "weighted_ce": _CLS_WEIGHT = torch.tensor([1.0, fg_weight], dtype=torch.float32) roi_heads.fastrcnn_loss = fastrcnn_loss_weighted return (f"class-weighted cross-entropy on ROI classification " f"(w_bg=1.0, w_polyp={fg_weight}) + unchanged smooth-L1 box loss") raise ValueError(f"unknown loss '{kind}' (expected ce | focal | weighted_ce)") # --------------------------------------------------------------------------- # Detector # --------------------------------------------------------------------------- def build_detector(arch: str = "fasterrcnn_resnet50_fpn", num_classes: int = NUM_CLASSES, pretrained_backbone: bool = True, img_size: int = 512, trainable_backbone_layers: int = 3, box_score_thresh: float = 0.05, box_nms_thresh: float = 0.5, box_detections_per_img: int = 20, anchor_sizes: Optional[Tuple] = None): """Faster R-CNN / RetinaNet builder. box_score_thresh is kept LOW (0.05) at inference so that the precision-recall curve and mAP are computed over the full score range; the reported operating point uses a separate, tuned confidence threshold at evaluation time. """ from torchvision.models import ResNet50_Weights from torchvision.models.detection import (FasterRCNN_ResNet50_FPN_Weights, RetinaNet_ResNet50_FPN_Weights, fasterrcnn_resnet50_fpn, retinanet_resnet50_fpn) from torchvision.models.detection.faster_rcnn import FastRCNNPredictor arch = arch.lower() # NOTE: torchvision defaults `weights_backbone` to ImageNet even when # `weights=None`, so it must be silenced explicitly for a true from-scratch run. wb = ResNet50_Weights.IMAGENET1K_V1 if pretrained_backbone else None if arch.startswith("fasterrcnn"): weights = FasterRCNN_ResNet50_FPN_Weights.COCO_V1 if pretrained_backbone else None model = fasterrcnn_resnet50_fpn( weights=weights, weights_backbone=wb, trainable_backbone_layers=trainable_backbone_layers, min_size=img_size, max_size=img_size, box_score_thresh=box_score_thresh, box_nms_thresh=box_nms_thresh, box_detections_per_img=box_detections_per_img) in_feat = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_feat, num_classes) return model if arch.startswith("retinanet"): weights = RetinaNet_ResNet50_FPN_Weights.COCO_V1 if pretrained_backbone else None model = retinanet_resnet50_fpn( weights=weights, weights_backbone=wb, num_classes=None if weights else num_classes, trainable_backbone_layers=trainable_backbone_layers, min_size=img_size, max_size=img_size, score_thresh=box_score_thresh, nms_thresh=box_nms_thresh, detections_per_img=box_detections_per_img) if weights is not None: from torchvision.models.detection.retinanet import RetinaNetClassificationHead num_anchors = model.head.classification_head.num_anchors model.head.classification_head = RetinaNetClassificationHead( in_channels=256, num_anchors=num_anchors, num_classes=num_classes) return model raise ValueError(f"unknown arch '{arch}'") def parameter_report(model) -> dict: tot = sum(p.numel() for p in model.parameters()) tr = sum(p.numel() for p in model.parameters() if p.requires_grad) return {"total_params": tot, "trainable_params": tr, "frozen_params": tot - tr} if __name__ == "__main__": print(set_roi_loss("focal", alpha=0.25, gamma=2.0)) print(set_roi_loss("ce"))