基于pytorch预训练模型进行目标检测

基于pytorch预训练模型的目标检测

本实验说明：

1. 用jupyter notebook实现
2. 代码基于pytorch2.0.1，torchvision0.16用的Faster RCNN FPN模型实现
3. 默认已经知晓基本的概念，如目标检测、微调等

写在前面

深度神经网络是一个灰箱，如果只是用这个模型，我们无需管它内部结构到底是什么样的，只要知道它的输入输出是什么就行了。

Faster RCNN FPN的输入输出

帮助文档

路径：pytorch官网-Docs帮助文档-torchvision0.16-Models and pre-trained weights模型和预训练权重-FASTER R-CNN-FASTERRCNN_RESNET50_FPN_V2

通过上述路径找到模型的帮助文档，可以知晓很多事情：

• 模型的原始论文
• 模型训练时的输入输出的格式和注意
• 模型预测时的输入输出的格式和注意
• Coding过程中各参数代表的含义
• 一个简单的示例

好吧！上述的很多内容需要到FASTERRCNN_RESNET50_FPN中找，因为这两个模型基本就是一个东西，torch官网也懒得再写一遍，直接链接过去看就行。

模型的输入应该是一个张量列表，每个张量的形状都是[C, H, W]，即 [Channel，Height，Width]。一个图像对应一个张量，数值范围0-1，对应FloatTensor格式。也就说图片在输入模型前需要进行归一化操作。

模型有两个模式，训练和预测，对应的输入输出都会有变化。

训练阶段的输入输出

输入

在训练过程中，模型期望输入张量和目标。

• 张量：这个上面已经提到了，就对应一张图片。
• 目标：用框和标签来描述目标，如图。
• • 框（FloatTensor格式[N, 4]）：用[x1, y1, x2, y2]表示1个实例，要求0 <= x1 < x2 <= W同时0 <= y1 < y2 <= H。
• • 标签（Int64Tensor格式[N]）：每个实例框对应的类标签。
    

假设我用Dataset已经写好了我所需要的数据集，现在把它打印出来看看。

train_set = CustomDataset('datasets/annotations/train.json')  # 生成训练数据集
train_set[0]  # 索引一个出来看看数据长什么样

(tensor([[[0.8902, 0.8902, 0.8902,  ..., 0.1765, 0.1843, 0.1882],
          [0.8902, 0.8902, 0.8902,  ..., 0.1804, 0.1882, 0.1922],
          [0.8902, 0.8902, 0.8902,  ..., 0.1686, 0.1765, 0.1804],
          ...,
          [0.7137, 0.7020, 0.6980,  ..., 0.6902, 0.6980, 0.7020],
          [0.7255, 0.7098, 0.7137,  ..., 0.6941, 0.7020, 0.7059],
          [0.7294, 0.7137, 0.7176,  ..., 0.6941, 0.7020, 0.7098]],
 
         [[1.0000, 1.0000, 1.0000,  ..., 0.2000, 0.2078, 0.2118],
          [1.0000, 1.0000, 1.0000,  ..., 0.2039, 0.2118, 0.2157],
          [1.0000, 1.0000, 1.0000,  ..., 0.1922, 0.2000, 0.2039],
          ...,
          [0.7725, 0.7608, 0.7686,  ..., 0.7686, 0.7765, 0.7804],
          [0.7804, 0.7686, 0.7765,  ..., 0.7725, 0.7804, 0.7843],
          [0.7804, 0.7725, 0.7804,  ..., 0.7725, 0.7804, 0.7882]],
 
         [[0.9882, 0.9882, 0.9882,  ..., 0.1529, 0.1608, 0.1647],
          [0.9882, 0.9882, 0.9882,  ..., 0.1490, 0.1569, 0.1608],
          [0.9882, 0.9882, 0.9882,  ..., 0.1373, 0.1451, 0.1490],
          ...,
          [0.7843, 0.7725, 0.7686,  ..., 0.7608, 0.7686, 0.7725],
          [0.7922, 0.7804, 0.7765,  ..., 0.7647, 0.7725, 0.7765],
          [0.8039, 0.7843, 0.7804,  ..., 0.7686, 0.7765, 0.7843]]]),
 {'boxes': tensor([[  0.0000, 275.5000, 116.1875, 394.5000],
          [222.5000, 275.2500, 307.0000, 328.5000],
          [211.5000, 279.7500, 268.2500, 304.0000],
          [356.5000, 272.0000, 397.0000, 305.2500],
          [384.7500, 276.2500, 438.0000, 315.2500],
          [470.0000, 262.5000, 503.0000, 287.0000],
          [500.5000, 259.0000, 527.0000, 279.5000],
          [478.7500, 242.2500, 505.2500, 270.2500],
          [522.5000, 254.1250, 539.0000, 271.5000],
          [534.0000, 269.7500, 580.5000, 309.0000],
          [564.0000, 273.0000, 723.0000, 409.5000],
          [636.0000, 234.7500, 679.0000, 285.0000],
          [617.0000, 237.5000, 646.0000, 276.0000],
          [585.5000, 251.6250, 605.0000, 273.5000],
          [546.0000, 239.3750, 560.0000, 257.5000]], dtype=torch.float16),
  'labels': tensor([16,  8, 16, 16, 16, 16, 16,  3,  3, 16, 16,  3,  3, 20,  7])})

可以看到tensor表示的就是一张图片，而一张图片中有多个实例，上述案例就包含了15个实例。接下来就可以用DataLoader把数据打包训练了。

