目录
摘要
Detection Transformer是一种基于Transformer架构的目标检测模型,它革新了传统目标检测的方法。传统目标检测模型通常依赖于手工设计的特征提取器和复杂的后处理步骤,而DETR则采用端到端的方式,直接预测目标物体的类别和边界框。DETR简化了目标检测流程,提高了检测精度,并实现了对多尺度目标的鲁棒检测。本篇博客将详细讲解DETR目标检测模型,以及附上PyTorch运行结果。
Abstract
Detection Transformer is an object detection model based on Transformer architecture, which revolutionizes traditional object detection methods. Traditional object detection models typically rely on manually designed feature extractors and complex post-processing steps, while DETR adopts an end-to-end approach to directly predict the category and bounding box of the target object. DETR simplifies the object detection process, improves detection accuracy, and achieves robust detection of multi-scale objects. This blog will provide a detailed explanation of the DETR object detection model and attach the PyTorch running results.
DETR
Object Detection with Transformers是端到端的目标检测框架,该网络模型没有使用非极大值抑制法去掉冗余框,也没有采用生成anchor的方式预测框,提高了目标检测的速度。DETR采用了新的目标函数,以二分图匹配的方式输出独一无二的预测框,我们来具体看看DETR是如何工作的吧!
1、网络结构
DETR的骨干网络采用的是ResNet-50进行图像特征提取,如上图所示CNN部分。将ResNet提取特征之后的特征图像拉直。然后,送入Transformer之中提取全局特征,全局特征的学习更有利于大目标的识别,能够很好的预测框。set of box predictions处会产生100个固定数量的预测框,这是作者人为设定的Query数(下文第2点会详细说到)。最后,将预测的100个框进行集合预测,确定出最终的预测框个数,再与图像标记的ground truth框进行误差计算,得以训练DETR。没有匹配的预测框则标记为背景类。
2、如何生成预测框
Transformer网络结构图,如下图所示:
我们知道在Transformer中encoder负责学习全局特征,这样有利于decoder更好的去进行解码预测。在decoder解码中,也需要输入,即上图中的Outputs处。Transformer是通过decoder处的输入的Query去和encoder输出的Key进行掩码自注意力得出预测结果。
我的理解是这样的,可能画图更好解释:
这里的Q为什么会选择100呢?因为作者采用COCO数据集进行训练,单张图片中最多出现几十种目标,所以100的预测框数量对于该数据集目标的预测是绰绰有余的。
但是,这样不就有多余的框出现了吗?如何做到端到端呢?多的冗余框不又和生成anchor的形式一样了,又需要去除多余的框。下文第3点会详细解释。
3、如何实现端到端
作者运用了一种新的目标函数,实现端到端的检测。上图中的 N 即为第2点中提到的Query,作者采用 N=100 ,是远大于图像目标的数量的。所以100个预测框,作者采用二分图匹配的形式与ground truth框进行一对一的匹配,就无需非极大值抑制处理。
具体来说,假设a、b、c点到达X、Y、Z点分别有着不同的代价,而它们分别到达每一点的代价图称为cost matrix。在scipy中的linear-sum-assignment函数能够计算出最优化匹配,使得abc到达XYZ的总价值最小。
在DETR中,我们可以理解为a、b、c代表着100个预测框,而X、Y、Z代表ground truth框。遍历所有预测框和ground truth框计算cost,得到最终的cost matrix。cost计算公式如下所示:
然后,利用scipy中的linear-sum-assignment函数计算出cost matrix的最优化匹配。这样九实现了预测框和真实框的一对一匹配,没有出现冗余的框。
最后,在将预测框和真实框进行类别预测和框预测的损失计算,即可反向传播优化模型。损失函数公式如下所示:
4、参数传递过程
DETR输入 800x1066x3 的图像,通过backbone提取图像特征,CNN模块下采样32倍,得到 25x34x2048 的特征图像,在与位置编码相加之前,会进行降维操作至 25x34x256 。位置编码维度和backbone输出相同,都是 25x34x256 ,因为便于相加操作。
在输入进encoder之前会将 25x34x256 的特征图像进行拉直变为 850x256 ,即850=25x34,共850个输入,每个输入256维。经过 6 个encoder之后的输出仍为 850x256 ,再将其传入decoder。
上图decoder中的object queries就是上文提到的Q=100,即向encoder的输出中解码出100个预测框。经过 6 个decoder之后输出 100x256 的特征向量,将其传入预测头prediction heads。
prediction heads会有很多个平行的全连接网络FFN,若需要预测类别(class box),则FFN是 1x91 的卷积层;若进行框回归(no object),则FFN是 1x4 的卷积层。
最后,在100个预测框都有了对应的类别预测和框预测,就需要和右图中的真实框进行最优匹配,得以去除冗余框;再通过反向传播,得以更新参数。
5、代码
作者提供了一个简短的代码,以了解DETR的整个模型结构,其代码十分简洁,如下图所示:
骨干网络采用ResNet-50,该骨干网络在ImageNet上预训练,DETR模型训练PyTorch代码如下:
#-------------------------------------#
# 对数据集进行训练
#-------------------------------------#
import datetime
import os
from functools import partial
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch import nn
from torch.utils.data import DataLoader
from nets.detr import DETR
from nets.detr_training import (build_loss, get_lr_scheduler, set_optimizer_lr,
weights_init)
from utils.callbacks import EvalCallback, LossHistory
from utils.dataloader import DetrDataset, detr_dataset_collate
from utils.utils import (get_classes, seed_everything, show_config,
worker_init_fn)
from utils.utils_fit import fit_one_epoch
if __name__ == "__main__":
#---------------------------------#
# Cuda 是否使用Cuda
# 没有GPU可以设置成False
#---------------------------------#
Cuda = True
#----------------------------------------------#
# Seed 用于固定随机种子
# 使得每次独立训练都可以获得一样的结果
#----------------------------------------------#
seed = 11
#---------------------------------------------------------------------#
# distributed 用于指定是否使用单机多卡分布式运行
# 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
# Windows系统下默认使用DP模式调用所有显卡,不支持DDP。
# DP模式:
# 设置 distributed = False
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
# DDP模式:
# 设置 distributed = True
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
#---------------------------------------------------------------------#
distributed = False
#---------------------------------------------------------------------#
# fp16 是否使用混合精度训练
# 可减少约一半的显存、需要pytorch1.7.1以上
#---------------------------------------------------------------------#
fp16 = False
#---------------------------------------------------------------------#
# classes_path 指向model_data下的txt,与自己训练的数据集相关
# 训练前一定要修改classes_path,使其对应自己的数据集
#---------------------------------------------------------------------#
classes_path = 'model_data/voc_classes.txt'
#----------------------------------------------------------------------------------------------------------------------------#
# 权值文件的下载请看README,可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的,因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分,用于进行特征提取。
# 预训练权重对于99%的情况都必须要用,不用的话主干部分的权值太过随机,特征提取效果不明显,网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作,可以将model_path设置成logs文件夹下的权值文件,将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数,来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重,因此是在train.py进行加载的,下面的pretrain不影响此处的权值加载。
# 如果想要让模型从主干的预训练权值开始训练,则设置model_path = '',下面的pretrain = True,此时仅加载主干。
# 如果想要让模型从0开始训练,则设置model_path = '',下面的pretrain = Fasle,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
#
# 一般来讲,网络从0开始的训练效果会很差,因为权值太过随机,特征提取效果不明显,因此非常、非常、非常不建议大家从0开始训练!
# 如果一定要从0开始,可以了解imagenet数据集,首先训练分类模型,获得网络的主干部分权值,分类模型的 主干部分 和该模型通用,基于此进行训练。
#----------------------------------------------------------------------------------------------------------------------------#
model_path = 'model_data/detr_resnet50_weights_coco.pth'
#------------------------------------------------------#
# input_shape 输入的shape大小
#------------------------------------------------------#
input_shape = [800, 800]
#---------------------------------------------#
# resnet50
# resnet101
#---------------------------------------------#
backbone = "resnet50"
#----------------------------------------------------------------------------------------------------------------------------#
# pretrained 是否使用主干网络的预训练权重,此处使用的是主干的权重,因此是在模型构建的时候进行加载的。
# 如果设置了model_path,则主干的权值无需加载,pretrained的值无意义。
# 如果不设置model_path,pretrained = True,此时仅加载主干开始训练。
# 如果不设置model_path,pretrained = False,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
pretrained = False
#----------------------------------------------------------------------------------------------------------------------------#
# 训练分为两个阶段,分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
# 冻结训练需要的显存较小,显卡非常差的情况下,可设置Freeze_Epoch等于UnFreeze_Epoch,此时仅仅进行冻结训练。
#
# 在此提供若干参数设置建议,各位训练者根据自己的需求进行灵活调整:
# (一)从整个模型的预训练权重开始训练:
# AdamW:
# Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 100,Freeze_Train = True,optimizer_type = 'adamw',Init_lr = 1e-4,weight_decay = 1e-4。(冻结)
# Init_Epoch = 0,UnFreeze_Epoch = 100,Freeze_Train = False,optimizer_type = 'adamw',Init_lr = 1e-4,weight_decay = 1e-4。(不冻结)
# 其中:UnFr