# Copyright (c) Facebook, Inc. and its affiliates.
import logging
import math
from typing import List, Tuple
import torch
from fvcore.nn import sigmoid_focal_loss_jit
from torch import Tensor, nn
from torch.nn import functional as F
from detectron2.config import configurable
from detectron2.layers import CycleBatchNormList, ShapeSpec, batched_nms, cat, get_norm
from detectron2.structures import Boxes, ImageList, Instances, pairwise_iou
from detectron2.utils.events import get_event_storage
from ..anchor_generator import build_anchor_generator
from ..backbone import Backbone, build_backbone
from ..box_regression import Box2BoxTransform, _dense_box_regression_loss
from ..matcher import Matcher
from .build import META_ARCH_REGISTRY
from .dense_detector import DenseDetector, permute_to_N_HWA_K # noqa
__all__ = ["RetinaNet"]
logger = logging.getLogger(__name__)
@META_ARCH_REGISTRY.register()
class RetinaNet(DenseDetector):
"""
Implement RetinaNet in :paper:`RetinaNet`.
"""
@configurable
def __init__(
self,
*,
backbone: Backbone,
head: nn.Module,
head_in_features,
anchor_generator,
box2box_transform,
anchor_matcher,
num_classes,
focal_loss_alpha=0.25,
focal_loss_gamma=2.0,
smooth_l1_beta=0.0,
box_reg_loss_type="smooth_l1",
test_score_thresh=0.05,
test_topk_candidates=1000,
test_nms_thresh=0.5,
max_detections_per_image=100,
pixel_mean,
pixel_std,
vis_period=0,
input_format="BGR",
):
"""
NOTE: this interface is experimental.
Args:
backbone: a backbone module, must follow detectron2's backbone interface
head (nn.Module): a module that predicts logits and regression deltas
for each level from a list of per-level features
head_in_features (Tuple[str]): Names of the input feature maps to be used in head
anchor_generator (nn.Module): a module that creates anchors from a
list of features. Usually an instance of :class:`AnchorGenerator`
box2box_transform (Box2BoxTransform): defines the transform from anchors boxes to
instance boxes
anchor_matcher (Matcher): label the anchors by matching them with ground truth.
num_classes (int): number of classes. Used to label background proposals.
# Loss parameters:
focal_loss_alpha (float): focal_loss_alpha
focal_loss_gamma (float): focal_loss_gamma
smooth_l1_beta (float): smooth_l1_beta
box_reg_loss_type (str): Options are "smooth_l1", "giou", "diou", "ciou"
# Inference parameters:
test_score_thresh (float): Inference cls score threshold, only anchors with
score > INFERENCE_TH are considered for inference (to improve speed)
test_topk_candidates (int): Select topk candidates before NMS
test_nms_thresh (float): Overlap threshold used for non-maximum suppression
(suppress boxes with IoU >= this threshold)
max_detections_per_image (int):
Maximum number of detections to return per image during inference
(100 is based on the limit established for the COCO dataset).
pixel_mean, pixel_std: see :class:`DenseDetector`.
"""
super().__init__(
backbone, head, head_in_features, pixel_mean=pixel_mean, pixel_std=pixel_std
)
self.num_classes = num_classes
# Anchors
self.anchor_generator = anchor_generator
self.box2box_transform = box2box_transform
self.anchor_matcher = anchor_matcher
# Loss parameters:
self.focal_loss_alpha = focal_loss_alpha
self.focal_loss_gamma = focal_loss_gamma
self.smooth_l1_beta = smooth_l1_beta
self.box_reg_loss_type = box_reg_loss_type
# Inference parameters:
self.test_score_thresh = test_score_thresh
self.test_topk_candidates = test_topk_candidates
self.test_nms_thresh = test_nms_thresh
self.max_detections_per_image = max_detections_per_image
# Vis parameters
self.vis_period = vis_period
self.input_format = input_format
@classmethod
def from_config(cls, cfg):
backbone = build_backbone(cfg)
backbone_shape = backbone.output_shape()
feature_shapes = [backbone_shape[f] for f in cfg.MODEL.RETINANET.IN_FEATURES]
head = RetinaNetHead(cfg, feature_shapes)
anchor_generator = build_anchor_generator(cfg, feature_shapes)
return {
"backbone": backbone,
"head": head,
"anchor_generator": anchor_generator,
"box2box_transform": Box2BoxTransform(weights=cfg.MODEL.RETINANET.BBOX_REG_WEIGHTS),
"anchor_matcher": Matcher(
cfg.MODEL.RETINANET.IOU_THRESHOLDS,
cfg.MODEL.RETINANET.IOU_LABELS,
allow_low_quality_matches=True,
),
"pixel_mean": cfg.MODEL.PIXEL_MEAN,
"pixel_std": cfg.MODEL.PIXEL_STD,
"num_classes": cfg.MODEL.RETINANET.NUM_CLASSES,
"head_in_features": cfg.MODEL.RETINANET.IN_FEATURES,
# Loss parameters:
"focal_loss_alpha": cfg.MODEL.RETINANET.FOCAL_LOSS_ALPHA,
"focal_loss_gamma": cfg.MODEL.RETINANET.FOCAL_LOSS_GAMMA,
"smooth_l1_beta": cfg.MODEL.RETINANET.SMOOTH_L1_LOSS_BETA,
"box_reg_loss_type": cfg.MODEL.RETINANET.BBOX_REG_LOSS_TYPE,
# Inference parameters:
"test_score_thresh": cfg.MODEL.RETINANET.SCORE_THRESH_TEST,
"test_topk_candidates": cfg.MODEL.RETINANET.TOPK_CANDIDATES_TEST,
"test_nms_thresh": cfg.MODEL.RETINANET.NMS_THRESH_TEST,
"max_detections_per_image": cfg.TEST.DETECTIONS_PER_IMAGE,
# Vis parameters
"vis_period": cfg.VIS_PERIOD,
"input_format": cfg.INPUT.FORMAT,
}
def forward_training(self, images, features, predictions, gt_instances):
# Transpose the Hi*Wi*A dimension to the middle:
pred_logits, pred_anchor_deltas = self._transpose_dense_predictions(
predictions, [self.num_classes, 4]
)
anchors = self.anchor_generator(features)
gt_labels, gt_boxes = self.label_anchors(anchors, gt_instances)
return self.losses(anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes)
def losses(self, anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes):
"""
Args:
anchors (list[Boxes]): a list of #feature level Boxes
gt_labels, gt_boxes: see output of :meth:`RetinaNet.label_anchors`.
Their shapes are (N, R) and (N, R, 4), respectively, where R is
the total number of anchors across levels, i.e. sum(Hi x Wi x Ai)
pred_logits, pred_anchor_deltas: both are list[Tensor]. Each element in the
list corresponds to one level and has shape (N, Hi * Wi * Ai, K or 4).
Where K is the number of classes used in `pred_logits`.
Returns:
dict[str, Tensor]:
mapping from a named loss to a scalar tensor storing the loss.
Used during training only. The dict keys are: "loss_cls" and "loss_box_reg"
"""
num_images = len(gt_labels)
gt_labels = torch.stack(gt_labels) # (N, R)
valid_mask = gt_labels >= 0
pos_mask = (gt_labels >= 0) & (gt_labels != self.num_classes)
num_pos_anchors = pos_mask.sum().item()
get_event_storage().put_scalar("num_pos_anchors", num_pos_anchors / num_images)
normalizer = self._ema_update("loss_normalizer", max(num_pos_anchors, 1), 100)
# classification and regression loss
gt_labels_target = F.one_hot(gt_labels[valid_mask], num_classes=self.num_classes + 1)[
:, :-1
] # no loss for the last (background) class
loss_cls = sigmoid_focal_loss_jit(
cat(pred_logits, dim=1)[valid_mask],
gt_labels_target.to(pred_logits[0].dtype),
alpha=self.focal_loss_alpha,
gamma=self.focal_loss_gamma,
reduction="sum",
)
loss_box_reg = _dense_box_regression_loss(
anchors,
self.box2box_transform,
pred_anchor_deltas,
gt_boxes,
pos_mask,
box_reg_loss_type=self.box_reg_loss_type,
smooth_l1_beta=self.smooth_l1_beta,
)
return {
"loss_cls": loss_cls / normalizer,
"loss_box_reg": loss_box_reg / normalizer,
}
@torch.no_grad()
def label_anchors(self, anchors, gt_instances):
"""
Args:
anchors (list[Boxes]): A list of #feature level Boxes.
The Boxes contains anchors of this image on the specific feature level.
gt_instances (list[Instances]): a list of N `Instances`s. The i-th
`Instances` contains the ground-truth per-instance annotations
for the i-th input image.
Returns:
list[Tensor]: List of #img tensors. i-th element is a vector of labels whose length is
the total number of anchors across all feature maps (sum(Hi * Wi * A)).
Label values are in {-1, 0, ..., K}, with -1 means ignore, and K means background.
list[Tensor]: i-th element is a Rx4 tensor, where R is the total number of anchors
across feature maps. The values are the matched gt boxes for each anchor.
Values are undefined for those anchors not labeled as foreground.
"""
anchors = Boxes.cat(anchors) # Rx4
gt_labels = []
matched_gt_boxes = []
for gt_per_image in gt_instances:
match_quality_matrix = pairwise_iou(gt_per_image.gt_boxes, anchors)
matched_idxs, anchor_labels = self.anchor_matcher(match_quality_matrix)
del match_quality_matrix
if len(gt_per_image) > 0:
matched_gt_boxes_i = gt_per_image.gt_boxes.tensor[matched_idxs]
gt_labels_i = gt_per_image.gt_classes[matched_idxs]
# Anchors with label 0 are treated as background.
gt_labels_i[anchor_labels == 0] = self.num_classes
# Anchors with label -1 are ignored.
gt_labels_i[anchor_labels == -1] = -1
else:
matched_gt_boxes_i = torch.zeros_like(anchors.tensor)
gt_labels_i = torch.zeros_like(matched_idxs) + self.num_classes
gt_labels.append(gt_labels_i)
matched_gt_boxes.append(matched_gt_boxes_i)
return gt_labels, matched_gt_boxes
def forward_inference(
self, images: ImageList, features: List[Tensor], predictions: List[List[Tensor]]
):
pred_logits, pred_anchor_deltas = self._transpose_dense_predictions(
predictions, [self.num_classes, 4]
)
anchors = self.anchor_generator(features)
results: List[Instances] = []
for img_idx, image_size in enumerate(images.image_sizes):
scores_per_image = [x[img_idx].sigmoid_() for x in pred_logits]
deltas_per_image = [x[img_idx] for x in pred_anchor_deltas]
results_per_image = self.inference_single_image(
anchors, scores_per_image, deltas_per_image, image_size
)
results.append(results_per_image)
return results
def inference_single_image(
self,
anchors: List[Boxes],
box_cls: List[Tensor],
box_delta: List[Tensor],
image_size: Tuple[int, int],
):
"""
Single-image inference. Return bounding-box detection results by thresholding
on scores and applying non-maximum suppression (NMS).
Arguments:
anchors (list[Boxes]): list of #feature levels. Each entry contains
a Boxes object, which contains all the anchors in that feature level.
box_cls (list[Tensor]): list of #feature levels. Each entry contains
tensor of size (H x W x A, K)
box_delta (list[Tensor]): Same shape as 'box_cls' except that K becomes 4.
image_size (tuple(H, W)): a tuple of the image height and width.
Returns:
Same as `inference`, but for only one image.
"""
pred = self._decode_multi_level_predictions(
anchors,
box_cls,
box_delta,
self.test_score_thresh,
self.test_topk_candidates,
image_size,
)
keep = batched_nms( # per-class NMS
pred.pred_boxes.tensor, pred.scores, pred.pred_classes, self.test_nms_thresh
)
return pred[keep[: self.max_detections_per_image]]
class RetinaNetHead(nn.Module):
"""
The head used in RetinaNet for object classification and box regression.
It has two subnets for the two tasks, with a common structure but separate parameters.
"""
@configurable
def __init__(
self,
*,
input_shape: List[ShapeSpec],
num_classes,
num_anchors,
conv_dims: List[int],
norm="",
prior_prob=0.01,
):
"""
NOTE: this interface is experimental.
Args:
input_shape (List[ShapeSpec]): input shape
num_classes (int): number of classes. Used to label background proposals.
num_anchors (int): number of generated anchors
conv_dims (List[int]): dimensions for each convolution layer
norm (str or callable):
Normalization for conv layers except for the two output layers.
See :func:`detectron2.layers.get_norm` for supported types.
prior_prob (float): Prior weight for computing bias
"""
super().__init__()
self._num_features = len(input_shape)
if norm == "BN" or norm == "SyncBN":
logger.info(
f"Using domain-specific {norm} in RetinaNetHead with len={self._num_features}."
)
bn_class = nn.BatchNorm2d if norm == "BN" else nn.SyncBatchNorm
def norm(c):
return CycleBatchNormList(
length=self._num_features, bn_class=bn_class, num_features=c
)
else:
norm_name = str(type(get_norm(norm, 32)))
if "BN" in norm_name:
logger.warning(
f"Shared BatchNorm (type={norm_name}) may not work well in RetinaNetHead."
)
cls_subnet = []
bbox_subnet = []
for in_channels, out_channels in zip(
[input_shape[0].channels] + list(conv_dims), conv_dims
):
cls_subnet.append(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
)
if norm:
cls_subnet.append(get_norm(norm, out_channels))
cls_subnet.append(nn.ReLU())
bbox_subnet.append(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
)
if norm:
bbox_subnet.append(get_norm(norm, out_channels))
bbox_subnet.append(nn.ReLU())
self.cls_subnet = nn.Sequential(*cls_subnet)
self.bbox_subnet = nn.Sequential(*bbox_subnet)
self.cls_score = nn.Conv2d(
conv_dims[-1], num_anchors * num_classes, kernel_size=3, stride=1, padding=1
)
self.bbox_pred = nn.Conv2d(
conv_dims[-1], num_anchors * 4, kernel_size=3, stride=1, padding=1
)
# Initialization
for modules in [self.cls_subnet, self.bbox_subnet, self.cls_score, self.bbox_pred]:
for layer in modules.modules():
if isinstance(layer, nn.Conv2d):
torch.nn.init.normal_(layer.weight, mean=0, std=0.01)
torch.nn.init.constant_(layer.bias, 0)
# Use prior in model initialization to improve stability
bias_value = -(math.log((1 - prior_prob) / prior_prob))
torch.nn.init.constant_(self.cls_score.bias, bias_value)
@classmethod
def from_config(cls, cfg, input_shape: List[ShapeSpec]):
num_anchors = build_anchor_generator(cfg, input_shape).num_cell_anchors
assert (
len(set(num_anchors)) == 1
), "Using different number of anchors between levels is not currently supported!"
num_anchors = num_anchors[0]
return {
"input_shape": input_shape,
"num_classes": cfg.MODEL.RETINANET.NUM_CLASSES,
"conv_dims": [input_shape[0].channels] * cfg.MODEL.RETINANET.NUM_CONVS,
"prior_prob": cfg.MODEL.RETINANET.PRIOR_PROB,
"norm": cfg.MODEL.RETINANET.NORM,
"num_anchors": num_anchors,
}
def forward(self, features: List[Tensor]):
"""
Arguments:
features (list[Tensor]): FPN feature map tensors in high to low resolution.
Each tensor in the list correspond to different feature levels.
Returns:
logits (list[Tensor]): #lvl tensors, each has shape (N, AxK, Hi, Wi).
The tensor predicts the classification probability
at each spatial position for each of the A anchors and K object
classes.
bbox_reg (list[Tensor]): #lvl tensors, each has shape (N, Ax4, Hi, Wi).
The tensor predicts 4-vector (dx,dy,dw,dh) box
regression values for every anchor. These values are the
relative offset between the anchor and the ground truth box.
"""
assert len(features) == self._num_features
logits = []
bbox_reg = []
for feature in features:
logits.append(self.cls_score(self.cls_subnet(feature)))
bbox_reg.append(self.bbox_pred(self.bbox_subnet(feature)))
return logits, bbox_reg