如何用DETR官方代码训练自己的数据集
DETR作为Transformer在目标检测中的开山之作，影响力自然不必多说，并且官方已经开源了代码，因此最好的学习方式便是利用论文和代码进行实践，论文地址和Github代码链接附上：
论文地址：https://arxiv.org/abs/2005.12872
Github地址：https://github.com/facebookresearch/detr
关于模型原理网上已经有各个大佬的讲解了，这里放上我觉得讲的最清晰的： 深度学习之目标检测（十一）–DETR详解
提示：仅供参考，具体命令行或代码按需执行
1. 代码结构
首先我们需要先Clone下来代码，代码的结构并不复杂，只有几个文件夹：
- datasets：源码是使用COCO2017作为基础数据集的，因此这里面包含有COCO的Dataset类以及用于数据增强的transforms.py文件。
- models：该文件夹下存储了DETR模型的基础结构，包括Transformer模块和主干网络等等，我们都可以不用去修改它。
- utils：该文件夹下存储的都是一些杂七杂八的工具类，也不用去修改它。
重点我们只需要关注以下几个文件：
- main.py：这是整个程序的入口文件，我们需要运行这个文件来设置各种超参数以及进入训练，验证的入口，也就是engine.py。
- engine.py：这是“引擎”文件，说白了就是从main文件会调用里面的函数去进行训练和验证。
- datasets/transforms.py：这是数据增强的部分，有很长时间我一直没有训练出好的结果，在看了这里的代码才知道是自己的数据格式出了问题，因此后面也会讲一下。
2. 构建自己的数据集
COCO数据集可以直接采用官方的实现方法，在Github里也有对应的Tutorial，这里我们讲一下如何使用自己的数据集进行训练。
2.1 Roboflow标注
首先是数据标注部分，我采用的是Roboflow平台进行标注，标注完成后选择导出格式为COCO数据集格式，选择“Download zip to computer”，然后在自己电脑上进行压缩，我们可以得到这样格式的文件，分了train，valid，test文件夹，每个文件夹下存储了对应的图片还有一个json文件，json文件里包含了所有标签的信息。
2.2 自定义Dataset类（PyTorch）
接下来我们要把数据喂入神经网络，官方代码是用PyTorch实现的，因此我们要自己实现自己数据集的Dataset类，创建一个新的py文件，放在datasets文件夹下面。这里放上我的代码，很基础的继承torch.utils.data.Dataset类的写法，也比较好理解，重要的部分我已经注释出来了。
class SmartFarm(Dataset):
def __init__(self, path, image_set='train', transform=None, n_cls=1):
super().__init__()
assert image_set in ['train', 'val', 'test'], 'set must be one of train, valid, test'
"""
官方原代码里是用的val代表valid，但是roboflow生成的文件夹叫valid，可以改直接改文件名，我这样写是为了方便复用
"""
if image_set == 'val':
self.path = os.path.join(path, 'valid')
else:
self.path = os.path.join(path, image_set)
self.json_path = os.path.join(self.path, '_annotations.coco.json')
self.transform = transform
self.n_cls = n_cls + 1 # 背景也算一个单独的类别
self.ids = list(self.image2id().keys())
def __len__(self):
with open(self.json_path, 'r') as f:
data = json.load(f)
image_mapping = self.image2id()
return len(image_mapping)
def __getitem__(self, index):
idx = self.ids[index]
img, target = self.read_json(self.json_path, idx)
if self.transform is not None:
img, target = self.transform(img, target)
return img, target
"""
对应每一个类和类id
Return: {1:'flower', 2:'leaf'}
"""
def label2id(self):
with open(self.json_path, 'r') as f:
data = json.load(f)
mapping = {}
for item in data:
mapping[item['id']] = item['name']
return mapping
"""
对应每一张图片的名字和对应id
Return: {1:'7711724741884_-pic_jpg.rf.b4bfcd10c478314fbeb7f136a5c33c46.jpg'}
"""
def image2id(self):
with open(self.json_path, 'r') as f:
data = json.load(f)
mapping = {}
for item in data['images']:
mapping[item['id']] = item['file_name']
return mapping
"""
读取Json文件，这一部分可以参考一下官方源码的datasets/coco.py中的ConvertCocoPolysToMask类
最终返回的target的格式应该是一个字典，像这样：
{"boxes":[], "labels":[], "area":[], "iscrowd":[], "orig_size":[], "size":[]}
因为roboflow生成的json文件只有一个，然而__getiten__中一次只返回一张图和标签，因此我们需要把各个部分都拆解开来
"""
def read_json(self, path, idx):
with open(path, 'r') as f:
data = json.load(f)
image_mapping = self.image2id()
target = {}
bboxes, classes, image_files, area, iscrowd = [], [], [], [], []
for ann in data['annotations']:
if ann['image_id'] == idx:
image_files.append(image_mapping[ann['image_id']])
bboxes.append(ann['bbox'])
classes.append(ann['category_id'])
area.append(ann['area'])
iscrowd.append(ann['iscrowd'])
bboxes = torch.as_tensor(bboxes, dtype=torch.float32)
# =====================================
# 这里需要注意，roboflow生成的是（x, y, w, h）格式的box，xy为左上角的坐标
# 但是官方源码的输入需要的是左上右下两个点的坐标
tmp = bboxes.reshape(-1, 2, 2)
left_corner = tmp[:, 0]
right_corner = tmp[:, 0] + tmp[:, 1]
bboxes = torch.cat((left_corner, right_corner), dim=-1)
# =====================================
classes = torch.as_tensor(classes, dtype=torch.int64)
area = torch.as_tensor(area, dtype=torch.float32)
iscrowd = torch.as_tensor(iscrowd)
target["boxes"] = bboxes
target["labels"] = classes
target["image_id"] = torch.tensor(idx)
target["area"] = area
target["iscrowd"] = iscrowd
image = Image.open(os.path.join(self.path, image_files[0]))
w, h = image.size
target["orig_size"] = torch.as_tensor([int(h), int(w)])
target["size"] = torch.as_tensor([int(h), int(w)])
return image, target
"""
调用数据集的入口函数，也是参考的原代码中datasets/coco.py的写法
"""
def build(image_set, args):
path = Path(args.other_dataset_path)
assert path.exists(), f'provided SF path {path} does not exist'
dataset = SmartFarm(path, image_set, transform=make_coco_transforms(image_set))
return dataset
写好Dataset类之后，要怎么去调用呢？来到datasets/__init__.py文件里，这个文件是用来选择数据集的，我们加上自己的数据集选择器，不用修改原有的代码，只用加一个if判断即可：
def build_dataset(image_set, args):
if args.dataset_file == 'coco':
return build_coco(image_set, args)
if args.dataset_file == 'coco_panoptic':
# to avoid making panopticapi required for coco
from .coco_panoptic import build as build_coco_panoptic
return build_coco_panoptic(image_set, args)
"""
我的数据集名字叫sf，当我传入这个参数的时候它会自动调用上面创建好的build入口函数创建SmartFarm的数据类：
"""
if args.dataset_file == 'sf':
from .sf import build
return build(image_set, args)
raise ValueError(f'dataset {args.dataset_file} not supported')
2.3 数据增强部分
如果你像我之前那样写数据的部分，那这一部分你可以直接跳过不用管，原代码中的数据增强就可以直接使用，但如果你想自己换更多的数据增强方式，这一部分有必要了解一下，因为我自己也踩了不少坑哈哈哈。
接下来我们了解一下原代码是如何实现数据增强的，重点关注datasets文件夹下的coco.py和transforms.py文件。coco.py中有一个make_coco_transforms函数，用于做数据增强，里面使用了简写T，一般来说第一反应都是T不就是代表的torchvision.transforms库嘛，那好，你被误导了！！
原作者自己重写了transforms库，在transforms.py文件里面，这个T是“import datasets.transforms as T”而非“import torchvision.transforms as T”。其实也不叫重写，只是添加了对bounding box的数据增强处理，因为torchvision只实现了对图像的增强，但对于图像分割的mask或者目标检测的bbox都没有对应的可以直接拿来用的函数（现在好像有了，但是为了能直接复用原代码，偷懒就不改了），具体怎么实现的你们可以直接去看transforms.py的源码，其实真的真的非常简单。
要注意的部分只有transforms.py下的Normalize类，这也是让我幡然醒悟的一段代码，从倒数第四行里的“box_xyxy_to_cxcywh”我们可以看出，原来它在数据增强之前是(x, y, x, y)的数据格式而不是roboflow自带的(x, y, w, h)啊！！，并且真正喂入网络的格式是(center_x, center_y, w, h)，有点绕，但是必须弄清楚！！不然就会像我之前那样loss飙到2w多。。。
并且最后模型预测的输出结果也会是(center_x, center_y, w, h)，如果把图像和目标检测框画出来，得以相反的方式把它还原。
class Normalize(object):
def __init__(self, mean, std):
self.mean = mean
self.std = std
def __call__(self, image, target=None):
image = F.normalize(image, mean=self.mean, std=self.std)
if target is None:
return image, None
target = target.copy()
h, w = image.shape[-2:]
if "boxes" in target:
boxes = target["boxes"]
boxes = box_xyxy_to_cxcywh(boxes)
boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
target["boxes"] = boxes
return image, target
3. 开始训练
3.1 模型参数设置
综上我们完成了数据的处理部分，接下来关注到main.py文件，开始进行模型的训练部分。同大多数源码一样，代码是在服务器上跑的因此采用了命令行启动的方式，具体来说就是它定义了很多参数，我们需要从命令行里发送给它，像这样：
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco --coco_panoptic_path /path/to/coco_panoptic --dataset_file coco_panoptic --output_dir /output/path/box_model
但由于我是在本地Pycharm上跑的，为了能方便调试，我把该有的参数都定义在了代码里面，比如下面几个比较重要的参数：
def get_args_parser():
# 设置训练超参数
parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
parser.add_argument('--lr', default=1e-4, type=float)
parser.add_argument('--lr_backbone', default=1e-5, type=float)
parser.add_argument('--batch_size', default=8, type=int)
parser.add_argument('--weight_decay', default=1e-4, type=float)
parser.add_argument('--epochs', default=300, type=int)
parser.add_argument('--lr_drop', default=200, type=int)
parser.add_argument('--clip_max_norm', default=0.1, type=float,
help='gradient clipping max norm')
...
# 数据集选择器，这个参数会在datasets/__init__.py中的build_dataset函数中调用，选择我们自己的自定义Dataset类
parser.add_argument('--dataset_file', default='sf')
# 这个是我自己新加的一个参数，原代码只提供了coco_path，这样写的话可以在不改变原代码的基础上选择我们自己的数据集路径
parser.add_argument('--other_dataset_path', default='/Users/shijunshen/Documents/Code/dataset/Smart_Farm_Detection.v1i.coco', type=str)
# 选择你的output文件夹，这里会存储checkpoint和log文件
parser.add_argument('--output_dir', default='./output',
help='path where to save, empty for no saving')
# 选择cpu或者cuda，因为我是MacOS，所以我使用的是苹果的硬件加速mps
parser.add_argument('--device', default='mps',
help='device to use for training / testing')
# 原代码中定义的是100，因为采用的coco数据集，这里的数值对应的是模型会预测出来多少个预测框，由于在我的数据集里大概一张图有十几个目标，因此我定的20。
parser.add_argument('--num_queries', default=20, type=int,
help="Number of query slots")
# 这个参数会调用之前训练好的checkpoint文件，你可以直接下载官网预训练好的文件
parser.add_argument('--resume', default='./output/checkpoint0599.pth', help='resume from checkpoint')
# ====================================================================
# 原代码只有eval一个参数，predict是我自己加上的，这两个参数在train的时候不用管它，在测试和模型预测的时候可以在后面直接加上default=True，不加True就没有用，这个和在命令行里加上--eval这个参数作用是一样的
parser.add_argument('--eval', action='store_true')
parser.add_argument('--predict', action='store_true')
# ====================================================================
# 多线程，我的cpu有10个核我就填了10，影响不大
parser.add_argument('--num_workers', default=10, type=int)
...
return parser
其它的代码也很常规，很容易看懂，还有一个部分就是engine.py，我们需要修改一段代码，在它的evaluate函数里面，把coco_evaluator置为None，不然它就会调用COCO的评估代码里去了。训练的时候每一轮都会调用这个函数去进行打印，这也是我觉得作者写的还不够好的一点，就是只能用他的原生代码去训练COCO数据集，没有给一个完善的接口让我们拿来即用，必须得自己来改改。
coco_evaluator = None
# coco_evaluator = CocoEvaluator(base_ds, iou_types)
同时我们还要找到models/detr.py下的build函数，把它的num_classes改为你自己的类别数量，这一段其实可以采用传参的模式传到模型中来，我也不知道为什么作者留在这里只能单独来设置。
num_classes = 1 if args.dataset_file != 'coco' else 91
if args.dataset_file == "coco_panoptic":
# for panoptic, we just add a num_classes that is large enough to hold
# max_obj_id + 1, but the exact value doesn't really matter
num_classes = 250
接下来只需要直接点击运行按钮就可以开始训练模型啦！记得–eval，和–predict不要设置为True喔。
3.2 加载预训练权重然后微调
这个部分可以跳过，当然我还是更建议大家采用这种做法，因为原代码中已经给出了训练好的权重啦。
原代码在COCO数据集上训练的权重我们也可以拿来做初始化，但是要注意你下载的权重必须和你的网络结构是一样的，比如你在main.py中设置了backbone为resnet50，你也要下载对应的权重文件，如果你本身就什么也没改的话，就可以直接下载它的第一个权重文件。
然后把resume设置为你下载好的权重文件路径，==同时也是最重要的一点==，由于我们修改了num_queries以及num_classes，模型的head部分已经不一样了，因此我们需要在main.py里加上几段代码：
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(
args.resume, map_location='cpu', check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
# ==============================================================
# 这一段是修改了的，去除多余的参数，并将load_state_dict设置为strict=False，这样它便会只加载模型结构相同部分的预训练参数
del checkpoint["model"]["class_embed.weight"]
del checkpoint["model"]["class_embed.bias"]
del checkpoint["model"]["query_embed.weight"]
model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
# ==============================================================
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
4. 可视化测试结果
训练之后我们肯定想把结果直观地显示出来，不然怎么去给老板交差🤡，源码没有实现这个，因此我添加了上面所说的predict参数，并且在main.py和engine.py里分别加入了这两段：
"""
main.py
"""
if args.eval:
test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
data_loader_val, base_ds, device, args.output_dir)
# if args.output_dir:
# utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
return
# 这一段是我加的
if args.predict:
predict(model, criterion, postprocessors,
data_loader_val, base_ds, device, args.output_dir)
"""
engine.py
这个只是我用来做测试用的，因此写得不是很完善，只可视化了第一张图像
但是其中对box的还原你们可以参考一下
"""
@torch.no_grad()
def predict(model, criterion, postprocessors, data_loader, base_ds, device, output_dir):
model.eval()
criterion.eval()
for samples, targets in data_loader:
samples = samples.to(device)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
outputs = model(samples)
for i, tensor in enumerate(samples.tensors):
img = tensor
break
# 反Normalization，这个mean和std是ImageNet数据集上做pretrain时候的常用参数，原代码也用了这两个值，可以在数据增强那部分看到
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
img = img.cpu().clone().detach().numpy()
img = np.transpose(img, (1, 2, 0))
img = img * np.array(std) + np.array(mean)
img = np.clip(img, 0, 1)
img = (img * 255).astype(np.uint8)
img_cv2 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # 将NumPy数组转换为OpenCV的Mat对象
bbox = outputs['pred_boxes'].cpu().detach().numpy()[0]
# bbox = targets[0]['boxes'].cpu().numpy()
img_height, img_width, _ = img_cv2.shape
[h, w] = targets[0]['size']
# 把（center_x, center_y, w, h）的格式转换为左上角和右下角两个点的坐标，用cv2.rectangle把它们画出来
for box in bbox:
print(box)
cx = int(box[0] * w)
cy = int(box[1] * h)
w_ = int(box[2] * w)
h_ = int(box[3] * h)
xmin = int(cx - w_ / 2)
ymin = int(cy - h_ / 2)
xmax = int(cx + w_ / 2)
ymax = int(cy + h_ / 2)
print(xmin, ymin, xmax, ymax)
cv2.rectangle(img_cv2, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
cv2.imshow('test', img_cv2)
cv2.waitKey()
在我的项目中，最终的结果是这样的，因为我在num_queries中设置了20，所以模型一共预测了20个bounding box，后面你们也可以做处理，比如只显示预测概率最高的前几个框，在模型预测结果的字典中不仅有‘pred_boxes’，也有‘pred_logits’，可以过滤出最高概率的几个框，不过这样已经可以看出，结果还是不错的啦。
5. 总结
总的来说，DETR训练的效率还是蛮高的，我只用了15张标注好的图片进行训练，并且下载了原代码中在COCO数据集上预训练好的模型然后在自己的数据集上进行训练，用Colab上的L4 GPU训练了十分钟左右，600个epoch便达到了这样的效果。
欢迎大家提问以及指正！