train_dl = DataLoader(dataset=train_set, batch_size=1, shuffle=True)  # 这个东西就可以丢给模型了

输出

该模型在训练过程中返回一个tensor格式的字典，其中包含了RPN和R-CNN的分类和回归损失。

还是打印出来看看更加直观。

for epoch in range(num_epochs):
    model.train()  # 将模型调整到训练模式
    for images, targets in train_dl:  # 从DataLoader中加载tensor格式的图片和字典
        # 此处省略若干代码
        loss_dict = model(images, targets)  # 这里才是模型的输入输出，loss_dict就是模型的输出
        # 此处省略若干代码
    # 此时打印出来的是训练完一轮之后的loss
    print('loss_classifier = ', loss_dict['loss_classifier']) 
    print('loss_box_reg = ', loss_dict['loss_box_reg'])
    print('loss_objectness = ', loss_dict['loss_objectness'])
    print('loss_rpn_box_reg = ', loss_dict['loss_rpn_box_reg'])
    break

loss_classifier =  tensor(0.2090, device='cuda:0', grad_fn=<NllLossBackward0>)
loss_box_reg =  tensor(0.2016, device='cuda:0', grad_fn=<DivBackward0>)
loss_objectness =  tensor(0.0405, device='cuda:0', grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
loss_rpn_box_reg =  tensor(0.0077, device='cuda:0', grad_fn=<DivBackward0>)

输出总共是4个loss，也就是说网络已经帮你算好loss了，不需要再设计新的loss。我们再引入优化器就好了。

预测阶段的输入输出

输入

在预测过程中，模型只需要输入张量。只需要把预测的图片输入进去就可以啦。

输出

模型以List[Dict[Tensor]]的形式返回后处理的预测。Dict的字段如下，其中N为检测次数：

• 框（FloatTensor格式[N, 4]）：用[x1, y1, x2, y2]表示1个实例，要求0 <= x1 < x2 <= W同时0 <= y1 < y2 <= H。
• 标签（Int64Tensor格式[N]）：预测的类别。
• 分数（Tensor格式[N]）：预测的分数。

搞清楚了模型的输入输出是怎么回事，就再来了解一下如何修改出自己的模型。

预训练和修改模型

帮助文档

你以为有其他的帮助文档，错了，还是原来的帮助文档🤪。

我们看一看FASTERRCNN_RESNET50_FPN_V2的参数：

• weights（FasterRCNN_ResNet50_FPN_V2_Weights，可选）：要使用的预训练权重。点击FasterRCNN_ResNet50_FPN_V2_Weights查看更多细节和可选的预训练参数。默认情况下不使用预训练权重。
• progress（bool，可选）：如果为True，显示下载到stderr的进度条。默认值为True。可能是指权重下载到本地时的进度，附赠一下权重的保存路径：C:\Users\Users Name.cache\torch\hub\checkpoints。
• num_classes（int，可选）：模型输出类别的数量（包括背景）。如果类别有10种，那么需要填11种，因为要加上背景。
• weights_backbone（ResNet50_Weights，可选）：骨干网络的预训练权重。
• trainable_backbone_layers（int，可选）：从最终块开始的可训练（未冻结）层的数量。有效值在0到5之间，5意味着所有backbone层都是可训练的。如果传入了None(默认值)，则该值设置为3。

实际操作

from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights
from torchvision.models import ResNet50_Weights


model = fasterrcnn_resnet50_fpn_v2(
    weights=FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT,
#    num_classes=23,
    weights_backbone=ResNet50_Weights.DEFAULT,
    trainable_backbone_layers=5,
    box_score_thresh=0.80  # 筛选出分数高于0.8的框
)

以上就是用预训练权重初始化模型的过程，可以自己调着试试。比如使用FasterRCNN_ResNet50_FPN_V2_Weights作为权重，或者只用骨干网路权重，其他的全都用自己的数据训练微调。

这里有个问题，完整的预训练网络是用COCO训练的，一共分为91个类别，通过帮助文档我们能获得更多细节：

属性	参数
box_map (on COCO-val2017)	46.7
categories	background, person, bicycle, … (88 omitted)
min_size	height=1, width=1
num_params	43712278
recipe	link
GFLOPS	280.37
File size	167.1 MB

也就是说如果用了预训练权重就要求分类数量为91，但我可能只需要22(前景) + 1(背景) = 23类。我想使用完整的权重，又想修改类别，就需要如下操作：

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


# 获取模型的分类器的输入特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features

# 将模型的分类器替换为新的分类器，适应你的目标类别数
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=23)

还是打印出来看看直观一些：

model

FasterRCNN(
  (transform): GeneralizedRCNNTransform(...)
  (backbone): BackboneWithFPN(...)
  (rpn): RegionProposalNetwork(...)
  (roi_heads): RoIHeads(
    ...
    (box_predictor): FastRCNNPredictor(
      (cls_score): Linear(in_features=1024, out_features=23, bias=True)
      (bbox_pred): Linear(in_features=1024, out_features=92, bias=True)
    )
  )
)

由于整个模型打印出来太长了，就省略了亿点点。可以看到cls_score中的out_features已经被修改成了23。而不是原来的91了。bbox_pred中的out_features是92，正好是23的4倍，对应[x1, y1, x2, y2]。

其它关键

• 数据与整理：知道了模型需要的输入输出，就按照输入输出去获取和整理数据。
• 训练与调参：优化器、参数的设置，然后把整理好的数据喂给模型训练，调整参数获取更好的结果。

附件

未整理的代码，看看就行。链接。