Pytorch官方Faster R-CNN源代码解析(一)——特征提取
作者：互联网
Pytorch官方使用的示例代码如下：
import torch
import torchvision
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# For training
images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)
boxes[:, :, 2:4] = boxes[:, :, 0:2] + boxes[:, :, 2:4]
labels = torch.randint(1, 91, (4, 11))
images = list(image for image in images)
targets = []
for i in range(len(images)):
d = {
boxes: boxes[i], labels: labels[i]}
targets.append(d)
output = model(images, targets)
# For inference
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)
# optionally, if you want to export the model to ONNX:
torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)
下面主要就示例代码进行详细说明。
首先，初始化 Faster R-CNN 模型。
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
可以看出，这里使用的是主干网络 Resnet-50-FPN 的 Faster R-CNN。接下来 Debug 进内部代码。
def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
num_classes=91, pretrained_backbone=True, trainable_backbone_layers=3, **kwargs):
"""
Constructs a Faster R-CNN model with a ResNet-50-FPN backbone.
构建一个主干网络为 ResNet-50-FPN 的 Faster R-CNN 模型。
The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
image, and should be in ``0-1`` range. Different images can have different sizes.
模型的输入应该为一个由tensors组成的列表，每个tensor的形状为[C,H,W]，对于每一个图像的元素值都应该在[0,1]的范围内，不同的图像有着不同的尺寸。
The behavior of the model changes depending if it is in training or evaluation mode.
模型有训练与评估两种模式，模型的表现取决于模型所处的模式。
During training, the model expects both the input tensors, as well as a targets (list of dictionary),
containing:
- boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with values of ``x``
between ``0`` and ``W`` and values of ``y`` between ``0`` and ``H``
- labels (``Int64Tensor[N]``): the class label for each ground-truth box
在训练过程中，模型需要输入图像的tensor，以及目标(字典组成的列表)，其包含：
- 边框（FloatTensor[N,4]）：真实框为[x1,y1,x2,y2]的形式，x 的值在 0~W 之间，y 的值在 0-H 之间。
- 标签（Int64Tensor[N]）：每个真实框的类别标签。
The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
losses for both the RPN and the R-CNN.
在训练期间，模型返回一个 ”Dict[Tensor]“，包含 RPN 与 R-CNN 阶段的分类与回归损失。
During inference, the model requires only the input tensors, and returns the post-processed
predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as
follows:
- boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with values of ``x``
between ``0`` and ``W`` and values of ``y`` between ``0`` and ``H``
- labels (``Int64Tensor[N]``): the predicted labels for each image
- scores (``Tensor[N]``): the scores or each prediction
在推理过程中，模型仅需要输入图像的tensor，然后返回经过后处理的预测结果以 "List[Dict[Tensor]]" 的形式，对于每一个输入图像，其 "Dict" 域如下：
- 边框（FloatTensor[N,4]）：预测框为[x1,y1,x2,y2]的形式，x 的值在 0~W 之间，y 的值在 0~H 之间。
- 标签（Int64Tensor[N]）：每个图像的预测标签。
- 分数（Tensor[N]）：每个预测的分数。
Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
Faster R—CNN 可以被导出为一个固定批大小域固定尺寸输入图像的 ONNX 格式。
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017
progress (bool): If True, displays a progress bar of the download to stderr
pretrained_backbone (bool): If True, returns a model with backbone pre-trained on Imagenet
num_classes (int): number of output classes of the model (including the background)
trainable_backbone_layers (int): number of trainable (not frozen) resnet layers starting from final block.
Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
参数：
pretrianed（bool）：如果为真，返回一个在 COCO train2017 上的预训练模型。
progress（bool）：如果为真，将下载进度条展示在屏幕。
pretrained_backbone（bool）：如果为真，返回一个在 Imagenet 上的主干网络预训练模型。
num_classes（int）：模型输出的种类数量（包括背景）。
trainable_backbone_layers（int）：从最后一个块开始可训练 ResNet 层的数量（未被冻结）。合法的值在 0~5 之间，5 意味着所有主干网络的层都是可训练的。
"""
# 使用 assert 判断 trainable_backbone_layers 的值是否合法
assert trainable_backbone_layers <= 5 and trainable_backbone_layers >= 0
# dont freeze any layers if pretrained model or backbone is not used
# 如果预训练模型或者预训练主干网络未被使用，不要冻结任何层。
if not (pretrained or pretrained_backbone):
trainable_backbone_layers = 5
if pretrained:
# no need to download the backbone if pretrained is set
# 如果预训练模型被使用，就不需要下载预训练主干网络
pretrained_backbone = False
# 获取 ResNet_FPN 主干网络
backbone = resnet_fpn_backbone(resnet50, pretrained_backbone, trainable_layers=trainable_backbone_layers)
# 获取 Faster R-CNN 模型
model = FasterRCNN(backbone, num_classes, **kwargs)
if pretrained:
# 如果使用预训练模型，就下载相关的预训练模型配置
state_dict = load_state_dict_from_url(model_urls[fasterrcnn_resnet50_fpn_coco],
progress=progress)
# 加载模型配置到模型中
model.load_state_dict(state_dict)
return model # 返回模型
Debug 进获取 ResNet_FPN 主干网络对应代码。
def resnet_fpn_backbone(
backbone_name,
pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d,
trainable_layers=3,
returned_layers=None,
extra_blocks=None
):
"""
Constructs a specified ResNet backbone with FPN on top. Freezes the specified number of layers in the backbone.
构建一个在顶端加入FPN的ResNet主干网络。冻结主干网络中指定数量的层。