File size: 8,865 Bytes
d0518d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """
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"))
|