0. 引言
0.1 如何找到PyTorch官方的模型源码
from torchvision.models.detection import faster_rcnn
但是只有模型代码，并没有和训练相关的代码。
0.2 如何找到PyTorch官方的训练相关代码
链接：https://github.com/pytorch/vision/tree/main/references
0.3 Faster R-CNN的环境配置
- Python 3.6以上
- PyTorch 1.5以上
- pycocotools
pip install pycocotools - Ubuntu或Centos（不建议Windows）
- GPU
- lxml
- matplotlib
- numpy
- tqdm
- Pillow
0.4 文件结构
├── backbone: 特征提取网络，可以根据自己的要求选择：①MobileNet v2；②ResNet-50 + FPN（Feature Pyramid Networks， 特征金字塔、池化金字塔）
├── network_files: Faster R-CNN网络（包括Fast R-CNN以及RPN等模块）
├── train_utils: 训练验证相关模块（包括cocotools）：https://github.com/pytorch/vision/tree/main/references/detection
├── my_dataset.py: 自定义dataset用于读取VOC数据集
├── train_mobilenet.py: 以MobileNetV2做为backbone进行训练（主讲，准确率没有下面的好）
├── train_resnet50_fpn.py: 以resnet50+FPN做为backbone进行训练（效果最好）
├── train_multi_GPU.py: 针对使用多GPU的用户使用（并行训练）
├── predict.py: 简易的预测脚本，使用训练好的权重进行预测
├── validation.py: 利用训练好的权重验证/测试数据的COCO指标，并生成record_mAP.txt文件
└── pascal_voc_classes.json: PASCAL VOC标签文件
0.5 pascal_voc_classes.json
内容如下：
{
"aeroplane": 1,
"bicycle": 2,
"bird": 3,
"boat": 4,
"bottle": 5,
"bus": 6,
"car": 7,
"cat": 8,
"chair": 9,
"cow": 10,
"diningtable": 11,
"dog": 12,
"horse": 13,
"motorbike": 14,
"person": 15,
"pottedplant": 16,
"sheep": 17,
"sofa": 18,
"train": 19,
"tvmonitor": 20
}
Q：类别为什么不从0开始？
A：在目标检测中，一般0是留给背景（负样本）的
虽然PASCAL VOC 2012有20个类别，但实际训练时给了21个类别（为背景专门设置了一个类别）。
0.6 预训练权重下载地址（下载后放入backbone文件夹中）：
- MobileNetV2 backbone: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
- ResNet50+FPN backbone: https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
- 注意，下载的预训练权重记得要重命名，比如在train_resnet50_fpn.py中读取的是
fasterrcnn_resnet50_fpn_coco.pth文件，
不是fasterrcnn_resnet50_fpn_coco-258fb6c6.pth- 对于MobileNet v2来说，因为预训练权重只有backbone，而Faster R-CNN除了backbone还有RPN网络以及FC层，所以这个训练权重是不完整的。
- 而ResNet50+FPN是包含backbone和RPN的，所以是一个完整的预训练权重。
0.7 数据集
- Pascal VOC2012 train/val数据集下载地址：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
0.8 训练方法
- 确保提前准备好数据集
- 确保提前下载好对应预训练模型权重
- 若要训练mobilenetv2+fasterrcnn，直接使用train_mobilenet.py训练脚本
- 若要训练resnet50+fpn+fasterrcnn，直接使用train_resnet50_fpn.py训练脚本
- 若要使用多GPU训练，使用
python -m torch.distributed.launch --nproc_per_node=8 --use_env train_multi_GPU.py指令,nproc_per_node参数为使用GPU数量 - 如果想指定使用哪些GPU设备可在指令前加上
CUDA_VISIBLE_DEVICES=0,3(例如我只要使用设备中的第1块和第4块GPU设备) CUDA_VISIBLE_DEVICES=0,3 python -m torch.distributed.launch --nproc_per_node=2 --use_env train_multi_GPU.py
- 学术研究使用mobilenetv2+fasterrcnn
- 实际应用使用resnet50+fpn+fasterrcnn
0.9 注意事项
- 在使用训练脚本时，注意要将
--data-path(VOC_root)设置为自己存放VOCdevkit文件夹所在的根目录 - 由于带有FPN结构的Faster RCNN很吃显存，如果GPU的显存不够(如果batch_size小于8的话)建议在create_model函数中使用默认的norm_layer，即不传递norm_layer变量，默认去使用
FrozenBatchNorm2d(即不会去更新参数的bn层)，使用中发现效果也很好。 - 在使用预测脚本时，要将
train_weights设置为你自己生成的权重路径。 - 使用validation文件时，注意确保你的验证集或者测试集中必须包含每个类别的目标，并且使用时只需要修改
--num-classes、--data-path和--weights-path即可，其他代码尽量不要改动
0.10 Faster R-CNN框架图
1. 自定义Dataset
https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html
Defining the Dataset，定义你的Dataset
The reference scripts for training object detection, instance segmentation and person keypoint detection allows for easily supporting adding new custom datasets. The dataset should inherit from the standard torch.utils.data.Dataset class, and implement __len__ and __getitem__.
用于训练对象检测、实例分割和人物关键点检测的参考脚本允许轻松支持添加新的自定义数据集。 数据集应该从标准的torch.utils.data.Dataset类继承，并实现__len__和__getitem__。
__len__：获取数据集长度__getitem__：返回图片和信息
The only specificity that we require is that the dataset __getitem__ should return:
- image: a PIL Image of size (
H, W) - target: a dict containing the following fields
boxes(FloatTensor[N, 4]): the coordinates of theNbounding boxes in[x0, y0, x1, y1]format, ranging from0toWand0toHlabels(Int64Tensor[N]): the label for each bounding box.0represents always the background class.image_id(Int64Tensor[1]): an image identifier. It should be unique between all the images in the dataset, and is used during evaluationarea(Tensor[N]): The area of the bounding box. This is used during evaluation with the COCO metric, to separate the metric scores between small, medium and large boxes.iscrowd(UInt8Tensor[N]): instances withiscrowd=Truewill be ignored during evaluation.- (optionally)
masks(UInt8Tensor[N, H, W]): The segmentation masks for each one of the objects - (optionally)
keypoints(FloatTensor[N, K, 3]): For each one of theNobjects, it contains theKkeypoints in[x, y, visibility]format, defining the object.visibility=0means that the keypoint is not visible. Note that for data augmentation, the notion of flipping a keypoint is dependent on the data representation, and you should probably adaptreferences/detection/transforms.pyfor your new keypoint representation
If your model returns the above methods, they will make it work for both training and evaluation, and will use the evaluation scripts from pycocotools which can be installed with pip install pycocotools.
如果您的模型返回上述方法，它们将使其适用于训练和评估，并将使用来自 pycocotools 的评估脚本，可以使用 pip install pycocotools 安装。
NOTE
One note on the labels. The model considers class0as background. If your dataset does not contain the background class, you should not have0in your labels. For example, assuming you have just two classes, cat and dog, you can define1(not0) to represent cats and2to represent dogs. So, for instance, if one of the images has both classes, your labels tensor should look like[1,2].
Additionally, if you want to use aspect ratio grouping during training (so that each batch only contains images with similar aspect ratios), then it is recommended to also implement a
get_height_and_widthmethod, which returns the height and the width of the image. If this method is not provided, we query all elements of the dataset via__getitem__, which loads the image in memory and is slower than if a custom method is provided.
另外，如果你想在训练时使用长宽比分组（这样每批只包含长宽比相似的图像），那么建议也实现一个 get_height_and_width 方法，它返回图像的高度和宽度。 如果没有提供这个方法，我们会通过__getitem__查询数据集的所有元素，这会将图像加载到内存中，并且比提供自定义方法要慢。
1.1 自定义dataset代码
import numpy as np
from torch.utils.data import Dataset
import os
import torch
import json
from PIL import Image
from lxml import etree
class VOCDataSet(Dataset):
"""
读取解析PASCAL VOC2007/2012数据集
需要实现两个方法：
1. __len__
2. __getitem__
[可选] 3. get_height_and_width
"""
def __init__(self, voc_root, year="2012", transforms=None, txt_name: str = "train.txt"):
assert year in ["2007", "2012"], "year must be in ['2007', '2012']"
# 增加容错能力
if "VOCdevkit" in voc_root:
self.root = os.path.join(voc_root, f"VOC{year}")
else:
self.root = os.path.join(voc_root, "VOCdevkit", f"VOC{year}")
self.img_root = os.path.join(self.root, "JPEGImages")
self.annotations_root = os.path.join(self.root, "Annotations")
# read train.txt or val.txt file
txt_path = os.path.join(self.root, "ImageSets", "Main", txt_name)
assert os.path.exists(txt_path), "not found {} file.".format(txt_name)
# 打开txt文件并读取每一行
"""
因为每一行后面都有一个换行符，所以使用line.strip()方法将换行符去掉
"""
with open(txt_path) as read:
# xml_list用list存储每一个图片信息文件的名称 xxx.xml
xml_list = [os.path.join(self.annotations_root, line.strip() + ".xml")
for line in read.readlines() if len(line.strip()) > 0]
self.xml_list = []
# check file
for xml_path in xml_list:
if os.path.exists(xml_path) is False: # 如果xxx.xml并不存在
print(f"Warning: not found '{xml_path}', skip this annotation file.")
continue
# check for targets
with open(xml_path) as fid:
xml_str = fid.read() # 读取标注文件的每一行
"""
etree.fromstring()
该方法是将xml格式转化为Element对象，Element 对象代表 XML 文档中的一个元素。
元素可以包含属性、其他元素或文本。如果一个元素包含文本，则在文本节点中表示该文本。
传入的为一个xml文件，经过该方法后变成一个Element对象<Element annotation at 0x24b46496680>
"""
xml = etree.fromstring(xml_str)
data = self.parse_xml_to_dict(xml)["annotation"] # 将xml文件解析成字典形式
if "object" not in data: # 标注信息中如果没有要对象
print(f"INFO: no objects in {xml_path}, skip this annotation file.")
continue
self.xml_list.append(xml_path)
assert len(self.xml_list) > 0, "in '{}' file does not find any information.".format(txt_path)
# read class_indict
json_file = './pascal_voc_classes.json'
assert os.path.exists(json_file), "{} file not exist.".format(json_file)
with open(json_file, 'r') as f:
self.class_dict = json.load(f) # 这里是把json转换为dict
self.transforms = transforms
def __len__(self):
# 获取所有数据文件（标签个数=图片）的个数
return len(self.xml_list)
def __getitem__(self, idx):
# read xml
xml_path = self.xml_list[idx] # 获取xml文件路径
with open(xml_path) as fid:
xml_str = fid.read()
xml = etree.fromstring(xml_str) # 读取xml文件: <Element annotation at 0x7f1e540bcc00>
"""
self.parse_xml_to_dict(xml)就可以得到字典，但因为<annotation>是父节点，所以我们需要去掉它，故["annotation"]
这样就可以得到我们想要的字典了：
{"folder": "VOC2012",
"filename": "2007_000063.jpg",
"source": {"database": "The VOC2007 Database", "annotation": "PASCAL VOC2007", "image": "flickr"},
"size": {"width": '500', "height": '375', "depth": '3'},
"segmented": 1,
"object": {"name": "dog", "pose": "Unspecified", "truncated": '0', "difficult": '0', "bndbox":
{"xmin": 123, "ymin": 115, "xmax": '379', "ymax": '275'}},
"object": {"name": "chair", "pose": "Frontal", "truncated": '1', "difficult": 0, "bndbox":
{"xmin": '75', "ymin": '1', "xmax": '428', "ymax": '375'}}
}
Note: 这样解析后，所有的值的数据类型均为str
"""
data = self.parse_xml_to_dict(xml)["annotation"] # 获取下一层字典的信息
img_path = os.path.join(self.img_root, data["filename"]) # 获取图片完整的路径
image = Image.open(img_path) # 使用Pillow打开图片（这里只是在内存中打开并没有通过GUI显示）
if image.format != "JPEG":
raise ValueError("Image '{}' format not JPEG".format(img_path))
boxes = []
labels = []
iscrowd = [] # coco数据独有的，表示目标之间是否有重叠，这里我们简单理解为“是否好检测”。0表示单目标 -> 比较好检测 -> 可以和
# VOC的difficult参数结合起来
assert "object" in data, "{} lack of object information.".format(xml_path)
# 遍历每一个<object>信息
for obj in data["object"]:
# 需将String转换为float
xmin = float(obj["bndbox"]["xmin"])
xmax = float(obj["bndbox"]["xmax"])
ymin = float(obj["bndbox"]["ymin"])
ymax = float(obj["bndbox"]["ymax"])
# 进一步检查数据，有的标注信息中可能有w或h为0的情况，这样的数据会导致计算回归loss为nan
if xmax <= xmin or ymax <= ymin:
print("Warning: in '{}' xml, there are some bbox w/h <=0".format(xml_path))
# 跳过本次迭代
continue
boxes.append([xmin, ymin, xmax, ymax])
labels.append(self.class_dict[obj["name"]]) # 根据目标的名字获取对应的标签索引值
if "difficult" in obj:
iscrowd.append(int(obj["difficult"]))
else:
iscrowd.append(0) # 表示容易检测
# convert everything into a torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
image_id = torch.tensor([idx]) # 当前数据对应的索引值
# area = (y_max - y_min) * (x_max - x_min)
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
image, target = self.transforms(image, target)
return image, target
def get_height_and_width(self, idx):
# read xml
xml_path = self.xml_list[idx]
with open(xml_path) as fid:
xml_str = fid.read()
xml = etree.fromstring(xml_str)
data = self.parse_xml_to_dict(xml)["annotation"] # 获取字典并剥掉最外层
data_height = int(data["size"]["height"]) # 获得图片的宽度（str -> int）
data_width = int(data["size"]["width"]) # 获得图片的高度（str -> int）
return data_height, data_width
def parse_xml_to_dict(self, xml):
"""
将xml文件解析成字典形式，参考tensorflow的recursive_parse_xml_to_dict
Args:
xml: xml tree obtained by parsing XML file contents using lxml.etree
Returns:
Python dictionary holding XML contents.
"""
# 这里的xml就是父节点<annotation>
if len(xml) == 0: # 遍历到底层，直接返回tag对应的信息
return {xml.tag: xml.text}
result = {}
for child in xml: # 遍历父节点返回子节点
child_result = self.parse_xml_to_dict(child) # 递归遍历标签信息
if child.tag != 'object': # 不是<object>节点
result[child.tag] = child_result[child.tag]
else: # 如果子节点为<object>
if child.tag not in result: # 因为object可能有多个，所以需要放入列表里
result[child.tag] = []
result[child.tag].append(child_result[child.tag])
return {xml.tag: result}
def coco_index(self, idx):
"""
该方法是专门为pycocotools统计标签信息准备，不对图像和标签作任何处理
由于不用去读取图片，可大幅缩减统计时间
Args:
idx: 输入需要获取图像的索引
"""
# read xml
xml_path = self.xml_list[idx]
with open(xml_path) as fid:
xml_str = fid.read()
xml = etree.fromstring(xml_str)
data = self.parse_xml_to_dict(xml)["annotation"]
data_height = int(data["size"]["height"])
data_width = int(data["size"]["width"])
# img_path = os.path.join(self.img_root, data["filename"])
# image = Image.open(img_path)
# if image.format != "JPEG":
# raise ValueError("Image format not JPEG")
boxes = []
labels = []
iscrowd = []
for obj in data["object"]:
xmin = float(obj["bndbox"]["xmin"])
xmax = float(obj["bndbox"]["xmax"])
ymin = float(obj["bndbox"]["ymin"])
ymax = float(obj["bndbox"]["ymax"])
boxes.append([xmin, ymin, xmax, ymax])
labels.append(self.class_dict[obj["name"]])
iscrowd.append(int(obj["difficult"]))
# convert everything into a torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
return (data_height, data_width), target
@staticmethod
def collate_fn(batch):
return tuple(zip(*batch))
if __name__ == '__main__':
import transforms
from draw_box_utils import draw_objs
from PIL import Image
import json
import matplotlib.pyplot as plt
import torchvision.transforms as ts
import random
import torchvision
# read class_indict
category_index = {}
try:
json_file = open('./pascal_voc_classes.json', 'r')
class_dict = json.load(json_file)
# 颠倒key和value的位置 "aeroplane": 1 -> 1: "aeroplane"
category_index = {str(v): str(k) for k, v in class_dict.items()}
except Exception as e:
print(e)
exit(-1)
data_transform = {
"train": transforms.Compose([transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val": transforms.Compose([transforms.ToTensor()])
}
# load train data set
train_data_set = VOCDataSet(voc_root=os.getcwd(), year="2012",
transforms=data_transform["train"],
txt_name="train.txt")
print(len(train_data_set))
for index in random.sample(range(0, len(train_data_set)), k=5): # 从训练集中随机采样5张图片
img, target = train_data_set[index] # 传入图像索引返回图像和target信息
img = torchvision.transforms.ToPILImage()(img) # 将tensor格式的图片转化为PIL的图片格式
plot_img = draw_objs(img,
target["boxes"].numpy(),
target["labels"].numpy(),
np.ones(target["labels"].shape[0]), # 因为读的是GT，所以概率是1
category_index=category_index,
box_thresh=0.5,
line_thickness=3, # 矩形框线的宽度
font='arial.ttf',
font_size=20)
plt.imshow(plot_img)
plt.show()
1.2 自定义Dataset的重点
-
首先明白自定义Dataset需要：
- 写一个类并继承自
torch.utils.data.Dataset - 实现
__len__方法（标注文件个数=训练文件个数） - 实现
__getitem__方法（返回值①image；②target）image：就是根据路径使用PIL加载图片并返回target：对于VOC 2012数据集，包含下面内容：- bbox（GT）的坐标 -> 转换为tensor
- 目标类别的索引值（1， 2， 3，…） -> 转换为tensor
- 图片的id
- bbox（GT）的面积
- 训练是否困难（0表示容易；1表示困难） -> 转换为tensor
- [可选] 实现
__get_height_and_width__方法（对于VOC 2012数据集可根据annotation中的信息可以直接获取）
- 写一个类并继承自
-
使用
line.strip()方法将换行符去掉` -
列表表达式和字典表达式的使用
-
读取
.xml文件的步骤- 先将
.xml文件写进内存
with open(xml_path) as fid: xml_str = fid.read() # 读取标注文件的每一行- 使用
etree.fromstring()将内存中的xml对象转换为Element对象
xml = etree.fromstring(xml_str)- 将
xml解析为字典形式并褪去最外面的字典壳
data = self.parse_xml_to_dict(xml)["annotation"]这样就可以得到我们想要的字典了：
{"folder": "VOC2012", "filename": "2007_000063.jpg", "source": {"database": "The VOC2007 Database", "annotation": "PASCAL VOC2007", "image": "flickr"}, "size": {"width": '500', "height": '375', "depth": '3'}, "segmented": 1, "object": {"name": "dog", "pose": "Unspecified", "truncated": '0', "difficult": '0', "bndbox": {"xmin": 123, "ymin": 115, "xmax": '379', "ymax": '275'}}, "object": {"name": "chair", "pose": "Frontal", "truncated": '1', "difficult": 0, "bndbox": {"xmin": '75', "ymin": '1', "xmax": '428', "ymax": '375'}} } Note: 这样解析后，所有的值的数据类型均为str - 先将
-
对于VOC来说，根据标注信息中的坐标计算面积公式如下：
a r e a = ( y m a x − y m i n ) ∗ ( x m a x − x m i n ) \mathrm{area} = (y_{\mathrm{max}} - y_{\mathrm{min}}) * (x_{\mathrm{max}} - x_{\mathrm{min}}) area=(ymax−ymin)∗(xmax−xmin) -
random.sample(range(0, len(train_data_set)), k=5) -
.xml的子父节点 -
torchvision.transforms.ToPILImage()(img)将tensor格式的图片转化为PIL的图片格式
1.3 自定义transform
在训练时只使用了随机旋转，但是要注意，图片旋转了，标注框的位置也需要旋转
代码如下：
import random
from torchvision.transforms import functional as F
class Compose(object):
"""组合多个transform函数"""
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
class ToTensor(object):
"""将PIL图像转为Tensor"""
def __call__(self, image, target):
image = F.to_tensor(image)
return image, target
class RandomHorizontalFlip(object):
"""随机水平翻转图像以及bboxes"""
def __init__(self, prob=0.5):
self.prob = prob
def __call__(self, image, target):
if random.random() < self.prob:
height, width = image.shape[-2:]
image = image.flip(-1) # 水平翻转图片
bbox = target["boxes"]
# bbox: xmin, ymin, xmax, ymax
bbox[:, [0, 2]] = width - bbox[:, [2, 0]] # 翻转对应bbox坐标信息
target["boxes"] = bbox
return image, target
2. Faster R-CNN框架图 —— faster_rcnn_framework.py
其中黄色的框表示在训练中才会有的，预测时没有。
import warnings
from collections import OrderedDict
from typing import Tuple, List, Dict, Optional, Union
import torch
from torch import nn, Tensor
import torch.nn.functional as F
from torchvision.ops import MultiScaleRoIAlign
from .roi_head import RoIHeads
from .transform import GeneralizedRCNNTransform
from .rpn_function import AnchorsGenerator, RPNHead, RegionProposalNetwork
class FasterRCNNBase(nn.Module):
"""
Main class for Generalized R-CNN.
Arguments:
backbone (nn.Module): 特征提取网络部分
rpn (nn.Module): 候选框生成部分
roi_heads (nn.Module): takes the features + the proposals from the RPN and computes
detections / masks from it.
transform (nn.Module): performs the data transformation from the inputs to feed into
the model
"""
def __init__(self, backbone, rpn, roi_heads, transform):
super(FasterRCNNBase, self).__init__()
self.transform = transform
self.backbone = backbone
self.rpn = rpn
self.roi_heads = roi_heads
# used only on torchscript mode
self._has_warned = False
@torch.jit.unused
def eager_outputs(self, losses, detections):
# type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]
if self.training:
return losses
return detections
def forward(self, images, targets=None):
# type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
"""
Arguments:
images (list[Tensor]): images to be processed
targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)
Returns:
result (list[BoxList] or dict[Tensor]): the output from the model.
During training, it returns a dict[Tensor] which contains the losses.
During testing, it returns list[BoxList] contains additional fields
like `scores`, `labels` and `mask` (for Mask R-CNN models).
"""
if self.training and targets is None:
raise ValueError("In training mode, targets should be passed")
if self.training:
assert targets is not None
for target in targets: # 进一步判断传入的target的boxes参数是否符合规定
boxes = target["boxes"]
if isinstance(boxes, torch.Tensor):
if len(boxes.shape) != 2 or boxes.shape[-1] != 4:
raise ValueError("Expected target boxes to be a tensor"
"of shape [N, 4], got {:}.".format(
boxes.shape))
else:
raise ValueError("Expected target boxes to be of type "
"Tensor, got {:}.".format(type(boxes)))
# original_image_sizes：存储每张图片原始的尺寸，为了之后可以映射到原图中
original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], [])
for img in images:
val = img.shape[-2:] # [H, W]
assert len(val) == 2 # 防止输入的是个一维向量
original_image_sizes.append((val[0], val[1]))
# original_image_sizes = [img.shape[-2:] for img in images]
images, targets = self.transform(images, targets) # 对图像进行预处理
# [(images, targets), (images, targets), ...]这才是需要送入网络的batch
# print(images.tensors.shape)
features = self.backbone(images.tensors) # 将图像输入backbone得到特征图
if isinstance(features, torch.Tensor): # 若只在一层特征层上预测，将feature放入有序字典中，并编号为‘0’
features = OrderedDict([('0', features)]) # 若在多层特征层上预测，传入的就是一个有序字典
# 将特征层以及标注target信息传入rpn中
# proposals: List[Tensor], Tensor_shape: [num_proposals, 4],
# 每个proposals是绝对坐标，且为(x1, y1, x2, y2)格式
proposals, proposal_losses = self.rpn(images, features, targets)
# 将rpn生成的数据以及标注target信息传入fast rcnn后半部分
detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
# 对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)
losses = {}
losses.update(detector_losses)
losses.update(proposal_losses)
if torch.jit.is_scripting():
if not self._has_warned:
warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")
self._has_warned = True
return losses, detections
else:
return self.eager_outputs(losses, detections)
# if self.training:
# return losses
#
# return detections
class TwoMLPHead(nn.Module):
"""
Standard heads for FPN-based models
Arguments:
in_channels (int): number of input channels
representation_size (int): size of the intermediate representation
"""
def __init__(self, in_channels, representation_size):
super(TwoMLPHead, self).__init__()
self.fc6 = nn.Linear(in_channels, representation_size)
self.fc7 = nn.Linear(representation_size, representation_size)
def forward(self, x):
x = x.flatten(start_dim=1)
x = F.relu(self.fc6(x))
x = F.relu(self.fc7(x))
return x
class FastRCNNPredictor(nn.Module):
"""
Standard classification + bounding box regression layers
for Fast R-CNN.
Arguments:
in_channels (int): number of input channels
num_classes (int): number of output classes (including background)
"""
def __init__(self, in_channels, num_classes):
super(FastRCNNPredictor, self).__init__()
self.cls_score = nn.Linear(in_channels, num_classes)
self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
def forward(self, x):
if x.dim() == 4:
assert list(x.shape[2:]) == [1, 1]
x = x.flatten(start_dim=1)
scores = self.cls_score(x)
bbox_deltas = self.bbox_pred(x)
return scores, bbox_deltas
class FasterRCNN(FasterRCNNBase):
"""
Implements 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.
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
between 0 and H and 0 and W
- labels (Int64Tensor[N]): the class label for each ground-truth box
The model returns a Dict[Tensor] during training, containing the classification and regression
losses for both the RPN and the 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 between
0 and H and 0 and W
- labels (Int64Tensor[N]): the predicted labels for each image
- scores (Tensor[N]): the scores or each prediction
Arguments:
backbone (nn.Module): the network used to compute the features for the model.
It should contain a out_channels attribute, which indicates the number of output
channels that each feature map has (and it should be the same for all feature maps).
The backbone should return a single Tensor or and OrderedDict[Tensor].
num_classes (int): number of output classes of the model (including the background).
If box_predictor is specified, num_classes should be None.
min_size (int): minimum size of the image to be rescaled before feeding it to the backbone
max_size (int): maximum size of the image to be rescaled before feeding it to the backbone
image_mean (Tuple[float, float, float]): mean values used for input normalization.
They are generally the mean values of the dataset on which the backbone has been trained
on
image_std (Tuple[float, float, float]): std values used for input normalization.
They are generally the std values of the dataset on which the backbone has been trained on
rpn_anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature
maps.
rpn_head (nn.Module): module that computes the objectness and regression deltas from the RPN
rpn_pre_nms_top_n_train (int): number of proposals to keep before applying NMS during training
rpn_pre_nms_top_n_test (int): number of proposals to keep before applying NMS during testing
rpn_post_nms_top_n_train (int): number of proposals to keep after applying NMS during training
rpn_post_nms_top_n_test (int): number of proposals to keep after applying NMS during testing
rpn_nms_thresh (float): NMS threshold used for postprocessing the RPN proposals
rpn_fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be
considered as positive during training of the RPN.
rpn_bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be
considered as negative during training of the RPN.
rpn_batch_size_per_image (int): number of anchors that are sampled during training of the RPN
for computing the loss
rpn_positive_fraction (float): proportion of positive anchors in a mini-batch during training
of the RPN
rpn_score_thresh (float): during inference, only return proposals with a classification score
greater than rpn_score_thresh
box_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in
the locations indicated by the bounding boxes
box_head (nn.Module): module that takes the cropped feature maps as input
box_predictor (nn.Module): module that takes the output of box_head and returns the
classification logits and box regression deltas.
box_score_thresh (float): during inference, only return proposals with a classification score
greater than box_score_thresh
box_nms_thresh (float): NMS threshold for the prediction head. Used during inference
box_detections_per_img (int): maximum number of detections per image, for all classes.
box_fg_iou_thresh (float): minimum IoU between the proposals and the GT box so that they can be
considered as positive during training of the classification head
box_bg_iou_thresh (float): maximum IoU between the proposals and the GT box so that they can be
considered as negative during training of the classification head
box_batch_size_per_image (int): number of proposals that are sampled during training of the
classification head
box_positive_fraction (float): proportion of positive proposals in a mini-batch during training
of the classification head
bbox_reg_weights (Tuple[float, float, float, float]): weights for the encoding/decoding of the
bounding boxes
"""
def __init__(self, backbone, num_classes=None, # num_classes是需要加上背景的
# transform parameter
min_size=800, max_size=1333, # 预处理resize时限制的最小尺寸与最大尺寸
image_mean=None, image_std=None, # 预处理normalize时使用的均值和方差
# RPN parameters
rpn_anchor_generator=None, rpn_head=None,
rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000, # rpn中在nms处理前保留的proposal数(根据score)
rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000, # rpn中在nms处理后保留的proposal数
rpn_nms_thresh=0.7, # rpn中进行nms处理时使用的iou阈值
rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3, # rpn计算损失时，采集正负样本设置的阈值
rpn_batch_size_per_image=256, rpn_positive_fraction=0.5, # rpn计算损失时采样的样本数，以及正样本占总样本的比例
rpn_score_thresh=0.0,
# Box parameters
box_roi_pool=None, box_head=None, box_predictor=None,
# 移除低目标概率 fast rcnn中进行nms处理的阈值 对预测结果根据score排序取前100个目标
box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5, # fast rcnn计算误差时，采集正负样本设置的阈值
box_batch_size_per_image=512, box_positive_fraction=0.25, # fast rcnn计算误差时采样的样本数，以及正样本占所有样本的比例
bbox_reg_weights=None):
if not hasattr(backbone, "out_channels"):
raise ValueError(
"backbone should contain an attribute out_channels"
"specifying the number of output channels (assumed to be the"
"same for all the levels"
)
assert isinstance(rpn_anchor_generator, (AnchorsGenerator, type(None))) # 传入None也可以，传入None后面会生成一个
assert isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None)))
if num_classes is not None:
if box_predictor is not None: # 如果box_predictor不为None，报错
raise ValueError("num_classes should be None when box_predictor "
"is specified")
else:
if box_predictor is None:
raise ValueError("num_classes should not be None when box_predictor "
"is not specified")
# 预测特征层的channels
out_channels = backbone.out_channels
# 若anchor生成器为空，则自动生成针对resnet50_fpn的anchor生成器
"""
>>> ((0.5, 1.0, 2.0),) * 5
((0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0))
"""
if rpn_anchor_generator is None:
anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
rpn_anchor_generator = AnchorsGenerator(
anchor_sizes, aspect_ratios
)
# 生成RPN通过滑动窗口预测网络部分
if rpn_head is None:
rpn_head = RPNHead(
out_channels, rpn_anchor_generator.num_anchors_per_location()[0]
)
# 默认rpn_pre_nms_top_n_train = 2000, rpn_pre_nms_top_n_test = 1000,
# 默认rpn_post_nms_top_n_train = 2000, rpn_post_nms_top_n_test = 1000,
rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
# 定义整个RPN框架
rpn = RegionProposalNetwork(
rpn_anchor_generator, rpn_head,
rpn_fg_iou_thresh, rpn_bg_iou_thresh,
rpn_batch_size_per_image, rpn_positive_fraction,
rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_nms_thresh,
score_thresh=rpn_score_thresh)
# Multi-scale RoIAlign pooling
if box_roi_pool is None:
box_roi_pool = MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3'], # 在哪些特征层进行roi pooling
output_size=[7, 7],
sampling_ratio=2)
# fast RCNN中roi pooling后的展平处理两个全连接层部分
if box_head is None:
resolution = box_roi_pool.output_size[0] # 默认等于7
representation_size = 1024
box_head = TwoMLPHead(
out_channels * resolution ** 2,
representation_size
)
# 在box_head的输出上预测部分
if box_predictor is None:
representation_size = 1024
box_predictor = FastRCNNPredictor(
representation_size,
num_classes)
# 将roi pooling, box_head以及box_predictor结合在一起
roi_heads = RoIHeads(
# box
box_roi_pool, box_head, box_predictor,
box_fg_iou_thresh, box_bg_iou_thresh, # 0.5 0.5
box_batch_size_per_image, box_positive_fraction, # 512 0.25
bbox_reg_weights,
box_score_thresh, box_nms_thresh, box_detections_per_img) # 0.05 0.5 100
if image_mean is None:
image_mean = [0.485, 0.456, 0.406]
if image_std is None:
image_std = [0.229, 0.224, 0.225]
# 对数据进行标准化，缩放，打包成batch等处理部分
transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std)
super(FasterRCNN, self).__init__(backbone, rpn, roi_heads, transform)
3. Transform —— transform.py
import math
from typing import List, Tuple, Dict, Optional
import torch
from torch import nn, Tensor
import torchvision
from .image_list import ImageList
@torch.jit.unused
def _resize_image_onnx(image, self_min_size, self_max_size):
# type: (Tensor, float, float) -> Tensor
from torch.onnx import operators
im_shape = operators.shape_as_tensor(image)[-2:]
min_size = torch.min(im_shape).to(dtype=torch.float32)
max_size = torch.max(im_shape).to(dtype=torch.float32)
scale_factor = torch.min(self_min_size / min_size, self_max_size / max_size)
image = torch.nn.functional.interpolate(
image[None], scale_factor=scale_factor, mode="bilinear", recompute_scale_factor=True,
align_corners=False)[0]
return image
def _resize_image(image, self_min_size, self_max_size):
# type: (Tensor, float, float) -> Tensor
im_shape = torch.tensor(image.shape[-2:])
min_size = float(torch.min(im_shape)) # 获取高宽中的最小值
max_size = float(torch.max(im_shape)) # 获取高宽中的最大值
scale_factor = self_min_size / min_size # 根据指定最小边长和图片最小边长计算缩放比例
# 如果使用该缩放比例计算的图片最大边长大于指定的最大边长
if max_size * scale_factor > self_max_size:
scale_factor = self_max_size / max_size # 将缩放比例设为指定最大边长和图片最大边长之比
# interpolate利用插值的方法缩放图片
# image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]
# bilinear只支持4D Tensor
image = torch.nn.functional.interpolate(
image[None], scale_factor=scale_factor, mode="bilinear", recompute_scale_factor=True,
align_corners=False)[0]
return image
class GeneralizedRCNNTransform(nn.Module):
"""
Performs input / target transformation before feeding the data to a GeneralizedRCNN
model.
The transformations it perform are:
- input normalization (mean subtraction and std division)
- input / target resizing to match min_size / max_size
It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
"""
def __init__(self, min_size, max_size, image_mean, image_std):
super(GeneralizedRCNNTransform, self).__init__()
if not isinstance(min_size, (list, tuple)):
min_size = (min_size,) # 转换为tuple类型
self.min_size = min_size # 指定图像的最小边长范围
self.max_size = max_size # 指定图像的最大边长范围
self.image_mean = image_mean # 指定图像在标准化处理中的均值
self.image_std = image_std # 指定图像在标准化处理中的方差
def normalize(self, image):
"""标准化处理"""
dtype, device = image.dtype, image.device
# 将传入的均值和方差转换为与输入图片相同的类型和设备
mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
# [:, None, None]: shape [3] -> [3, 1, 1]
"""
因为将mean和std转换为torch.tensor后，它俩的shape为[3]，而mean[:, None, None]看似是在切片，但实际上给mean和std添加了
两个维度，即shape从[3]变为了[3, 1, 1]
>>> x = torch.randint(1, (3, 112, 112)).cuda()
>>> x.shape
torch.Size([3, 112, 112])
>>> mean = [0.1, 0.2, 0.3]
>>> mean = torch.as_tensor(mean, dtype=x.dtype, device=x.device)
>>> mean.shape
torch.Size([3])
>>> mean[:, None, None].shape
torch.Size([3, 1, 1])
"""
return (image - mean[:, None, None]) / std[:, None, None]
def torch_choice(self, k):
# type: (List[int]) -> int
"""
Implements `random.choice` via torch ops so it can be compiled with
TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
is fixed.
"""
index = int(torch.empty(1).uniform_(0., float(len(k))).item())
return k[index]
def resize(self, image, target):
# type: (Tensor, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]
"""
将图片缩放到指定的大小范围内，并对应缩放bboxes信息
Args:
image: 输入的图片
target: 输入图片的相关信息（包括bboxes信息）
Returns:
image: 缩放后的图片
target: 缩放bboxes后的图片相关信息
"""
# image shape is [channel, height, width]
h, w = image.shape[-2:]
if self.training:
size = float(self.torch_choice(self.min_size)) # 指定输入图片的最小边长,注意是self.min_size不是min_size
else:
# FIXME assume for now that testing uses the largest scale
size = float(self.min_size[-1]) # 指定输入图片的最小边长,注意是self.min_size不是min_size
if torchvision._is_tracing():
image = _resize_image_onnx(image, size, float(self.max_size))
else:
image = _resize_image(image, size, float(self.max_size))
if target is None:
return image, target
bbox = target["boxes"]
# 根据图像的缩放比例来缩放bbox
bbox = resize_boxes(boxes=bbox, original_size=[h, w], new_size=image.shape[-2:])
target["boxes"] = bbox
return image, target
# _onnx_batch_images() is an implementation of
# batch_images() that is supported by ONNX tracing.
@torch.jit.unused
def _onnx_batch_images(self, images, size_divisible=32):
# type: (List[Tensor], int) -> Tensor
max_size = []
for i in range(images[0].dim()):
max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)
max_size.append(max_size_i)
stride = size_divisible
max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)
max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)
max_size = tuple(max_size)
# work around for
# pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
# which is not yet supported in onnx
padded_imgs = []
for img in images:
padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
padded_img = torch.nn.functional.pad(img, [0, padding[2], 0, padding[1], 0, padding[0]])
padded_imgs.append(padded_img)
return torch.stack(padded_imgs)
def max_by_axis(self, the_list):
# type: (List[List[int]]) -> List[int]
# the_list: [BS, C, H, W]
maxes = the_list[0] # [C, H, W]
for sublist in the_list[1:]: # 得到每一个图片的shape
for index, item in enumerate(sublist): # 迭代每一张图片的shape
maxes[index] = max(maxes[index], item)
return maxes
def batch_images(self, images, size_divisible=32):
# type: (List[Tensor], int) -> Tensor
"""
将一批图像打包成一个batch返回（注意batch中每个tensor的shape是相同的）
Args:
images: 输入的一批图片
size_divisible: 将图像高和宽调整到该数的整数倍
Returns:
batched_imgs: 打包成一个batch后的tensor数据
"""
if torchvision._is_tracing():
# batch_images() does not export well to ONNX
# call _onnx_batch_images() instead
return self._onnx_batch_images(images, size_divisible)
# 分别计算一个batch中所有图片中的最大channel, height, width
"""
>>> y.shape
torch.Size([2, 3, 112, 112])
>>> list(y.shape)
[2, 3, 112, 112]
"""
max_size = self.max_by_axis([list(img.shape) for img in images])
stride = float(size_divisible)
# max_size = list(max_size)
# 将height向上调整到stride的整数倍（最靠近stride整数倍的数值）
# math.ceil：向上取整
max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)
# 将width向上调整到stride的整数倍
max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)
# [batch, channel, height, width]
batch_shape = [len(images)] + max_size
# 创建shape为batch_shape且值全部为0的tensor
batched_imgs = images[0].new_full(batch_shape, 0)
for img, pad_img in zip(images, batched_imgs):
# 将输入images中的每张图片复制到新的batched_imgs的每张图片中，对齐左上角，保证bboxes的坐标不变
# 这样保证输入到网络中一个batch的每张图片的shape相同
# copy_: Copies the elements from src into self tensor and returns self
pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) # pad_img[C, H, W]
return batched_imgs
def postprocess(self,
result, # type: List[Dict[str, Tensor]]
image_shapes, # type: List[Tuple[int, int]]
original_image_sizes # type: List[Tuple[int, int]]
):
# type: (...) -> List[Dict[str, Tensor]]
"""
对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
Args:
result: list(dict), 网络的预测结果, len(result) == batch_size
image_shapes: list(torch.Size), 图像预处理缩放后的尺寸, len(image_shapes) == batch_size
original_image_sizes: list(torch.Size), 图像的原始尺寸, len(original_image_sizes) == batch_size
Returns:
"""
if self.training:
return result
# 遍历每张图片的预测信息，将boxes信息还原回原尺度
for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
boxes = pred["boxes"]
boxes = resize_boxes(boxes, im_s, o_im_s) # 将bboxes缩放回原图像尺度上
result[i]["boxes"] = boxes
return result
def __repr__(self):
"""自定义输出实例化对象的信息，可通过print打印实例信息"""
format_string = self.__class__.__name__ + '('
_indent = '\n '
format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std)
format_string += "{0}Resize(min_size={1}, max_size={2}, mode='bilinear')".format(_indent, self.min_size,
self.max_size)
format_string += '\n)'
return format_string
def forward(self,
images, # type: List[Tensor]
targets=None # type: Optional[List[Dict[str, Tensor]]]
):
# type: (...) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]
images = [img for img in images]
for i in range(len(images)):
image = images[i]
target_index = targets[i] if targets is not None else None
if image.dim() != 3:
raise ValueError("images is expected to be a list of 3d tensors "
"of shape [C, H, W], got {}".format(image.shape))
image = self.normalize(image) # 对图像进行标准化处理
image, target_index = self.resize(image, target_index) # 对图像和对应的bboxes缩放到指定范围
images[i] = image
if targets is not None and target_index is not None:
targets[i] = target_index
# 记录resize后的图像尺寸
image_sizes = [img.shape[-2:] for img in images]
images = self.batch_images(images) # 将images打包成一个batch
image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])
for image_size in image_sizes:
assert len(image_size) == 2
image_sizes_list.append((image_size[0], image_size[1]))
image_list = ImageList(images, image_sizes_list)
return image_list, targets
def resize_boxes(boxes, original_size, new_size):
# type: (Tensor, List[int], List[int]) -> Tensor
"""
将boxes参数根据图像的缩放情况进行相应缩放
Arguments:
original_size: 图像缩放前的尺寸
new_size: 图像缩放后的尺寸
"""
ratios = [
torch.tensor(s, dtype=torch.float32, device=boxes.device) /
torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
for s, s_orig in zip(new_size, original_size)
]
ratios_height, ratios_width = ratios
# Removes a tensor dimension, boxes [minibatch, 4]
# Returns a tuple of all slices along a given dimension, already without it.
xmin, ymin, xmax, ymax = boxes.unbind(1)
xmin = xmin * ratios_width
xmax = xmax * ratios_width
ymin = ymin * ratios_height
ymax = ymax * ratios_height
return torch.stack((xmin, ymin, xmax, ymax), dim=1)
4. RPN源码讲解 —— rpn_function.py
4.1 RPN Head
class RPNHead(nn.Module):
"""
add a RPN head with classification and regression
通过滑动窗口计算预测目标概率与bbox regression参数
Arguments:
in_channels: number of channels of the input feature
num_anchors: number of anchors to be predicted
"""
def __init__(self, in_channels, num_anchors):
super(RPNHead, self).__init__()
# 3x3 滑动窗口 [BS, C, H, W] -> [BS, C, H, W]
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
# 计算预测的目标分数（这里的目标只是指前景或者背景） [BS, C, H, W] -> [BS, num_anchor, H, W]
self.cls_logits = nn.Conv2d(in_channels, num_anchors, kernel_size=1, stride=1) # []
# 计算预测的目标bbox regression参数 [BS, C, H, W] -> [BS, num_anchor*4, H, W]
self.bbox_pred = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=1, stride=1)
# 参数初始化
for layer in self.children():
if isinstance(layer, nn.Conv2d):
torch.nn.init.normal_(layer.weight, std=0.01)
torch.nn.init.constant_(layer.bias, 0)
def forward(self, x):
# type: (List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
# x为图像经过backbone生成的预测特征层
logits = []
bbox_reg = []
for i, feature in enumerate(x): # 迭代预测特征层
t = F.relu(self.conv(feature))
logits.append(self.cls_logits(t)) # 预测分类结果
bbox_reg.append(self.bbox_pred(t)) # 预测bbox结果
return logits, bbox_reg
4.2 Anchors Generator
class AnchorsGenerator(nn.Module):
__annotations__ = {
"cell_anchors": Optional[List[torch.Tensor]],
"_cache": Dict[str, List[torch.Tensor]]
}
"""
anchors生成器
Module that generates anchors for a set of feature maps and
image sizes.
The module support computing anchors at multiple sizes and aspect ratios
per feature map.
sizes and aspect_ratios should have the same number of elements, and it should
correspond to the number of feature maps.
sizes[i] and aspect_ratios[i] can have an arbitrary number of elements,
and AnchorGenerator will output a set of sizes[i] * aspect_ratios[i] anchors
per spatial location for feature map i.
Arguments:
sizes (Tuple[Tuple[int]]): anchor的scale
aspect_ratios (Tuple[Tuple[float]]): 每一个anchor所采用的比例
"""
def __init__(self, sizes=(128, 256, 512), aspect_ratios=(0.5, 1.0, 2.0)):
super(AnchorsGenerator, self).__init__()
# 检查sizes和aspect_ratios里面第一个元素是不是list或tuple数据类型，不是就对其进行修改
if not isinstance(sizes[0], (list, tuple)):
# TODO change this
sizes = tuple((s,) for s in sizes)
if not isinstance(aspect_ratios[0], (list, tuple)):
aspect_ratios = (aspect_ratios,) * len(sizes)
assert len(sizes) == len(aspect_ratios) # 二者的元素个数是否相同
self.sizes = sizes
self.aspect_ratios = aspect_ratios
self.cell_anchors = None
self._cache = {}
def generate_anchors(self, scales, aspect_ratios, dtype=torch.float32, device=torch.device("cpu")):
# type: (List[int], List[float], torch.dtype, torch.device) -> Tensor
"""
compute anchor sizes
Arguments:
scales: sqrt(anchor_area)
aspect_ratios: h/w ratios
dtype: float32
device: cpu/gpu
"""
scales = torch.as_tensor(scales, dtype=dtype, device=device)
aspect_ratios = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)
h_ratios = torch.sqrt(aspect_ratios)
w_ratios = 1.0 / h_ratios
# [r1, r2, r3]' * [s1, s2, s3]
# number of elements is len(ratios)*len(scales)
ws = (w_ratios[:, None] * scales[None, :]).view(-1)
hs = (h_ratios[:, None] * scales[None, :]).view(-1)
# left-top, right-bottom coordinate relative to anchor center(0, 0)
# 生成的anchors模板都是以（0, 0）为中心的, shape [len(ratios)*len(scales), 4]
base_anchors = torch.stack([-ws, -hs, ws, hs], dim=1) / 2
return base_anchors.round() # round 四舍五入
def set_cell_anchors(self, dtype, device):
# type: (torch.dtype, torch.device) -> None
if self.cell_anchors is not None:
cell_anchors = self.cell_anchors
assert cell_anchors is not None
# suppose that all anchors have the same device
# which is a valid assumption in the current state of the codebase
if cell_anchors[0].device == device:
return
# 根据提供的sizes和aspect_ratios生成anchors模板
# anchors模板都是以(0, 0)为中心的anchor
cell_anchors = [
self.generate_anchors(sizes, aspect_ratios, dtype, device)
for sizes, aspect_ratios in zip(self.sizes, self.aspect_ratios)
]
self.cell_anchors = cell_anchors
def num_anchors_per_location(self):
# 计算每个预测特征层上每个滑动窗口的预测目标数
return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)]
# For every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:2),
# output g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a.
def grid_anchors(self, grid_sizes, strides):
# type: (List[List[int]], List[List[Tensor]]) -> List[Tensor]
"""
anchors position in grid coordinate axis map into origin image
计算预测特征图对应原始图像上的所有anchors的坐标
Args:
grid_sizes: 预测特征矩阵的height和width
strides: 预测特征矩阵上一步对应原始图像上的步距
"""
anchors = []
cell_anchors = self.cell_anchors # 预测特征图上anchors的模板
assert cell_anchors is not None
# 遍历每个预测特征层的grid_size，strides和cell_anchors
for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors):
grid_height, grid_width = size # 每一个预测特征图的高度和宽度
stride_height, stride_width = stride # 每一个预测特征图一个cell对应原图的高度和宽度的尺度信息
device = base_anchors.device # 设备信息
"""
>>> torch.arange(0, 3, dtype=torch.float32, device="cuda")
tensor([0., 1., 2.], device='cuda:0')
"""
# For output anchor, compute [x_center, y_center, x_center, y_center]
# shifts_x: shape->[grid_width] 对应原图上的x坐标(列)
shifts_x = torch.arange(0, grid_width, dtype=torch.float32, device=device) * stride_width
# shifts_y: shape->[grid_height] 对应原图上的y坐标(行)
shifts_y = torch.arange(0, grid_height, dtype=torch.float32, device=device) * stride_height
# 计算预测特征矩阵上每个点对应原图上的坐标(anchors模板的坐标偏移量)
# torch.meshgrid函数分别传入行坐标和列坐标，生成网格行坐标矩阵和网格列坐标矩阵
"""
>>> x = torch.arange(0, 5, dtype=torch.float32, device="cuda")
>>> y = torch.arange(5, 10, dtype=torch.float32, device="cuda")
>>> grid = torch.meshgrid(x, y)
>>> grid
(tensor([[0., 0., 0., 0., 0.],
[1., 1., 1., 1., 1.],
[2., 2., 2., 2., 2.],
[3., 3., 3., 3., 3.],
[4., 4., 4., 4., 4.]], device='cuda:0'), tensor([[5., 6., 7., 8., 9.],
[5., 6., 7., 8., 9.],
[5., 6., 7., 8., 9.],
[5., 6., 7., 8., 9.],
[5., 6., 7., 8., 9.]], device='cuda:0'))
"""
# shape: [grid_height, grid_width]
shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
shift_x = shift_x.reshape(-1) # [850]
shift_y = shift_y.reshape(-1)
# 计算anchors坐标(xmin, ymin, xmax, ymax)在原图上的坐标偏移量
# shape: [grid_width*grid_height, 4]
shifts = torch.stack([shift_x, shift_y, shift_x, shift_y], dim=1) # [850, 4]
# For every (base anchor, output anchor) pair,
# offset each zero-centered base anchor by the center of the output anchor.
# 将anchors模板与原图上的坐标偏移量相加得到原图上所有anchors的坐标信息(shape不同时会使用广播机制)
shifts_anchor = shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4) # [850, 15, 4]
# 其中850为feature map的cell个数；15为每个cell生成的anchor个数；4为每个anchor的坐标
anchors.append(shifts_anchor.reshape(-1, 4)) # [12750， 4]：每个anchors的坐标信息
return anchors # List[Tensor(all_num_anchors, 4)]
def cached_grid_anchors(self, grid_sizes, strides):
# type: (List[List[int]], List[List[Tensor]]) -> List[Tensor]
"""将计算得到的所有anchors信息进行缓存"""
key = str(grid_sizes) + str(strides)
# self._cache是字典类型
if key in self._cache:
return self._cache[key]
anchors = self.grid_anchors(grid_sizes, strides) # [12750, 4]得到每一个anchor的坐标信息
self._cache[key] = anchors
return anchors
def forward(self, image_list, feature_maps):
# type: (ImageList, List[Tensor]) -> List[Tensor]
# image_list: ① batch；② 图像尺寸信息
# feature_maps：预测特征层的信息，数据类型为List[Tensor]
# 获取每个预测特征层的尺寸(height, width)
grid_sizes = list([feature_map.shape[-2:] for feature_map in feature_maps])
# 获取输入图像的height和width
image_size = image_list.tensors.shape[-2:]
# 获取变量类型和设备类型
dtype, device = feature_maps[0].dtype, feature_maps[0].device
# one step in feature map equate n pixel stride in origin image
# 计算特征层上的一步等于原始图像上的步长
strides = [[torch.tensor(image_size[0] // g[0], dtype=torch.int64, device=device),
torch.tensor(image_size[1] // g[1], dtype=torch.int64, device=device)] for g in grid_sizes]
# 根据提供的sizes和aspect_ratios生成anchors模板
self.set_cell_anchors(dtype, device)
# 计算/读取所有anchors的坐标信息（这里的anchors信息是映射到原图上的所有anchors信息，不是anchors模板）
# 得到的是一个list列表，对应每张预测特征图映射回原图的anchors坐标信息
anchors_over_all_feature_maps = self.cached_grid_anchors(grid_sizes, strides)
anchors = torch.jit.annotate(List[List[torch.Tensor]], [])
# 遍历一个batch中的每张图像
for i, (image_height, image_width) in enumerate(image_list.image_sizes):
anchors_in_image = []
# 遍历每张预测特征图映射回原图的anchors坐标信息
for anchors_per_feature_map in anchors_over_all_feature_maps:
anchors_in_image.append(anchors_per_feature_map)
anchors.append(anchors_in_image)
# 将每一张图像的所有预测特征层的anchors坐标信息拼接在一起
# anchors是个list，每个元素为一张图像的所有anchors信息
anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors]
# Clear the cache in case that memory leaks.
self._cache.clear()
return anchors
4.3 Region Proposal Network
class RegionProposalNetwork(torch.nn.Module):
"""
Implements Region Proposal Network (RPN).
Arguments:
anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature
maps.
head (nn.Module): module that computes the objectness and regression deltas
fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be
considered as positive during training of the RPN.
bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be
considered as negative during training of the RPN.
batch_size_per_image (int): number of anchors that are sampled during training of the RPN
for computing the loss
positive_fraction (float): proportion of positive anchors in a mini-batch during training
of the RPN
pre_nms_top_n (Dict[str]): number of proposals to keep before applying NMS. It should
contain two fields: training and testing, to allow for different values depending
on training or evaluation
post_nms_top_n (Dict[str]): number of proposals to keep after applying NMS. It should
contain two fields: training and testing, to allow for different values depending
on training or evaluation
nms_thresh (float): NMS threshold used for postprocessing the RPN proposals
"""
__annotations__ = { # 对__init__函数中使用的变量进行注释，该部分不是必须的，只是为了方便理解每个变量的含义
'box_coder': det_utils.BoxCoder,
'proposal_matcher': det_utils.Matcher,
'fg_bg_sampler': det_utils.BalancedPositiveNegativeSampler,
'pre_nms_top_n': Dict[str, int],
'post_nms_top_n': Dict[str, int],
}
def __init__(self, anchor_generator, head,
fg_iou_thresh, bg_iou_thresh,
batch_size_per_image, positive_fraction,
pre_nms_top_n, post_nms_top_n, nms_thresh, score_thresh=0.0):
super(RegionProposalNetwork, self).__init__()
self.anchor_generator = anchor_generator
self.head = head
self.box_coder = det_utils.BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))
# use during training
# 计算anchors与真实bbox的iou
self.box_similarity = box_ops.box_iou
self.proposal_matcher = det_utils.Matcher( # 实例化Matcher类
fg_iou_thresh, # 当iou大于fg_iou_thresh(0.7)时视为正样本
bg_iou_thresh, # 当iou小于bg_iou_thresh(0.3)时视为负样本
allow_low_quality_matches=True
)
self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(
batch_size_per_image, positive_fraction # 256, 0.5
)
# use during testing
self._pre_nms_top_n = pre_nms_top_n
self._post_nms_top_n = post_nms_top_n
self.nms_thresh = nms_thresh
self.score_thresh = score_thresh
self.min_size = 1.
def pre_nms_top_n(self):
if self.training:
return self._pre_nms_top_n['training']
return self._pre_nms_top_n['testing']
def post_nms_top_n(self):
if self.training:
return self._post_nms_top_n['training']
return self._post_nms_top_n['testing']
def assign_targets_to_anchors(self, anchors, targets):
# type: (List[Tensor], List[Dict[str, Tensor]]) -> Tuple[List[Tensor], List[Tensor]]
"""
计算每个anchors最匹配的gt，并划分为正样本，背景以及废弃的样本
Args：
anchors: (List[Tensor])： 每张图像上预测的anchors
targets: (List[Dict[Tensor]) # 包含了GT等信息
Returns:
labels: 标记anchors归属类别（1, 0, -1分别对应正样本，背景，废弃的样本）
注意，在RPN中只有前景和背景，所有正样本的类别都是1，0代表背景
matched_gt_boxes：与anchors匹配的gt
"""
labels = [] # 存储anchors匹配的标签
matched_gt_boxes = [] # 存储anchors匹配的GT
# 遍历每张图像的anchors和targets
for anchors_per_image, targets_per_image in zip(anchors, targets):
gt_boxes = targets_per_image["boxes"] # 只提取GT信息（里面原本包含boxes, labels, image_id, area, iscrowd）
"""
torch.numel(input) → int
返回tensor所有元素的个数（包括0元素）
例子：
>>> a = torch.randint(100, (3, 112, 112))
>>> a.numel()
37632
>>> a = torch.zeros(3, 112, 112)
>>> a.numel()
37632
"""
if gt_boxes.numel() == 0: # 当前图片中没有GT
device = anchors_per_image.device
matched_gt_boxes_per_image = torch.zeros(anchors_per_image.shape, dtype=torch.float32, device=device)
labels_per_image = torch.zeros((anchors_per_image.shape[0],), dtype=torch.float32, device=device)
else: # 当前图片中有GT
# 计算anchors与真实bbox的iou信息
# set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
match_quality_matrix = box_ops.box_iou(gt_boxes, anchors_per_image) # [当前图片GT的个数，该图像生成anchors的总个数]
# 计算每个anchors与gt匹配iou最大的索引（如果iou<0.3索引置为-1，0.3<iou<0.7索引为-2）
matched_idxs = self.proposal_matcher(match_quality_matrix)
# get the targets corresponding GT for each proposal
# NB: need to clamp the indices because we can have a single
# GT in the image, and matched_idxs can be -2, which goes
# out of bounds
# 这里使用clamp设置下限0是为了方便取每个anchors对应的gt_boxes信息
# 负样本和舍弃的样本都是负值，所以为了防止越界直接置为0
# 因为后面是通过labels_per_image变量来记录正样本位置的，
# 所以负样本和舍弃的样本对应的gt_boxes信息并没有什么意义，
# 反正计算目标边界框回归损失时只会用到正样本。
"""
torch.clamp(input, min=None, max=None) → Tensor
参数：
input: 输入tensor
min：元素大小的下限
max：元素大小的上限
返回值：
经过裁剪后的tensor
例子：
>>> a = torch.linspace(-1, 1, 4)
>>> a
tensor([-1.0000, -0.3333, 0.3333, 1.0000])
>>> torch.clamp(a, min=-0.5, max=0.5)
tensor([-0.5000, -0.3333, 0.3333, 0.5000])
"""
matched_gt_boxes_per_image = gt_boxes[matched_idxs.clamp(min=0)] # 将元素<0的全部设置为0
# 记录所有anchors匹配后的标签(正样本处标记为1，负样本处标记为0，丢弃样本处标记为-2)
labels_per_image = matched_idxs >= 0
labels_per_image = labels_per_image.to(dtype=torch.float32) # True -> 1; False -> 0 -> 正样本的位置值为1.0
# background (negative examples)
bg_indices = matched_idxs == self.proposal_matcher.BELOW_LOW_THRESHOLD # -1
labels_per_image[bg_indices] = 0.0 # 负样本的位置值为0.0
# discard indices that are between thresholds
inds_to_discard = matched_idxs == self.proposal_matcher.BETWEEN_THRESHOLDS # -2
labels_per_image[inds_to_discard] = -1.0 # 丢弃样本的位置值为-1.0
labels.append(labels_per_image)
matched_gt_boxes.append(matched_gt_boxes_per_image)
return labels, matched_gt_boxes
def _get_top_n_idx(self, objectness, num_anchors_per_level):
# type: (Tensor, List[int]) -> Tensor
"""
获取每张预测特征图上预测概率排前pre_nms_top_n的anchors索引值
Args:
objectness: Tensor(每张图像的预测目标概率信息)
num_anchors_per_level: List（每个预测特征层上的预测的anchors个数）
Returns:
"""
r = [] # 记录每个预测特征层上预测目标概率前pre_nms_top_n的索引信息
offset = 0
# 遍历每个预测特征层上的预测目标概率信息
"""
tensor.split(长度, dim)
>>> objectness = torch.randint(1, (2, 217413))
>>> num_anchors_per_level = [163200, 40800, 10200, 2550, 663]
>>> for ob in objectness.split(num_anchors_per_level, 1):
... print(ob.shape)
...
torch.Size([2, 163200])
torch.Size([2, 40800])
torch.Size([2, 10200])
torch.Size([2, 2550])
torch.Size([2, 663])
>>> print(f"{a.shape}\n{b.shape}\n{c.shape}\n{d.shape}\n{e.shape}")
torch.Size([2, 163200])
torch.Size([2, 40800])
torch.Size([2, 10200])
torch.Size([2, 2550])
torch.Size([2, 663])
"""
for ob in objectness.split(num_anchors_per_level, 1):
if torchvision._is_tracing():
num_anchors, pre_nms_top_n = _onnx_get_num_anchors_and_pre_nms_top_n(ob, self.pre_nms_top_n())
else:
num_anchors = ob.shape[1] # 预测特征层上的预测的anchors个数
# self.pre_nms_top_n()训练时为2000，测试时为1000
pre_nms_top_n = min(self.pre_nms_top_n(), num_anchors)
# Returns the k largest elements of the given input tensor along a given dimension
"""
tensor.topk(k, 维度, 是否从大到小排序（默认Ture）, 是否排序)
torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)
参数：
① k -> top-k
② dim
③ largest=True： 是否按照从大到小的顺序排序
④ sorted：控制是否按排序顺序返回元素
返回值有两个：
① 返回top-k排序后的数值
② 返回top-k排序后的索引
Example:
>>> x = torch.arange(1, 10)
>>> x
tensor([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> a, b = x.topk(k=5, dim=0)
>>> a
tensor([9, 8, 7, 6, 5])
>>> b
tensor([8, 7, 6, 5, 4])
>>> x.topk(k=5, dim=0)
torch.return_types.topk(
values=tensor([9, 8, 7, 6, 5]),
indices=tensor([8, 7, 6, 5, 4]))
"""
_, top_n_idx = ob.topk(pre_nms_top_n, dim=1) # 只要排序后的索引，不要值
# 这里是将每一个预测特征层的anchors进行遍历，这里对于每一层r.append(top_n_idx)，对于第一层是对的，但对于后面层来说，里面存储
# 的数就不对了，因为里面存储的idx是根据topk得到的，而topk返回的是这一层的idx。
# 简单说，后面层的索引起点位置不应该是0。第二层idx的起点位置应该是第一层最后一个anchors的idx+1。
# 为了达到这个目的，这里使用了offset（偏移量），让这一层结束后让下一层的idx的起点处于正确的位置（而不是从0开始的）
r.append(top_n_idx + offset)
offset += num_anchors
return torch.cat(r, dim=1)
def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_level):
# type: (Tensor, Tensor, List[Tuple[int, int]], List[int]) -> Tuple[List[Tensor], List[Tensor]]
"""
筛除小boxes框，nms处理，根据预测概率获取前post_nms_top_n个目标
Args:
proposals: 预测的bbox坐标
objectness: 预测的目标概率
image_shapes: batch中每张图片的size信息
num_anchors_per_level: 每个预测特征层上预测anchors的数目
Returns:
"""
num_images = proposals.shape[0] # BS
device = proposals.device
# do not backprop throught objectness
objectness = objectness.detach() # 丢弃objectness原有的梯度信息（只获取它的数值信息）
objectness = objectness.reshape(num_images, -1) # [BS, anchor个数, 4个坐标] -> [BS，anchor个数*4个坐标]
# Returns a tensor of size size filled with fill_value
# levels负责记录分隔不同预测特征层上的anchors索引信息
# idx：预测特征层的索引
# n：该预测特征层anchors的个数
"""
torch.full(生成tensor的shape, 填充值, 数据类型, 设备)
>>> a = torch.full((19380, ), 0, dtype=torch.int64, device="cuda")
>>> a.shape
torch.Size([19380])
>>> a
tensor([0, 0, 0, ..., 0, 0, 0], device='cuda:0')
"""
levels = [torch.full((n, ), idx, dtype=torch.int64, device=device)
for idx, n in enumerate(num_anchors_per_level)]
"""
>>> levels = [torch.full((n, ), idx, dtype=torch.int64, device="cuda") for idx, n in enumerate([1000, 600, 300])]
>>> levels = torch.cat(levels, dim=0)
>>> levels.shape
torch.Size([1900])
>>> levels
tensor([0, 0, 0, ..., 2, 2, 2], device='cuda:0')
这样我们就可以用不同的数值（0，1，2，3...）来区分不同的proposals是属于哪一个特征提取层了！
"""
levels = torch.cat(levels, 0)
# Expand this tensor to the same size as objectness
# [所有特征提取层所有anchor的个数] -> [1, 所有特征提取层所有anchor的个数] -> [BS，anchor个数*4个坐标]
levels = levels.reshape(1, -1).expand_as(objectness)
# select top_n boxes independently per level before applying nms
# 获取每张预测特征图上预测概率排前pre_nms_top_n的anchors索引值
top_n_idx = self._get_top_n_idx(objectness, num_anchors_per_level)
image_range = torch.arange(num_images, device=device)
batch_idx = image_range[:, None] # [batch_size, 1]
# 根据每个预测特征层预测概率排前pre_nms_top_n的anchors索引值获取相应概率信息
objectness = objectness[batch_idx, top_n_idx]
levels = levels[batch_idx, top_n_idx]
# 预测概率排前pre_nms_top_n的anchors索引值获取相应bbox坐标信息
proposals = proposals[batch_idx, top_n_idx]
objectness_prob = torch.sigmoid(objectness)
final_boxes = []
final_scores = []
# 遍历每张图像的相关预测信息
for boxes, scores, lvl, img_shape in zip(proposals, objectness_prob, levels, image_shapes):
# 调整预测的boxes信息，将越界的坐标调整到图片边界上
boxes = box_ops.clip_boxes_to_image(boxes, img_shape)
# 返回boxes满足宽，高都大于min_size的索引
keep = box_ops.remove_small_boxes(boxes, self.min_size)
# 获取滤除小目标后的proposal
boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]
# 移除小概率boxes，参考下面这个链接
# https://github.com/pytorch/vision/pull/3205
keep = torch.where(torch.ge(scores, self.score_thresh))[0] # ge: >=
boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]
# non-maximum suppression, independently done per level
# keep是执行nms处理后且按照目标类别分数进行排序后输出的idx
keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
# keep only topk scoring predictions
# 获取前post_nms_top_n个索引
keep = keep[: self.post_nms_top_n()]
# 得到最终的proposal和scores
boxes, scores = boxes[keep], scores[keep]
# 添加到列表中
final_boxes.append(boxes)
final_scores.append(scores)
return final_boxes, final_scores
def compute_loss(self, objectness, pred_bbox_deltas, labels, regression_targets):
# type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor]
"""
计算RPN损失，包括类别损失（前景与背景），bbox regression损失
Arguments:
objectness (Tensor)：预测的前景概率
pred_bbox_deltas (Tensor)：预测的bbox regression
labels (List[Tensor])：真实的标签 1, 0, -1（batch中每一张图片的labels对应List的一个元素中）
regression_targets (List[Tensor])：真实的bbox regression
Returns:
objectness_loss (Tensor) : 类别损失
box_loss (Tensor)：边界框回归损失
"""
# 按照给定的batch_size_per_image, positive_fraction选择正负样本
sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
# 将一个batch中的所有正负样本List(Tensor)分别拼接在一起，并获取非零位置的索引
# sampled_pos_inds = torch.nonzero(torch.cat(sampled_pos_inds, dim=0)).squeeze(1)
sampled_pos_inds = torch.where(torch.cat(sampled_pos_inds, dim=0))[0]
# sampled_neg_inds = torch.nonzero(torch.cat(sampled_neg_inds, dim=0)).squeeze(1)
sampled_neg_inds = torch.where(torch.cat(sampled_neg_inds, dim=0))[0]
# 将所有正负样本索引拼接在一起
sampled_inds = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0)
objectness = objectness.flatten()
labels = torch.cat(labels, dim=0)
regression_targets = torch.cat(regression_targets, dim=0)
# 计算边界框回归损失 -> 只需计算正样本的损失
box_loss = det_utils.smooth_l1_loss(
pred_bbox_deltas[sampled_pos_inds],
regression_targets[sampled_pos_inds],
beta=1 / 9,
size_average=False,
) / (sampled_inds.numel())
# 计算目标预测概率损失，损失函数为BCE；logits表明传入的分数不需要进行任何预处理
objectness_loss = F.binary_cross_entropy_with_logits(
objectness[sampled_inds], labels[sampled_inds]
)
return objectness_loss, box_loss
def forward(self,
images, # type: ImageList
features, # type: Dict[str, Tensor]
targets=None # type: Optional[List[Dict[str, Tensor]]]
):
# type: (...) -> Tuple[List[Tensor], Dict[str, Tensor]]
"""
Arguments:
images (ImageList): images for which we want to compute the predictions
features (Dict[Tensor]): features computed from the images that are
used for computing the predictions. Each tensor in the list
correspond to different feature levels
targets (List[Dict[Tensor]): ground-truth boxes present in the image (optional).
If provided, each element in the dict should contain a field `boxes`,
with the locations of the ground-truth boxes.
Returns:
boxes (List[Tensor]): the predicted boxes from the RPN, one Tensor per
image.
losses (Dict[Tensor]): the losses for the model during training. During
testing, it is an empty dict.
"""
# RPN uses all feature maps that are available
# features是所有预测特征层组成的OrderedDict
features = list(features.values()) # 其中每一个预测特征图层中元素的大小为：[BS, C, H, W]
# 计算每个预测特征层上的预测目标概率objectness和bboxes regression参数pred_bbox_deltas
# objectness和pred_bbox_deltas都是list
# objectness: [BS, 15, H, W]
# pred_bbox_deltas: [BS, 15*4, H, W] = [BS, 60, H, W]
objectness, pred_bbox_deltas = self.head(features)
# 生成一个batch图像的所有anchors信息,list(tensor)元素个数等于batch_size
# images: 两部分 -> ① image_sizes：这1个batch中每张图片的H和W； ② tensors：[BS, C, H, W]
anchors = self.anchor_generator(images, features)
# batch_size
# anchors：是一个list，list中每一个元素代表了每一个图片对应的anchors信息，因为batch=8，所以有8个元素（每个元素的shape为：[14625, 4]）
num_images = len(anchors) # 计算一个batch中有多少张图片（这里为8）
# numel() Returns the total number of elements in the input tensor.
# 计算每个预测特征层上的对应的anchors数量
num_anchors_per_level_shape_tensors = [o[0].shape for o in objectness] # [15, H, W]
# 每个特征矩阵的每个cell会生成15个anchors，而特征矩阵的长度和宽度分别为H和W，所以一共会生成15*H*W个anchors
num_anchors_per_level = [s[0] * s[1] * s[2] for s in num_anchors_per_level_shape_tensors] # [14625]
# 调整内部tensor格式以及shape
objectness, pred_bbox_deltas = concat_box_prediction_layers(objectness,
pred_bbox_deltas)
# apply pred_bbox_deltas to anchors to obtain the decoded proposals
# note that we detach the deltas because Faster R-CNN do not backprop through
# the proposals
# 将预测的bbox regression参数应用到anchors上得到最终预测bbox坐标
proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)
proposals = proposals.view(num_images, -1, 4) # [BS, anchor数量，4]
# 筛除小boxes框，nms处理，根据预测概率获取前post_nms_top_n个目标
boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
"""
如果是训练模式，则计算损失
"""
losses = {}
if self.training:
assert targets is not None
# 计算每个anchors最匹配的gt，并将anchors进行分类，前景，背景以及废弃的anchors
labels, matched_gt_boxes = self.assign_targets_to_anchors(anchors, targets)
# 结合anchors以及对应的gt，计算regression参数
# matched_gt_boxes：每个anchor所匹配的GT
# anchors：每个anchors的坐标
# 根据这两个参数计算回归损失
regression_targets = self.box_coder.encode(matched_gt_boxes, anchors)
loss_objectness, loss_rpn_box_reg = self.compute_loss(
objectness, pred_bbox_deltas, labels, regression_targets
)
losses = {
"loss_objectness": loss_objectness,
"loss_rpn_box_reg": loss_rpn_box_reg
}
return boxes, losses
4.4 Detection utils
import torch
import math
from typing import List, Tuple
from torch import Tensor
class BalancedPositiveNegativeSampler(object):
"""
This class samples batches, ensuring that they contain a fixed proportion of positives
"""
def __init__(self, batch_size_per_image, positive_fraction):
# type: (int, float) -> None
"""
Arguments:
batch_size_per_image (int): number of elements to be selected per image
positive_fraction (float): percentage of positive elements per batch
"""
self.batch_size_per_image = batch_size_per_image
self.positive_fraction = positive_fraction
def __call__(self, matched_idxs):
# type: (List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
"""
Arguments:
matched idxs: list of tensors containing -1, 0 or positive values.
Each tensor corresponds to a specific image.
-1 values are ignored, 0 are considered as negatives and > 0 as
positives.
Returns:
pos_idx (list[tensor])
neg_idx (list[tensor])
Returns two lists of binary masks for each image.
The first list contains the positive elements that were selected,
and the second list the negative example.
"""
pos_idx = []
neg_idx = []
# 遍历每张图像的matched_idxs
for matched_idxs_per_image in matched_idxs:
# >= 1的为正样本, nonzero返回非零元素索引
# positive = torch.nonzero(matched_idxs_per_image >= 1).squeeze(1)
positive = torch.where(torch.ge(matched_idxs_per_image, 1))[0]
# = 0的为负样本
# negative = torch.nonzero(matched_idxs_per_image == 0).squeeze(1)
negative = torch.where(torch.eq(matched_idxs_per_image, 0))[0]
# 指定正样本的数量
num_pos = int(self.batch_size_per_image * self.positive_fraction)
# protect against not enough positive examples
# 如果正样本数量不够就直接采用所有正样本
num_pos = min(positive.numel(), num_pos)
# 指定负样本数量
num_neg = self.batch_size_per_image - num_pos
# protect against not enough negative examples
# 如果负样本数量不够就直接采用所有负样本
num_neg = min(negative.numel(), num_neg)
# randomly select positive and negative examples
# Returns a random permutation of integers from 0 to n - 1.
# 随机选择指定数量的正负样本
"""
>>> perm1 = torch.randperm(200)[:7]
>>> perm2 = torch.randperm(500)[:249]
>>> perm1
tensor([185, 148, 155, 2, 37, 160, 56])
>>> perm2.numel()
249
"""
perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]
pos_idx_per_image = positive[perm1]
neg_idx_per_image = negative[perm2]
# create binary mask from indices
pos_idx_per_image_mask = torch.zeros_like(
matched_idxs_per_image, dtype=torch.uint8
)
neg_idx_per_image_mask = torch.zeros_like(
matched_idxs_per_image, dtype=torch.uint8
)
pos_idx_per_image_mask[pos_idx_per_image] = 1
neg_idx_per_image_mask[neg_idx_per_image] = 1
pos_idx.append(pos_idx_per_image_mask)
neg_idx.append(neg_idx_per_image_mask)
return pos_idx, neg_idx
@torch.jit._script_if_tracing
def encode_boxes(reference_boxes, proposals, weights):
# type: (torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
"""
Encode a set of proposals with respect to some
reference boxes
Arguments:
reference_boxes (Tensor): reference boxes(gt)
proposals (Tensor): boxes to be encoded(anchors)
weights:
"""
# perform some unpacking to make it JIT-fusion friendly
wx = weights[0]
wy = weights[1]
ww = weights[2]
wh = weights[3]
# unsqueeze()
# Returns a new tensor with a dimension of size one inserted at the specified position.
proposals_x1 = proposals[:, 0].unsqueeze(1)
proposals_y1 = proposals[:, 1].unsqueeze(1)
proposals_x2 = proposals[:, 2].unsqueeze(1)
proposals_y2 = proposals[:, 3].unsqueeze(1)
reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1)
reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1)
reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1)
reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1)
# implementation starts here
# parse widths and heights
ex_widths = proposals_x2 - proposals_x1
ex_heights = proposals_y2 - proposals_y1
# parse coordinate of center point
ex_ctr_x = proposals_x1 + 0.5 * ex_widths
ex_ctr_y = proposals_y1 + 0.5 * ex_heights
gt_widths = reference_boxes_x2 - reference_boxes_x1
gt_heights = reference_boxes_y2 - reference_boxes_y1
gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths
gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights
targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
targets_dw = ww * torch.log(gt_widths / ex_widths)
targets_dh = wh * torch.log(gt_heights / ex_heights)
targets = torch.cat((targets_dx, targets_dy, targets_dw, targets_dh), dim=1)
return targets
class BoxCoder(object):
"""
This class encodes and decodes a set of bounding boxes into
the representation used for training the regressors.
"""
def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16)):
# type: (Tuple[float, float, float, float], float) -> None
"""
Arguments:
weights (4-element tuple)
bbox_xform_clip (float)
"""
self.weights = weights
self.bbox_xform_clip = bbox_xform_clip
def encode(self, reference_boxes, proposals):
# type: (List[Tensor], List[Tensor]) -> List[Tensor]
"""
结合anchors和与之对应的gt计算regression参数
Args:
reference_boxes: List[Tensor] 每个proposal/anchor对应的gt_boxes
proposals: List[Tensor] anchors/proposals
Returns: regression parameters
"""
# 统计每张图像的anchors个数，方便后面拼接在一起处理后在分开
# reference_boxes和proposal数据结构相同
boxes_per_image = [len(b) for b in reference_boxes]
reference_boxes = torch.cat(reference_boxes, dim=0)
proposals = torch.cat(proposals, dim=0)
# targets_dx, targets_dy, targets_dw, targets_dh
targets = self.encode_single(reference_boxes, proposals)
return targets.split(boxes_per_image, 0)
def encode_single(self, reference_boxes, proposals):
"""
Encode a set of proposals with respect to some
reference boxes
Arguments:
reference_boxes (Tensor): reference boxes
proposals (Tensor): boxes to be encoded
"""
dtype = reference_boxes.dtype
device = reference_boxes.device
weights = torch.as_tensor(self.weights, dtype=dtype, device=device)
targets = encode_boxes(reference_boxes, proposals, weights)
return targets
def decode(self, rel_codes, boxes):
# type: (Tensor, List[Tensor]) -> Tensor
"""
Args:
rel_codes: bbox regression parameters
boxes: anchors/proposals
Returns:
"""
assert isinstance(boxes, (list, tuple))
assert isinstance(rel_codes, torch.Tensor)
boxes_per_image = [b.size(0) for b in boxes]
concat_boxes = torch.cat(boxes, dim=0)
box_sum = 0
for val in boxes_per_image:
box_sum += val
# 将预测的bbox回归参数应用到对应anchors上得到预测bbox的坐标
pred_boxes = self.decode_single(
rel_codes, concat_boxes
)
# 防止pred_boxes为空时导致reshape报错
if box_sum > 0:
pred_boxes = pred_boxes.reshape(box_sum, -1, 4)
return pred_boxes
def decode_single(self, rel_codes, boxes):
"""
From a set of original boxes and encoded relative box offsets,
get the decoded boxes.
Arguments:
rel_codes (Tensor): encoded boxes (bbox regression parameters)
boxes (Tensor): reference boxes (anchors/proposals)
"""
boxes = boxes.to(rel_codes.dtype)
# xmin, ymin, xmax, ymax
widths = boxes[:, 2] - boxes[:, 0] # anchor/proposal宽度
heights = boxes[:, 3] - boxes[:, 1] # anchor/proposal高度
ctr_x = boxes[:, 0] + 0.5 * widths # anchor/proposal中心x坐标
ctr_y = boxes[:, 1] + 0.5 * heights # anchor/proposal中心y坐标
wx, wy, ww, wh = self.weights # RPN中为[1,1,1,1], fastrcnn中为[10,10,5,5]
dx = rel_codes[:, 0::4] / wx # 预测anchors/proposals的中心坐标x回归参数
dy = rel_codes[:, 1::4] / wy # 预测anchors/proposals的中心坐标y回归参数
dw = rel_codes[:, 2::4] / ww # 预测anchors/proposals的宽度回归参数
dh = rel_codes[:, 3::4] / wh # 预测anchors/proposals的高度回归参数
# limit max value, prevent sending too large values into torch.exp()
# self.bbox_xform_clip=math.log(1000. / 16) 4.135
dw = torch.clamp(dw, max=self.bbox_xform_clip)
dh = torch.clamp(dh, max=self.bbox_xform_clip)
pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
pred_w = torch.exp(dw) * widths[:, None]
pred_h = torch.exp(dh) * heights[:, None]
# xmin
pred_boxes1 = pred_ctr_x - torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w
# ymin
pred_boxes2 = pred_ctr_y - torch.tensor(0.5, dtype=pred_ctr_y.dtype, device=pred_h.device) * pred_h
# xmax
pred_boxes3 = pred_ctr_x + torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w
# ymax
pred_boxes4 = pred_ctr_y + torch.tensor(0.5, dtype=pred_ctr_y.dtype, device=pred_h.device) * pred_h
pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=2).flatten(1)
return pred_boxes
class Matcher(object):
BELOW_LOW_THRESHOLD = -1
BETWEEN_THRESHOLDS = -2
__annotations__ = {
'BELOW_LOW_THRESHOLD': int,
'BETWEEN_THRESHOLDS': int,
}
def __init__(self, high_threshold, low_threshold, allow_low_quality_matches=False):
# type: (float, float, bool) -> None
"""
Args:
high_threshold (float): quality values greater than or equal to
this value are candidate matches.
low_threshold (float): a lower quality threshold used to stratify
matches into three levels:
1) matches >= high_threshold
2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold)
3) BELOW_LOW_THRESHOLD matches in [0, low_threshold)
allow_low_quality_matches (bool): if True, produce additional matches
for predictions that have only low-quality match candidates. See
set_low_quality_matches_ for more details.
"""
self.BELOW_LOW_THRESHOLD = -1
self.BETWEEN_THRESHOLDS = -2
assert low_threshold <= high_threshold
self.high_threshold = high_threshold # 0.7
self.low_threshold = low_threshold # 0.3
self.allow_low_quality_matches = allow_low_quality_matches
def __call__(self, match_quality_matrix):
"""
计算anchors与每个gtboxes匹配的iou最大值，并记录索引，
iou<low_threshold索引值为-1， low_threshold<=iou<high_threshold索引值为-2
Args:
match_quality_matrix (Tensor[float]): an MxN tensor, containing the
pairwise quality between M ground-truth elements and N predicted elements.
Returns:
matches (Tensor[int64]): an N tensor where N[i] is a matched gt in
[0, M - 1] or a negative value indicating that prediction i could not
be matched.
"""
if match_quality_matrix.numel() == 0:
# empty targets or proposals not supported during training
if match_quality_matrix.shape[0] == 0:
raise ValueError(
"No ground-truth boxes available for one of the images "
"during training")
else:
raise ValueError(
"No proposal boxes available for one of the images "
"during training")
# match_quality_matrix is M (gt) x N (predicted)
# Max over gt elements (dim 0) to find best gt candidate for each prediction
# M x N 的每一列代表一个anchors与所有gt的匹配iou值
# matched_vals代表每列的最大值，即每个anchors与所有gt匹配的最大iou值
# matches对应最大值所在的索引
matched_vals, matches = match_quality_matrix.max(dim=0) # the dimension to reduce.
if self.allow_low_quality_matches:
all_matches = matches.clone() # 这里没有用=而是clone，如果用=则是引用，list引用修改的话，list本身也会修改，所以用的是clone
else:
all_matches = None
# Assign candidate matches with low quality to negative (unassigned) values
# 计算iou小于low_threshold的索引
below_low_threshold = matched_vals < self.low_threshold # 得到一个boolean蒙版
# 计算iou在low_threshold与high_threshold之间的索引值
between_thresholds = (matched_vals >= self.low_threshold) & (
matched_vals < self.high_threshold
)
# iou小于low_threshold的matches索引置为-1
matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD # -1
# iou在[low_threshold, high_threshold]之间的matches索引置为-2
matches[between_thresholds] = self.BETWEEN_THRESHOLDS # -2
if self.allow_low_quality_matches:
assert all_matches is not None
self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)
return matches
def set_low_quality_matches_(self, matches, all_matches, match_quality_matrix):
"""
Produce additional matches for predictions that have only low-quality matches.
Specifically, for each ground-truth find the set of predictions that have
maximum overlap with it (including ties); for each prediction in that set, if
it is unmatched, then match it to the ground-truth with which it has the highest
quality value.
为只有低质量匹配的预测生成额外的匹配。
具体来说，对于每个ground-truth，找到与其有最大重叠（包括关系）的一组预测；
对于该集合中的每个预测，如果不匹配，则将其与具有最高质量值的gt实况进行匹配。
"""
# For each gt, find the prediction with which it has highest quality
# 对于每个gt boxes寻找与其iou最大的anchor，
# highest_quality_foreach_gt为匹配到的最大iou值 -> [0.8, 0.85, 0.9, 0.65]
highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1) # the dimension to reduce.
# Find highest quality match available, even if it is low, including ties
# 寻找每个gt boxes与其iou最大的anchor索引，一个gt匹配到的最大iou可能有多个anchor
# gt_pred_pairs_of_highest_quality = torch.nonzero(
# match_quality_matrix == highest_quality_foreach_gt[:, None]
# )
"""
torch.where(condition，a，b)
其中：
输入参数condition：条件限制，如果满足条件，则选择a，否则选择b作为输出。
注意：
a和b是tensor
torch.where(condition) is identical to torch.nonzero(condition, as_tuple=True).
这里的torch.where等价于torch.nonzero
torch.nonzero(input_tensor)：返回输入tensor非零元素的坐标
例子：
>>> torch.nonzero(torch.Tensor([[0.6, 0.0, 0.0, 0.0],
... [0.0, 0.4, 0.0, 0.0],
... [0.0, 0.0, 1.2, 0.0],
... [0.0, 0.0, 0.0,-0.4]]))
==输出==
0 0
1 1
2 2
3 3
对于输出需要一行一行的解读：
第一行：0 0 输入tensor的[0, 0]是一个非零元素
第二行：1 1 输入tensor的[1, 1]是一个非零元素
...
"""
gt_pred_pairs_of_highest_quality = torch.where(
torch.eq(match_quality_matrix, highest_quality_foreach_gt[:, None])
) # [[0, 3], [1, 1], [2, 0], [3, 5]]
# Example gt_pred_pairs_of_highest_quality:
# tensor([[ 0, 39796],
# [ 1, 32055],
# [ 1, 32070],
# [ 2, 39190],
# [ 2, 40255],
# [ 3, 40390],
# [ 3, 41455],
# [ 4, 45470],
# [ 5, 45325],
# [ 5, 46390]])
# Each row is a (gt index, prediction index)
# Note how gt items 1, 2, 3, and 5 each have two ties
# gt_pred_pairs_of_highest_quality[:, 0]代表是对应的gt index(不需要)
# pre_inds_to_update = gt_pred_pairs_of_highest_quality[:, 1]
pre_inds_to_update = gt_pred_pairs_of_highest_quality[1]
# 保留该anchor匹配gt最大iou的索引，即使iou低于设定的阈值
matches[pre_inds_to_update] = all_matches[pre_inds_to_update]
def smooth_l1_loss(input, target, beta: float = 1. / 9, size_average: bool = True):
"""
very similar to the smooth_l1_loss from pytorch, but with
the extra beta parameter
"""
n = torch.abs(input - target)
# cond = n < beta
cond = torch.lt(n, beta)
loss = torch.where(cond, 0.5 * n ** 2 / beta, n - 0.5 * beta)
if size_average:
return loss.mean()
return loss.sum()
4.5 ROI Align、Two MLP Head、 Faster R-CNN Predictor
4.5.1 ROI Align
ROI Align对应的图中的ROI Pooling，代码如下：
from torchvision.ops import MultiScaleRoIAlign
# Multi-scale RoIAlign pooling
if box_roi_pool is None:
box_roi_pool = MultiScaleRoIAlign( # 这里就是ROI Pooling
featmap_names=['0', '1', '2', '3'], # 在哪些特征层进行roi pooling
output_size=[7, 7], # 这里给出了ROI Pooling后输出的shape
sampling_ratio=2)
该方法PyTorch官方已经实现，且已被封装，直接用就可以
4.5.2 Two MLP Head
这里是将经过ROI Pooling后的proposal feature map送入这两个全连接层，代码如下：
from typing import Optional, List, Dict, Tuple
import torch
from torch import Tensor
import torch.nn.functional as F
from . import det_utils
from . import boxes as box_ops
def fastrcnn_loss(class_logits, box_regression, labels, regression_targets):
# type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor]
"""
Computes the loss for Faster R-CNN.
Arguments:
class_logits : 预测类别概率信息，shape=[num_anchors, num_classes]
box_regression : 预测边目标界框回归信息
labels : 真实类别信息
regression_targets : 真实目标边界框信息
Returns:
classification_loss (Tensor)
box_loss (Tensor)
"""
labels = torch.cat(labels, dim=0)
regression_targets = torch.cat(regression_targets, dim=0)
# 计算类别损失信息
classification_loss = F.cross_entropy(class_logits, labels)
# get indices that correspond to the regression targets for
# the corresponding ground truth labels, to be used with
# advanced indexing
# 返回标签类别大于0的索引
# sampled_pos_inds_subset = torch.nonzero(torch.gt(labels, 0)).squeeze(1)
sampled_pos_inds_subset = torch.where(torch.gt(labels, 0))[0]
# 返回标签类别大于0位置的类别信息
labels_pos = labels[sampled_pos_inds_subset]
# shape=[num_proposal, num_classes]
N, num_classes = class_logits.shape
box_regression = box_regression.reshape(N, -1, 4)
# 计算边界框损失信息
box_loss = det_utils.smooth_l1_loss(
# 获取指定索引proposal的指定类别box信息
box_regression[sampled_pos_inds_subset, labels_pos],
regression_targets[sampled_pos_inds_subset],
beta=1 / 9,
size_average=False,
) / labels.numel()
return classification_loss, box_loss
class RoIHeads(torch.nn.Module):
__annotations__ = {
'box_coder': det_utils.BoxCoder,
'proposal_matcher': det_utils.Matcher,
'fg_bg_sampler': det_utils.BalancedPositiveNegativeSampler,
}
def __init__(self,
box_roi_pool, # Multi-scale RoIAlign pooling
box_head, # TwoMLPHead
box_predictor, # FastRCNNPredictor
# Faster R-CNN training
fg_iou_thresh, bg_iou_thresh, # default: 0.5, 0.5
batch_size_per_image, positive_fraction, # default: 512, 0.25
bbox_reg_weights, # None
# Faster R-CNN inference
score_thresh, # default: 0.05
nms_thresh, # default: 0.5
detection_per_img): # default: 100
super(RoIHeads, self).__init__()
"""
.box_iou是一个方法，用来计算IoU值
"""
self.box_similarity = box_ops.box_iou # 将计算IoU的方法赋值给self.box_similarity
# assign ground-truth boxes for each proposal
"""
Matcher在RPN中也使用到了，作用是将proposal划分到正负样本当中
"""
self.proposal_matcher = det_utils.Matcher( # 将Matcher类赋值给self.proposal_matcher
fg_iou_thresh, # default: 0.5
bg_iou_thresh, # default: 0.5
allow_low_quality_matches=False)
"""
BalancedPositiveNegativeSampler在RPN时也用到了，作用是将划分好的正负样本进行采样
参数：
batch_size_per_image： 总共采样512个样本
positive_fraction：正样本占25%
"""
self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(
batch_size_per_image, # default: 512
positive_fraction) # default: 0.25
if bbox_reg_weights is None:
bbox_reg_weights = (10., 10., 5., 5.) # 超参数
self.box_coder = det_utils.BoxCoder(bbox_reg_weights)
self.box_roi_pool = box_roi_pool # Multi-scale RoIAlign pooling
self.box_head = box_head # TwoMLPHead
self.box_predictor = box_predictor # FastRCNNPredictor
self.score_thresh = score_thresh # default: 0.05
self.nms_thresh = nms_thresh # default: 0.5
self.detection_per_img = detection_per_img # default: 100
def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels):
# type: (List[Tensor], List[Tensor], List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
"""
为每个proposal匹配对应的gt_box，并划分到正负样本中
Args:
proposals:
gt_boxes:
gt_labels:
Returns:
"""
matched_idxs = []
labels = []
# 遍历每张图像的proposals, gt_boxes, gt_labels信息
for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(proposals, gt_boxes, gt_labels):
if gt_boxes_in_image.numel() == 0: # 该张图像中没有gt框，为背景
# background image
device = proposals_in_image.device
clamped_matched_idxs_in_image = torch.zeros(
(proposals_in_image.shape[0],), dtype=torch.int64, device=device
)
labels_in_image = torch.zeros(
(proposals_in_image.shape[0],), dtype=torch.int64, device=device
)
else:
# set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
# 计算proposal与每个gt_box的iou重合度
match_quality_matrix = box_ops.box_iou(gt_boxes_in_image, proposals_in_image)
# 计算proposal与每个gt_box匹配的iou最大值，并记录索引，
# iou < low_threshold索引值为 -1， low_threshold <= iou < high_threshold索引值为 -2
matched_idxs_in_image = self.proposal_matcher(match_quality_matrix)
# 限制最小值，防止匹配标签时出现越界的情况
# 注意-1, -2对应的gt索引会调整到0,获取的标签类别为第0个gt的类别（实际上并不是）,后续会进一步处理
clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(min=0)
# 获取proposal匹配到的gt对应标签
labels_in_image = gt_labels_in_image[clamped_matched_idxs_in_image]
labels_in_image = labels_in_image.to(dtype=torch.int64)
# label background (below the low threshold)
# 将gt索引为-1的类别设置为0，即背景，负样本
bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD # -1
labels_in_image[bg_inds] = 0
# label ignore proposals (between low and high threshold)
# 将gt索引为-2的类别设置为-1, 即废弃样本
ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS # -2
labels_in_image[ignore_inds] = -1 # -1 is ignored by sampler
matched_idxs.append(clamped_matched_idxs_in_image)
labels.append(labels_in_image)
return matched_idxs, labels
def subsample(self, labels):
# type: (List[Tensor]) -> List[Tensor]
# BalancedPositiveNegativeSampler
sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
sampled_inds = []
# 遍历每张图片的正负样本索引
for img_idx, (pos_inds_img, neg_inds_img) in enumerate(zip(sampled_pos_inds, sampled_neg_inds)):
# 记录所有采集样本索引（包括正样本和负样本）
# img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img).squeeze(1)
img_sampled_inds = torch.where(pos_inds_img | neg_inds_img)[0]
sampled_inds.append(img_sampled_inds)
return sampled_inds
def add_gt_proposals(self, proposals, gt_boxes):
# type: (List[Tensor], List[Tensor]) -> List[Tensor]
"""
将gt_boxes拼接到proposal后面
Args:
proposals: 一个batch中每张图像rpn预测的boxes
gt_boxes: 一个batch中每张图像对应的真实目标边界框
Returns:
"""
proposals = [
torch.cat((proposal, gt_box))
for proposal, gt_box in zip(proposals, gt_boxes)
]
return proposals
def check_targets(self, targets):
# type: (Optional[List[Dict[str, Tensor]]]) -> None
assert targets is not None
assert all(["boxes" in t for t in targets])
assert all(["labels" in t for t in targets])
def select_training_samples(self,
proposals, # type: List[Tensor]
targets # type: Optional[List[Dict[str, Tensor]]]
):
# type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
"""
划分正负样本，统计对应gt的标签以及边界框回归信息
list元素个数为batch_size
Args:
proposals: rpn预测的boxes
targets:
Returns:
"""
# 检查target数据是否为空
self.check_targets(targets)
# 如果不加这句，jit.script会不通过(看不懂)
assert targets is not None
dtype = proposals[0].dtype
device = proposals[0].device
# 获取标注好的boxes以及labels信息
gt_boxes = [t["boxes"].to(dtype) for t in targets]
gt_labels = [t["labels"] for t in targets]
# append ground-truth bboxes to proposal
# 将gt_boxes拼接到proposal后面
proposals = self.add_gt_proposals(proposals, gt_boxes)
# get matching gt indices for each proposal
# 为每个proposal匹配对应的gt_box，并划分到正负样本中
matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)
# sample a fixed proportion of positive-negative proposals
# 按给定数量和比例采样正负样本
sampled_inds = self.subsample(labels)
matched_gt_boxes = []
num_images = len(proposals)
# 遍历每张图像
for img_id in range(num_images):
# 获取每张图像的正负样本索引
img_sampled_inds = sampled_inds[img_id]
# 获取对应正负样本的proposals信息
proposals[img_id] = proposals[img_id][img_sampled_inds]
# 获取对应正负样本的真实类别信息
labels[img_id] = labels[img_id][img_sampled_inds]
# 获取对应正负样本的gt索引信息
matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds]
gt_boxes_in_image = gt_boxes[img_id]
if gt_boxes_in_image.numel() == 0:
gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device)
# 获取对应正负样本的gt box信息
matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]])
# 根据gt和proposal计算边框回归参数（针对gt的）
regression_targets = self.box_coder.encode(matched_gt_boxes, proposals)
return proposals, labels, regression_targets
def postprocess_detections(self,
class_logits, # type: Tensor
box_regression, # type: Tensor
proposals, # type: List[Tensor]
image_shapes # type: List[Tuple[int, int]]
):
# type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
"""
对网络的预测数据进行后处理，包括
（1）根据proposal以及预测的回归参数计算出最终bbox坐标
（2）对预测类别结果进行softmax处理
（3）裁剪预测的boxes信息，将越界的坐标调整到图片边界上
（4）移除所有背景信息
（5）移除低概率目标
（6）移除小尺寸目标
（7）执行nms处理，并按scores进行排序
（8）根据scores排序返回前topk个目标
Args:
class_logits: 网络预测类别概率信息
box_regression: 网络预测的边界框回归参数
proposals: rpn输出的proposal
image_shapes: 打包成batch前每张图像的宽高
Returns:
"""
device = class_logits.device
# 预测目标类别数
num_classes = class_logits.shape[-1]
# 获取每张图像的预测bbox数量
boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]
# 根据proposal以及预测的回归参数计算出最终bbox坐标
pred_boxes = self.box_coder.decode(box_regression, proposals)
# 对预测类别结果进行softmax处理
pred_scores = F.softmax(class_logits, -1)
# split boxes and scores per image
# 根据每张图像的预测bbox数量分割结果
pred_boxes_list = pred_boxes.split(boxes_per_image, 0)
pred_scores_list = pred_scores.split(boxes_per_image, 0)
all_boxes = []
all_scores = []
all_labels = []
# 遍历每张图像预测信息
for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):
# 裁剪预测的boxes信息，将越界的坐标调整到图片边界上
boxes = box_ops.clip_boxes_to_image(boxes, image_shape)
# create labels for each prediction
labels = torch.arange(num_classes, device=device)
labels = labels.view(1, -1).expand_as(scores)
# remove prediction with the background label
# 移除索引为0的所有信息（0代表背景）
boxes = boxes[:, 1:]
scores = scores[:, 1:]
labels = labels[:, 1:]
# batch everything, by making every class prediction be a separate instance
boxes = boxes.reshape(-1, 4)
scores = scores.reshape(-1)
labels = labels.reshape(-1)
# remove low scoring boxes
# 移除低概率目标，self.scores_thresh=0.05
# gt: Computes input > other element-wise.
# inds = torch.nonzero(torch.gt(scores, self.score_thresh)).squeeze(1)
inds = torch.where(torch.gt(scores, self.score_thresh))[0]
boxes, scores, labels = boxes[inds], scores[inds], labels[inds]
# remove empty boxes
# 移除小目标
keep = box_ops.remove_small_boxes(boxes, min_size=1.)
boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
# non-maximun suppression, independently done per class
# 执行nms处理，执行后的结果会按照scores从大到小进行排序返回
keep = box_ops.batched_nms(boxes, scores, labels, self.nms_thresh)
# keep only topk scoring predictions
# 获取scores排在前topk个预测目标
keep = keep[:self.detection_per_img]
boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
all_boxes.append(boxes)
all_scores.append(scores)
all_labels.append(labels)
return all_boxes, all_scores, all_labels
def forward(self,
features, # type: Dict[str, Tensor]
proposals, # type: List[Tensor]
image_shapes, # type: List[Tuple[int, int]]
targets=None # type: Optional[List[Dict[str, Tensor]]]
):
# type: (...) -> Tuple[List[Dict[str, Tensor]], Dict[str, Tensor]]
"""
Arguments:
features (List[Tensor])： 输入图片经过backbone生成的特征图
proposals (List[Tensor[N, 4]])： RPN生成的proposals
image_shapes (List[Tuple[H, W]])： 图像预处理之后的大小
targets (List[Dict])： GT的annotation信息
"""
# 检查targets的数据类型是否正确
if targets is not None:
for t in targets:
floating_point_types = (torch.float, torch.double, torch.half)
assert t["boxes"].dtype in floating_point_types, "target boxes must of float type"
assert t["labels"].dtype == torch.int64, "target labels must of int64 type"
"""
如果是训练模式，则需选取使用的样本。这是因为RPN在训练模式下会保留2000个proposal，但在训练时只需从中采样512个即可；
如果是验证模式，则RPN仅会保留1000个proposal
"""
if self.training:
# 划分正负样本，统计对应gt的标签以及边界框回归信息
proposals, labels, regression_targets = self.select_training_samples(proposals, targets)
else: # eval模式下没有GT
labels = None
regression_targets = None
# 将采集样本通过Multi-scale RoIAlign pooling层
# box_features_shape: [num_proposals, channel, height, width]
"""
这里的box_roi_pool就是ROI Pooling，经过它的proposal.shape变为(7, 7)
参数：
1. features：输入预测层（如果是MobileNet v2那么仅有一个，如果是ResNet50+FPN则有4个）
2. proposals：每张图保留的512个proposal
3. image_shapes：每张图缩放之后得到的尺寸
这里图片缩放的并不是简单的Resize，而是将图片与想要尺寸在左上角进行对齐，不足的地方用0填充
这里图片的缩放后的尺寸是我们在dataloader里定义的尺寸，程序里是（224, 224）
(7,7)是输入预测特征层的proposal经过ROI Pooling后的尺寸
返回值：
box_features: 经过ROI Pooling返回的每个proposal对应的feature maps，shape为torch.Size([1024, 256, 7, 7])
1024是表示这是每个batch的累加（512 + 512 = 1024）
"""
box_features = self.box_roi_pool(features, proposals, image_shapes)
# 通过roi_pooling后的两层全连接层
# box_features_shape: [num_proposals, representation_size]
"""
box_head: 对应着图片上的Two MLP Head
返回值：
经过两个FC返回的一维向量。前一个表示BS，后一个表示FC的输出维度
"""
box_features = self.box_head(box_features)
# 接着分别预测目标类别和边界框回归参数（并行结构）
"""
再将Two MLP Head的值分别经过两个并行的FC得到：
1. 类别分数 torch.Size([1024, 21])
21 = 20(NC) + 1(负样本)
2. 预测回归参数 torch.Size([1024, 84])
84 = 21 * 4
"""
class_logits, box_regression = self.box_predictor(box_features)
result = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) # 定义一个空的list列表
losses = {} # 定义一个空的字典
"""
如果是训练模式，则会计算Fast R-CNN的损失，存入到losses这个dict中
如果是eval模式，则不需要计算损失，直接对结果进行后处理即可
将低概率的目标剔除、NMS处理等等
Note:
训练模式我们不需要看框预测的效果，因为没有必要，我们只需要知道loss就行，根据loss优化网络才是训练模式的目的
看预测框效果那是eval模式下才应该做的
"""
if self.training:
assert labels is not None and regression_targets is not None
loss_classifier, loss_box_reg = fastrcnn_loss(
class_logits, box_regression, labels, regression_targets)
losses = {
"loss_classifier": loss_classifier,
"loss_box_reg": loss_box_reg
}
else:
"""
boxes: 最终预测的目标边界框
scores：每个类别的分数
labels: 对应标签
"""
boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)
num_images = len(boxes)
for i in range(num_images):
result.append(
{
"boxes": boxes[i],
"labels": labels[i],
"scores": scores[i],
}
)
return result, losses
4.5.3 Faster R-CNN Predictor
这部分就是将Two MLP Head的输出结果送入这个预测头，得到：
- 训练模式
- 类别损失和预测框回归损失
- eval模式
- 类别分数和预测框
代码如下：
class FastRCNNPredictor(nn.Module):
"""
Standard classification + bounding box regression layers
for Fast R-CNN.
Arguments:
in_channels (int): number of input channels
num_classes (int): number of output classes (including background)
"""
def __init__(self, in_channels, num_classes):
super(FastRCNNPredictor, self).__init__()
self.cls_score = nn.Linear(in_channels, num_classes) # 1024 -> 21(VOC)
self.bbox_pred = nn.Linear(in_channels, num_classes * 4) # 1024 -> 21*4=84(VOC)
def forward(self, x): # x.shape: torch.Size([1024, 1024])
"""
>>> x = torch.randint(1, (3, 112, 112))
>>> x.shape
torch.Size([3, 112, 112])
>>> x.dim()
3
"""
if x.dim() == 4:
assert list(x.shape[2:]) == [1, 1]
x = x.flatten(start_dim=1) # 这里的flatten其实没有什么必要
scores = self.cls_score(x) # 预测目标概率分数 torch.Size([1024, 21])
bbox_deltas = self.bbox_pred(x) # 预测目标回归参数 torch.Size([1024, 84])
return scores, bbox_deltas
4.6 Fast R-CNN正负样本划分及采样
在训练模式中，并不是使用RPN网络生成的所有proposal，而是从中选取一部分proposal用于Fast R-CNN的损失计算。如图黄色的部分：
4.6.1 select_training_samples
if self.training:
# 划分正负样本，统计对应gt的标签以及边界框回归信息
"""
proposals：RPN生成的proposal进行了正负样本的划分
labels：对应的标签
regression_targets：对应GT box的回归参数
"""
proposals, labels, regression_targets = self.select_training_samples(proposals, targets)
else: # eval模式下没有GT
labels = None
regression_targets = None
def select_training_samples(self,
proposals, # type: List[Tensor]
targets # type: Optional[List[Dict[str, Tensor]]]
):
# type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
"""
划分正负样本，统计对应gt的标签以及边界框回归信息
list元素个数为batch_size
Args:
proposals: rpn预测的boxes
targets: GT信息
Returns:
"""
# 检查target数据是否为空
self.check_targets(targets)
# 如果不加这句，jit.script会不通过(看不懂)
assert targets is not None
dtype = proposals[0].dtype
device = proposals[0].device
# 获取标注好的boxes以及labels信息
"""
gt_boxes = [t["boxes"].to(dtype) for t in targets]
gt_labels = [t["labels"] for t in targets]
分别提取targets中的"boxes"和"labels"信息，并分别用[]包裹
"""
gt_boxes = [t["boxes"].to(dtype) for t in targets]
gt_labels = [t["labels"] for t in targets]
# append ground-truth bboxes to proposal
# 将gt_boxes拼接到proposal后面
proposals = self.add_gt_proposals(proposals, gt_boxes)
# get matching gt indices for each proposal
# 为每个proposal匹配对应的gt_box，并划分到正负样本中
"""
通过assign_targets_to_proposals这个方法为刚刚添加gt_boxes的proposals划分到正负样本中
参数：
1. proposals：self.add_gt_proposals(proposals, gt_boxes)后的proposals
2. gt_boxes: GT的boxes
3. gt_labels: GT的labels
返回值：
matched_idxs：每个proposal所匹配到的gt boxes的索引
labels：每个proposal所匹配到的gt boxes的标签
"""
matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)
# sample a fixed proportion of positive-negative proposals
# 按给定数量和比例采样正负样本
"""
在训练时并不是将RPN生成的proposal都用来计算损失，还需经过采样（一般为512）
"""
sampled_inds = self.subsample(labels)
matched_gt_boxes = []
num_images = len(proposals)
# 遍历每张图像
for img_id in range(num_images):
# 获取每张图像的正负样本索引
img_sampled_inds = sampled_inds[img_id]
# 获取对应正负样本的proposals信息
proposals[img_id] = proposals[img_id][img_sampled_inds] # [2003, 4] -> [512, 4]
# 获取对应正负样本的真实类别信息
labels[img_id] = labels[img_id][img_sampled_inds] # [2003,] -> [512,]
# 获取对应正负样本的gt索引信息
matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds]
# 提取gt boxes的信息； img_id为图片的id
gt_boxes_in_image = gt_boxes[img_id] # torch.Size([3, 4])， 其中3表示该图片有3个gt boxes，4为对应的坐标
if gt_boxes_in_image.numel() == 0:
gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device)
# 获取对应正负样本的gt box信息
"""
matched_gt_boxes: torch.Size([512, 4])
其中：
512为采样的proposal对应的gt box
4为对应gt box的坐标
Note:
是proposal对应的gt box而不是proposal
"""
matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]])
# 根据gt和proposal计算边框回归参数
# Note: 这里的回归参数是 针对gt而言的
"""
1. matched_gt_boxes: proposal匹配到的gt box
2. proposals: 采样后的proposal, shape: [512, 4]
regression_targets: 根据gt box和proposal得到的关于x,y,h,w的回归参数
用list保存的，每一张图片为一个元素
"""
regression_targets = self.box_coder.encode(matched_gt_boxes, proposals)
"""
Return
proposals： 添加了gt的proposals
labels： 对应的标签
regression_targets: 根据gt box和proposal得到的关于x,y,h,w的回归参数
"""
return proposals, labels, regression_targets
4.7 Fast R-CNN的损失计算
import warnings
from collections import OrderedDict
from typing import Tuple, List, Dict, Optional, Union
import torch
from torch import nn, Tensor
import torch.nn.functional as F
from torchvision.ops import MultiScaleRoIAlign
from .roi_head import RoIHeads
from .transform import GeneralizedRCNNTransform
from .rpn_function import AnchorsGenerator, RPNHead, RegionProposalNetwork
class FasterRCNNBase(nn.Module):
"""
Main class for Generalized R-CNN.
Arguments:
backbone (nn.Module): 特征提取网络部分
rpn (nn.Module): 候选框生成部分
roi_heads (nn.Module): takes the features + the proposals from the RPN and computes
detections / masks from it.
transform (nn.Module): performs the data transformation from the inputs to feed into
the model
"""
def __init__(self, backbone, rpn, roi_heads, transform):
super(FasterRCNNBase, self).__init__()
self.transform = transform
self.backbone = backbone
self.rpn = rpn
self.roi_heads = roi_heads
# used only on torchscript mode
self._has_warned = False
@torch.jit.unused
def eager_outputs(self, losses, detections):
# type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]
if self.training:
return losses
return detections
def forward(self, images, targets=None):
# type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
"""
Arguments:
images (list[Tensor]): images to be processed
targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)
Returns:
result (list[BoxList] or dict[Tensor]): the output from the model.
During training, it returns a dict[Tensor] which contains the losses.
During testing, it returns list[BoxList] contains additional fields
like `scores`, `labels` and `mask` (for Mask R-CNN models).
"""
if self.training and targets is None:
raise ValueError("In training mode, targets should be passed")
if self.training:
assert targets is not None
for target in targets: # 进一步判断传入的target的boxes参数是否符合规定
boxes = target["boxes"]
if isinstance(boxes, torch.Tensor):
if len(boxes.shape) != 2 or boxes.shape[-1] != 4:
raise ValueError("Expected target boxes to be a tensor"
"of shape [N, 4], got {:}.".format(
boxes.shape))
else:
raise ValueError("Expected target boxes to be of type "
"Tensor, got {:}.".format(type(boxes)))
# original_image_sizes：存储每张图片原始的尺寸，为了之后可以映射到原图中
original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], [])
for img in images:
val = img.shape[-2:] # [H, W]
assert len(val) == 2 # 防止输入的是个一维向量
original_image_sizes.append((val[0], val[1]))
# original_image_sizes = [img.shape[-2:] for img in images]
images, targets = self.transform(images, targets) # 对图像进行预处理
# [(images, targets), (images, targets), ...]这才是需要送入网络的batch
# print(images.tensors.shape)
features = self.backbone(images.tensors) # 将图像输入backbone得到特征图
if isinstance(features, torch.Tensor): # 若只在一层特征层上预测，将feature放入有序字典中，并编号为‘0’
features = OrderedDict([('0', features)]) # 若在多层特征层上预测，传入的就是一个有序字典
# 将特征层以及标注target信息传入rpn中
"""
Input:
1. images：输入图片(shape -> torch.Size([2, 3, 800, 1088]))
2. features: 预测特征层（对于MobileNet v2只有一个，对于Resnet50+FPN有4个 -> "0", "1", "2", "3", "pool"）
"0": torch.Size([2, 256, 200, 272])
"1": torch.Size([2, 256, 100, 136])
"2": torch.Size([2, 256, 50, 68])
"3": torch.Size([2, 256, 25, 34])
"pool": torch.Size([2, 256, 13, 17])
3. targets：annotation的一些信息
这里batch size设置为2，故有该list由两部分组成，每一个均有下面的属性
"boxes": torch.Size([1, 4])： 坐标信息
"label": 4
"image_id": 4
"area": 4
"iscrowd": 4
Output:
1. proposal：预测框，是一个list，最外层元素数量等于batch_size。
每一个元素的shape为：[2000, 4] -> RPN生成2000个预测框（proposal）和其4个坐标
2. proposal_losses：预测框与GT的损失，也是分为两个部分
1. loss_objectness：类别损失（是一个标量）
2. loss_rpn_box_reg：回归损失（是一个标量）
"""
# proposals: List[Tensor], Tensor_shape: [num_proposals, 4],
# 每个proposals是绝对坐标，且为(x1, y1, x2, y2)格式
proposals, proposal_losses = self.rpn(images, features, targets)
# 将rpn生成的数据以及标注target信息传入fast rcnn后半部分
"""
detections: 最终的框
detector_losses： 最终框的损失
"""
detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
# 对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)
losses = {}
losses.update(detector_losses)
losses.update(proposal_losses)
if torch.jit.is_scripting():
if not self._has_warned:
warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")
self._has_warned = True
return losses, detections
else:
return self.eager_outputs(losses, detections)
# if self.training:
# return losses
#
# return detections
class TwoMLPHead(nn.Module):
"""
Standard heads for FPN-based models
Arguments:
in_channels (int): number of input channels
representation_size (int): size of the intermediate representation
"""
def __init__(self, in_channels, representation_size):
super(TwoMLPHead, self).__init__()
self.fc6 = nn.Linear(in_channels, representation_size)
self.fc7 = nn.Linear(representation_size, representation_size) # 输入等于输出，representation_size为固定值1024
def forward(self, x):
"""
这里的x就是通过ROI Pooling的输出，即每个proposal对应的特征矩阵
x.shape: torch.Size([1024, 256, 7, 7])
1024: 整个batch（这里bs=2）代表的总的proposal的个数（512 + 512）
256： channel
7,7: height and width
"""
x = x.flatten(start_dim=1) # [BS, C, H, W] -> [BS, C*H*W] torch.Size([1024, 12544])
x = F.relu(self.fc6(x)) # 图中FC_1
x = F.relu(self.fc7(x)) # 图中FC_2
return x
class FastRCNNPredictor(nn.Module):
"""
Standard classification + bounding box regression layers
for Fast R-CNN.
Arguments:
in_channels (int): number of input channels
num_classes (int): number of output classes (including background)
"""
def __init__(self, in_channels, num_classes):
super(FastRCNNPredictor, self).__init__()
self.cls_score = nn.Linear(in_channels, num_classes) # 1024 -> 21(VOC)
self.bbox_pred = nn.Linear(in_channels, num_classes * 4) # 1024 -> 21*4=84(VOC)
def forward(self, x): # x.shape: torch.Size([1024, 1024])
"""
>>> x = torch.randint(1, (3, 112, 112))
>>> x.shape
torch.Size([3, 112, 112])
>>> x.dim()
3
"""
if x.dim() == 4:
assert list(x.shape[2:]) == [1, 1]
x = x.flatten(start_dim=1) # 这里的flatten其实没有什么必要
scores = self.cls_score(x) # 预测目标概率分数 torch.Size([1024, 21])
bbox_deltas = self.bbox_pred(x) # 预测目标回归参数 torch.Size([1024, 84])
return scores, bbox_deltas
class FasterRCNN(FasterRCNNBase):
"""
Implements 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.
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
between 0 and H and 0 and W
- labels (Int64Tensor[N]): the class label for each ground-truth box
The model returns a Dict[Tensor] during training, containing the classification and regression
losses for both the RPN and the 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 between
0 and H and 0 and W
- labels (Int64Tensor[N]): the predicted labels for each image
- scores (Tensor[N]): the scores or each prediction
Arguments:
backbone (nn.Module): the network used to compute the features for the model.
It should contain a out_channels attribute, which indicates the number of output
channels that each feature map has (and it should be the same for all feature maps).
The backbone should return a single Tensor or and OrderedDict[Tensor].
num_classes (int): number of output classes of the model (including the background).
If box_predictor is specified, num_classes should be None.
min_size (int): minimum size of the image to be rescaled before feeding it to the backbone
max_size (int): maximum size of the image to be rescaled before feeding it to the backbone
image_mean (Tuple[float, float, float]): mean values used for input normalization.
They are generally the mean values of the dataset on which the backbone has been trained
on
image_std (Tuple[float, float, float]): std values used for input normalization.
They are generally the std values of the dataset on which the backbone has been trained on
rpn_anchor_generator (AnchorGenerator): module that generates the anchors for a set of feature
maps.
rpn_head (nn.Module): module that computes the objectness and regression deltas from the RPN
rpn_pre_nms_top_n_train (int): number of proposals to keep before applying NMS during training
rpn_pre_nms_top_n_test (int): number of proposals to keep before applying NMS during testing
rpn_post_nms_top_n_train (int): number of proposals to keep after applying NMS during training
rpn_post_nms_top_n_test (int): number of proposals to keep after applying NMS during testing
rpn_nms_thresh (float): NMS threshold used for postprocessing the RPN proposals
rpn_fg_iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be
considered as positive during training of the RPN.
rpn_bg_iou_thresh (float): maximum IoU between the anchor and the GT box so that they can be
considered as negative during training of the RPN.
rpn_batch_size_per_image (int): number of anchors that are sampled during training of the RPN
for computing the loss
rpn_positive_fraction (float): proportion of positive anchors in a mini-batch during training
of the RPN
rpn_score_thresh (float): during inference, only return proposals with a classification score
greater than rpn_score_thresh
box_roi_pool (MultiScaleRoIAlign): the module which crops and resizes the feature maps in
the locations indicated by the bounding boxes
box_head (nn.Module): module that takes the cropped feature maps as input
box_predictor (nn.Module): module that takes the output of box_head and returns the
classification logits and box regression deltas.
box_score_thresh (float): during inference, only return proposals with a classification score
greater than box_score_thresh
box_nms_thresh (float): NMS threshold for the prediction head. Used during inference
box_detections_per_img (int): maximum number of detections per image, for all classes.
box_fg_iou_thresh (float): minimum IoU between the proposals and the GT box so that they can be
considered as positive during training of the classification head
box_bg_iou_thresh (float): maximum IoU between the proposals and the GT box so that they can be
considered as negative during training of the classification head
box_batch_size_per_image (int): number of proposals that are sampled during training of the
classification head
box_positive_fraction (float): proportion of positive proposals in a mini-batch during training
of the classification head
bbox_reg_weights (Tuple[float, float, float, float]): weights for the encoding/decoding of the
bounding boxes
"""
def __init__(self, backbone, num_classes=None, # num_classes是需要加上背景的
# transform parameter
min_size=800, max_size=1333, # 预处理resize时限制的最小尺寸与最大尺寸
image_mean=None, image_std=None, # 预处理normalize时使用的均值和方差
# RPN parameters
rpn_anchor_generator=None, rpn_head=None,
rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000, # rpn中在nms处理前保留的proposal数(根据score)
rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000, # rpn中在nms处理后保留的proposal数
rpn_nms_thresh=0.7, # rpn中进行nms处理时使用的iou阈值
rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3, # rpn计算损失时，采集正负样本设置的阈值
rpn_batch_size_per_image=256, rpn_positive_fraction=0.5, # rpn计算损失时采样的样本数，以及正样本占总样本的比例
rpn_score_thresh=0.0,
# Box parameters
box_roi_pool=None, box_head=None, box_predictor=None,
# 移除低目标概率 fast rcnn中进行nms处理的阈值 对预测结果根据score排序取前100个目标
box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5, # fast rcnn计算误差时，采集正负样本设置的阈值
box_batch_size_per_image=512, box_positive_fraction=0.25, # fast rcnn计算误差时采样的样本数，以及正样本占所有样本的比例
bbox_reg_weights=None):
if not hasattr(backbone, "out_channels"):
raise ValueError(
"backbone should contain an attribute out_channels"
"specifying the number of output channels (assumed to be the"
"same for all the levels"
)
assert isinstance(rpn_anchor_generator, (AnchorsGenerator, type(None))) # 传入None也可以，传入None后面会生成一个
assert isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None)))
if num_classes is not None:
if box_predictor is not None: # 如果box_predictor不为None，报错
raise ValueError("num_classes should be None when box_predictor "
"is specified")
else:
if box_predictor is None:
raise ValueError("num_classes should not be None when box_predictor "
"is not specified")
# 预测特征层的channels
out_channels = backbone.out_channels
# 若anchor生成器为空，则自动生成针对resnet50_fpn的anchor生成器
"""
>>> ((0.5, 1.0, 2.0),) * 5
((0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0), (0.5, 1.0, 2.0))
"""
if rpn_anchor_generator is None:
anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
rpn_anchor_generator = AnchorsGenerator(
anchor_sizes, aspect_ratios
)
# 生成RPN通过滑动窗口预测网络部分
if rpn_head is None:
rpn_head = RPNHead(
out_channels, rpn_anchor_generator.num_anchors_per_location()[0]
)
# 默认rpn_pre_nms_top_n_train = 2000, rpn_pre_nms_top_n_test = 1000,
# 默认rpn_post_nms_top_n_train = 2000, rpn_post_nms_top_n_test = 1000,
rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
# 定义整个RPN框架
# rpn_batch_size_per_image: RPN算计损失时采用正负样本的总个数
# rpn_positive_fraction：正样本占所有用于计算损失所有样本的比例
# rpn_pre_nms_top_n： 在进行NMS处理之前，针对每一个预测特征层所保留的目标个数
# rpn_post_nms_top_n： 在进行NMS处理之后，针对每一个预测特征层所剩余的目标个数（即RPN输出候选框的数目）
# rpn_nms_thresh：NMS处理时所指定的阈值
rpn = RegionProposalNetwork(
rpn_anchor_generator, rpn_head,
rpn_fg_iou_thresh, rpn_bg_iou_thresh,
rpn_batch_size_per_image, rpn_positive_fraction,
rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_nms_thresh,
score_thresh=rpn_score_thresh)
# Multi-scale RoIAlign pooling
if box_roi_pool is None:
box_roi_pool = MultiScaleRoIAlign( # 这里就是ROI Pooling
featmap_names=['0', '1', '2', '3'], # 在哪些特征层进行roi pooling
output_size=[7, 7], # 这里给出了ROI Pooling后输出的shape
sampling_ratio=2)
# fast RCNN中roi pooling后的展平处理两个全连接层部分
if box_head is None:
resolution = box_roi_pool.output_size[0] # 默认等于7
representation_size = 1024
box_head = TwoMLPHead(
out_channels * resolution ** 2, # flatten层的输出，即tensor展平后的元素个数，也是FC_1的输入
representation_size # FC_1的输出，也是FC_2的输入和输出
) # flatten -> in_channels -> representation_size -> representation_size
# 在box_head的输出上预测部分
if box_predictor is None:
representation_size = 1024 # Two MLP Head的输出，为固定值1024
box_predictor = FastRCNNPredictor(
representation_size, # Two MLP Head的输出，为固定值1024
num_classes) # num_classes为NC（加了背景的，VOC为20+1=21）
# 将roi pooling, box_head以及box_predictor结合在一起
"""
1. box_roi_pool: ROI Pooling
2. box_head: Two MLP Head
3. box_predictor: Fast R-CNN Predictor
function:
将Two MLP Head的结果并行通过两个FC
out:
1. cls_logits：每一个proposal的类别分数 [1024, 21](bs=2)
2. box_pred：每一个proposal对应的目标框的回归参数 [1024, 21*4=84](bs=2)
4. box_fg_iou_thresh（0.5）
5. box_bg_iou_thresh（0.5）
4和5：在前面匹配正负样本
如果proposal与GT的IoU>阈值 -> 正样本
如果proposal与GT的IoU<阈值 -> 负样本
6. box_batch_size_per_image（512）：每张图片会选取box_batch_size_per_image个proposal用于计算fast r-cnn的损失
7. box_positive_fraction（0.25）：指定样本中（即512个样本），正样本所占的比例
在训练过程中，并不是直接使用RPN生成的2000个proposal，而是从中采样、选取512个proposal
8. bbox_reg_weights：超参数
9. box_score_thresh, box_nms_thresh, box_detections_per_img：对最终预测的结果进行post-processing（后处理）的时候使用
到的一些阈值
"""
roi_heads = RoIHeads(
# box
box_roi_pool, box_head, box_predictor,
box_fg_iou_thresh, box_bg_iou_thresh, # 0.5 0.5
box_batch_size_per_image, box_positive_fraction, # 512 0.25
bbox_reg_weights,
box_score_thresh, box_nms_thresh, box_detections_per_img) # 0.05 0.5 100
if image_mean is None:
image_mean = [0.485, 0.456, 0.406]
if image_std is None:
image_std = [0.229, 0.224, 0.225]
# 对数据进行标准化，缩放，打包成batch等处理部分
transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std)
super(FasterRCNN, self).__init__(backbone, rpn, roi_heads, transform)
5. 换backbone
5.1 不带FPN
import os
import datetime
import torch
import transforms
from network_files import FasterRCNN, AnchorsGenerator
from my_dataset import VOCDataSet
from train_utils import GroupedBatchSampler, create_aspect_ratio_groups
from train_utils import train_eval_utils as utils
def create_model(num_classes):
import torchvision
from torchvision.models.feature_extraction import create_feature_extractor
"""
torchvision中的模型一般是针对imagenet进行训练的
不建议进行32倍下采样，因为下采样倍数大了，小目标检测效果就很差，所以建议使用到下采样16倍的特征层就可以了
"""
# vgg16
backbone = torchvision.models.vgg16_bn(pretrained=True)
# print(backbone)
"""
使用torchvision提供的create_feature_extractor提取指定中间层的输出
原理：通过create_feature_extractor重构backbone，通过重构的backbone就能获取中间层的输出
create_feature_extractor:
创建一个新的图形模块，该模块将给定模型中的中间节点作为字典返回，用户指定的键作为字符串，请求的输出作为值。
这是通过 FX 重写模型的计算图以返回所需节点作为输出来实现的。
删除所有未使用的节点及其相应的参数。
-------------
这个方法会根据我们传入的return_nodes参数找到对应输出的节点，再反向寻找该节点利用到之前哪些节点。之后会将没有使用到的节点全部删除。
return_nodes是一个字典：
"features.42": 对应节点的输出层
Q: 如何获取"features.42"？
A: 简单的方法，直接print(backbone)，然后再找到我们想要终止层的名称
"0": 重构的backbone的返回值，这个"0"是重构的backbone的key（不是value）
"""
backbone = create_feature_extractor(backbone, return_nodes={"features.42": "0"})
"""
Q: 怎么知道这个512的
A: 让tensor走一遍就知道了
out = backbone(torch.rand(1, 3, 224, 224))
print(out["0"].shape)
"""
backbone.out_channels = 512
# resnet50 backbone
# backbone = torchvision.models.resnet50(pretrained=True)
# # print(backbone)
# backbone = create_feature_extractor(backbone, return_nodes={"layer3": "0"})
# # out = backbone(torch.rand(1, 3, 224, 224))
# # print(out["0"].shape)
# backbone.out_channels = 1024
# EfficientNetB0
# backbone = torchvision.models.efficientnet_b0(pretrained=True)
# # print(backbone)
# backbone = create_feature_extractor(backbone, return_nodes={"features.5": "0"})
# # out = backbone(torch.rand(1, 3, 224, 224))
# # print(out["0"].shape)
# backbone.out_channels = 112
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
"""
ROI Pooling用的是mask rcnn中的ROIAlign
"""
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], # 在哪些特征层上进行RoIAlign pooling
output_size=[7, 7], # RoIAlign pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def main(args):
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
data_transform = {
"train": transforms.Compose([transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val": transforms.Compose([transforms.ToTensor()])
}
VOC_root = args.data_path
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_dataset = VOCDataSet(VOC_root, "2012", data_transform["train"], "train.txt")
train_sampler = None
# 是否按图片相似高宽比采样图片组成batch
# 使用的话能够减小训练时所需GPU显存，默认使用
if args.aspect_ratio_group_factor >= 0:
train_sampler = torch.utils.data.RandomSampler(train_dataset)
# 统计所有图像高宽比例在bins区间中的位置索引
group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor)
# 每个batch图片从同一高宽比例区间中取
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = args.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using %g dataloader workers' % nw)
if train_sampler:
# 如果按照图片高宽比采样图片，dataloader中需要使用batch_sampler
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_sampler=train_batch_sampler,
pin_memory=True,
num_workers=nw,
collate_fn=train_dataset.collate_fn)
else:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_dataset.collate_fn)
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_dataset = VOCDataSet(VOC_root, "2012", data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=val_dataset.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=args.num_classes + 1)
# print(model)
model.to(device)
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scaler = torch.cuda.amp.GradScaler() if args.amp else None
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.33)
# 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练
if args.resume != "":
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.amp and "scaler" in checkpoint:
scaler.load_state_dict(checkpoint["scaler"])
print("the training process from epoch{}...".format(args.start_epoch))
train_loss = []
learning_rate = []
val_map = []
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch, printing every 10 iterations
mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
device=device, epoch=epoch,
print_freq=50, warmup=True,
scaler=scaler)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [f"{i:.4f}" for i in coco_info + [mean_loss.item()]] + [f"{lr:.6f}"]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch}
if args.amp:
save_files["scaler"] = scaler.state_dict()
torch.save(save_files, "./save_weights/resNetFpn-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description=__doc__)
# 训练设备类型
parser.add_argument('--device', default='cuda:0', help='device')
# 训练数据集的根目录(VOCdevkit)
parser.add_argument('--data-path', default='./', help='dataset')
# 检测目标类别数(不包含背景)
parser.add_argument('--num-classes', default=20, type=int, help='num_classes')
# 文件保存地址
parser.add_argument('--output-dir', default='./save_weights', help='path where to save')
# 若需要接着上次训练，则指定上次训练保存权重文件地址
parser.add_argument('--resume', default='', type=str, help='resume from checkpoint')
# 指定接着从哪个epoch数开始训练
parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
# 训练的总epoch数
parser.add_argument('--epochs', default=15, type=int, metavar='N',
help='number of total epochs to run')
# 学习率
parser.add_argument('--lr', default=0.005, type=float,
help='initial learning rate, 0.02 is the default value for training '
'on 8 gpus and 2 images_per_gpu')
# SGD的momentum参数
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
# SGD的weight_decay参数
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay')
# 训练的batch size
parser.add_argument('--batch_size', default=4, type=int, metavar='N',
help='batch size when training.')
parser.add_argument('--aspect-ratio-group-factor', default=3, type=int)
# 是否使用混合精度训练(需要GPU支持混合精度)
parser.add_argument("--amp", default=False, help="Use torch.cuda.amp for mixed precision training")
args = parser.parse_args()
print(args)
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
main(args)
5.2 带FPN
import os
import datetime
import torch
import transforms
from network_files import FasterRCNN, AnchorsGenerator
from my_dataset import VOCDataSet
from train_utils import GroupedBatchSampler, create_aspect_ratio_groups
from train_utils import train_eval_utils as utils
from backbone import BackboneWithFPN, LastLevelMaxPool
def create_model(num_classes):
import torchvision
from torchvision.models.feature_extraction import create_feature_extractor
# --- mobilenet_v3_large fpn backbone --- #
backbone = torchvision.models.mobilenet_v3_large(pretrained=True)
# print(backbone)
"""
指定backbone中需要使用FPN的层
key: 对应的层
value: 返回的key
"""
return_layers = {"features.6": "0", # stride 8
"features.12": "1", # stride 16
"features.16": "2"} # stride 32
# 提供给fpn的每个特征层channel
in_channels_list = [40, 112, 960]
new_backbone = create_feature_extractor(backbone, return_layers)
# img = torch.randn(1, 3, 224, 224)
# outputs = new_backbone(img)
# [print(f"{k} shape: {v.shape}") for k, v in outputs.items()]
# --- efficientnet_b0 fpn backbone --- #
# backbone = torchvision.models.efficientnet_b0(pretrained=True)
# # print(backbone)
# return_layers = {"features.3": "0", # stride 8
# "features.4": "1", # stride 16
# "features.8": "2"} # stride 32
# # 提供给fpn的每个特征层channel
# in_channels_list = [40, 80, 1280]
# new_backbone = create_feature_extractor(backbone, return_layers)
# # img = torch.randn(1, 3, 224, 224)
# # outputs = new_backbone(img)
# # [print(f"{k} shape: {v.shape}") for k, v in outputs.items()]
backbone_with_fpn = BackboneWithFPN(new_backbone,
return_layers=return_layers,
in_channels_list=in_channels_list,
out_channels=256,
extra_blocks=LastLevelMaxPool(),
re_getter=False)
anchor_sizes = ((64,), (128,), (256,), (512,))
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
anchor_generator = AnchorsGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0', '1', '2'], # 在哪些特征层上进行RoIAlign pooling
output_size=[7, 7], # RoIAlign pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone_with_fpn,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def main(args):
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
data_transform = {
"train": transforms.Compose([transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val": transforms.Compose([transforms.ToTensor()])
}
VOC_root = args.data_path
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_dataset = VOCDataSet(VOC_root, "2012", data_transform["train"], "train.txt")
train_sampler = None
# 是否按图片相似高宽比采样图片组成batch
# 使用的话能够减小训练时所需GPU显存，默认使用
if args.aspect_ratio_group_factor >= 0:
train_sampler = torch.utils.data.RandomSampler(train_dataset)
# 统计所有图像高宽比例在bins区间中的位置索引
group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor)
# 每个batch图片从同一高宽比例区间中取
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = args.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using %g dataloader workers' % nw)
if train_sampler:
# 如果按照图片高宽比采样图片，dataloader中需要使用batch_sampler
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_sampler=train_batch_sampler,
pin_memory=True,
num_workers=nw,
collate_fn=train_dataset.collate_fn)
else:
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_dataset.collate_fn)
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_dataset = VOCDataSet(VOC_root, "2012", data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=val_dataset.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=args.num_classes + 1)
# print(model)
model.to(device)
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scaler = torch.cuda.amp.GradScaler() if args.amp else None
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.33)
# 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练
if args.resume != "":
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.amp and "scaler" in checkpoint:
scaler.load_state_dict(checkpoint["scaler"])
print("the training process from epoch{}...".format(args.start_epoch))
train_loss = []
learning_rate = []
val_map = []
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch, printing every 10 iterations
mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
device=device, epoch=epoch,
print_freq=50, warmup=True,
scaler=scaler)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [f"{i:.4f}" for i in coco_info + [mean_loss.item()]] + [f"{lr:.6f}"]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch}
if args.amp:
save_files["scaler"] = scaler.state_dict()
torch.save(save_files, "./save_weights/resNetFpn-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description=__doc__)
# 训练设备类型
parser.add_argument('--device', default='cuda:0', help='device')
# 训练数据集的根目录(VOCdevkit)
parser.add_argument('--data-path', default='./', help='dataset')
# 检测目标类别数(不包含背景)
parser.add_argument('--num-classes', default=20, type=int, help='num_classes')
# 文件保存地址
parser.add_argument('--output-dir', default='./save_weights', help='path where to save')
# 若需要接着上次训练，则指定上次训练保存权重文件地址
parser.add_argument('--resume', default='', type=str, help='resume from checkpoint')
# 指定接着从哪个epoch数开始训练
parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
# 训练的总epoch数
parser.add_argument('--epochs', default=15, type=int, metavar='N',
help='number of total epochs to run')
# 学习率
parser.add_argument('--lr', default=0.005, type=float,
help='initial learning rate, 0.02 is the default value for training '
'on 8 gpus and 2 images_per_gpu')
# SGD的momentum参数
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
# SGD的weight_decay参数
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay')
# 训练的batch size
parser.add_argument('--batch_size', default=4, type=int, metavar='N',
help='batch size when training.')
parser.add_argument('--aspect-ratio-group-factor', default=3, type=int)
# 是否使用混合精度训练(需要GPU支持混合精度)
parser.add_argument("--amp", default=False, help="Use torch.cuda.amp for mixed precision training")
args = parser.parse_args()
print(args)
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
main(args)
参考
- https://www.bilibili.com/video/BV1of4y1m7nj?spm_id_from=333.999.0.0
扫一扫
1782
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
