| """ |
| 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 |
|
|
| |
| |
| |
|
|
| _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) |
| 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)") |
|
|
|
|
| |
| |
| |
|
|
| 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() |
| |
| |
| 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")) |
|
|
