【目标检测】RetinaNet/FCOS loss优化、multi scale训练、mosaic数据增强
文章目录
- RetinaNet/FCOS loss优化
- multi scale训练
- mosaic数据增强
所有代码已上传到本人github repository：https://github.com/zgcr/pytorch-ImageNet-CIFAR-COCO-VOC-training
如果觉得有用，请点个star哟！
下列代码均在pytorch1.4版本中测试过，确认正确无误。
RetinaNet/FCOS loss优化
在前面的RetinaNet/FCOS实现中，计算loss时采用的是先计算每张图的分类/回归loss再加起来平均的方式。这种方式认为每张图的权重是一样的。但是，目标检测模型实际学习的是每张图标注的框，而每张图标注的框的数量是不一样的，因此每张图的权重也不一样。所以，我将RetinaNet/FCOS的各个loss改为整个batch计算loss再求平均，平均的分母是这个batch的所有positive anchor/point的数量，这个数量是和标注框的数量线性正相关的。这样修改后，计算loss时的并行度得到进一步提高，计算而来的梯度也更加合理，模型训练一个epoch的时间减少了40%，且训练的模型性能更好。
各个模型的训练结果如下：
| Network | resize | batch | gpu-num | apex | syncbn | epoch | mAP-mAR-loss | training-time(hours) |
|---|---|---|---|---|---|---|---|---|
| ResNet50-RetinaNet | 667 | 20 | 2 RTX2080Ti | yes | no | 12 | 0.305,0.421,0.56 | 17.43 |
| ResNet101-RetinaNet | 667 | 16 | 2 RTX2080Ti | yes | no | 12 | 0.306,0.420,0.55 | 22.06 |
| ResNet50-RetinaNet | 1000 | 16 | 4 RTX2080Ti | yes | no | 12 | 0.332,0.458,0.57 | 26.25 |
| ResNet50-FCOS | 667 | 24 | 2 RTX2080Ti | yes | no | 12 | 0.318,0.452,1.09 | 14.17 |
| ResNet101-FCOS | 667 | 20 | 2 RTX2080Ti | yes | no | 12 | 0.342,0.475,1.07 | 19.20 |
| ResNet50-FCOS | 1000 | 20 | 4 RTX2080Ti | yes | no | 12 | 0.361,0.502,1.10 | 18.92 |
| ResNet50-FCOS | 1333 | 12 | 4 RTX2080Ti | yes | no | 24 | 0.381,0.534,1.03 | 37.73 |
这里训练时resize均采用的时yolov3的resize方法，按照flops大小对标，我的resize=667 flops等同于RetinaNet文章中的resize=400，resize=1000等同于RetinaNet文章中的resize=600，resize=1333RetinaNet文章中的resize=800。可以看到FCOS在resize=1333时的mAP已经到了0.381，和原文中集成所有改进之后的mAP0.386已经非常接近了。
multi scale训练
multi scale训练在目标检测中使用的非常广泛。这里采用yolov3/yolov5中的multi scale方法。首先确定一个最小到最大尺寸的范围，比如没用multi scale前resize=416，那么可以确定一个范围[0.5,1.5]，乘以416之后就是最小到最大尺寸的范围。然后，选择一个stride长度，比如yolov3/yolov5是32，找出最小到最大尺寸中所有能被stride长度整除的尺寸。最后，随机选择一个尺寸，将所有图片resize到该尺寸再填充成正方形图片就可以送入网络训练了。
multi scale的代码实现如下（注意用multi scale时transform里不再使用resize方法）：
class MultiScaleCollater():
def __init__(self,
resize=512,
multi_scale_range=[0.5, 1.5],
stride=32,
use_multi_scale=False):
self.resize = resize
self.multi_scale_range = multi_scale_range
self.stride = stride
self.use_multi_scale = use_multi_scale
def next(self, data):
if self.use_multi_scale:
min_resize = int(
((self.resize + self.stride) * self.multi_scale_range[0]) //
self.stride * self.stride)
max_resize = int(
((self.resize + self.stride) * self.multi_scale_range[1]) //
self.stride * self.stride)
final_resize = random.choice(
range(min_resize, max_resize, self.stride))
else:
final_resize = self.resize
imgs = [s['img'] for s in data]
annots = [s['annot'] for s in data]
scales = [s['scale'] for s in data]
padded_img = torch.zeros((len(imgs), final_resize, final_resize, 3))
for i, image in enumerate(imgs):
height, width, _ = image.shape
max_image_size = max(height, width)
resize_factor = final_resize / max_image_size
resize_height, resize_width = int(height * resize_factor), int(
width * resize_factor)
image = cv2.resize(image, (resize_width, resize_height))
padded_img[i, 0:resize_height,
0:resize_width] = torch.from_numpy(image)
annots[i][:, :4] *= resize_factor
scales[i] = scales[i] * resize_factor
max_num_annots = max(annot.shape[0] for annot in annots)
if max_num_annots > 0:
annot_padded = torch.ones((len(annots), max_num_annots, 5)) * (-1)
if max_num_annots > 0:
for idx, annot in enumerate(annots):
if annot.shape[0] > 0:
annot_padded[
idx, :annot.shape[0], :] = torch.from_numpy(annot)
else:
annot_padded = torch.ones((len(annots), 1, 5)) * (-1)
padded_img = padded_img.permute(0, 3, 1, 2)
return {'img': padded_img, 'annot': annot_padded, 'scale': scales}
mosaic数据增强
在前面的文章中，我按照自己的理解写了一版mosaic数据增强的代码。但是最近在阅读yolov3/yolov5的代码时发现官方的mosaic数据增强有些不一样。假设不使用mosaic数据增强时的resize=416，最终送入网络训练的图片尺寸是416x416，那么使用mosaic数据增强后最终送入网络训练的图片尺寸就是832x832(长宽都变为416的两倍)。首先从数据集随机取4张图，然后把它们都resize到416，再用值114填充一个832x832的新图片，在416x0.5到416x1.5的范围随机取一个中心点，把这个新图片分成四块，即左上、右上、右下、左下。原来从数据集取的4张图先resize到416，然后会以右下角、左下角、左上角、右上角对齐中心点的方式按在新图片上的对应部分的大小进行裁剪，剪去超出新图片上相应部分的区域，然后把该图片填充到新图片上对应的部分。对于标注框，超出区域的标注框统统抛弃。
mosaic数据增强的代码实现如下：
import os
import cv2
import torch
import numpy as np
import random
import math
from torch.utils.data import Dataset
from pycocotools.coco import COCO
import torch.nn.functional as F
class CocoDetection(Dataset):
def __init__(self,
image_root_dir,
annotation_root_dir,
set='train2017',
resize=416,
use_mosaic=False,
mosaic_center_range=[0.5, 1.5],
transform=None):
self.image_root_dir = image_root_dir
self.annotation_root_dir = annotation_root_dir
self.set_name = set
self.resize = resize
self.use_mosaic = use_mosaic
self.mosaic_center_range = mosaic_center_range
self.transform = transform
self.coco = COCO(
os.path.join(self.annotation_root_dir,
'instances_' + self.set_name + '.json'))
self.load_classes()
def load_classes(self):
self.image_ids = self.coco.getImgIds()
self.cat_ids = self.coco.getCatIds()
self.categories = self.coco.loadCats(self.cat_ids)
self.categories.sort(key=lambda x: x['id'])
# category_id is an original id,coco_id is set from 0 to 79
self.category_id_to_coco_label = {
category['id']: i
for i, category in enumerate(self.categories)
}
self.coco_label_to_category_id = {
v: k
for k, v in self.category_id_to_coco_label.items()
}
def __len__(self):
return len(self.image_ids)
def __getitem__(self, idx):
if self.use_mosaic:
# mosaic center x, y
x_ctr, y_ctr = [
int(
random.uniform(self.resize * self.mosaic_center_range[0],
self.resize * self.mosaic_center_range[1]))
for _ in range(2)
]
# all 4 image indices
imgs_indices = [idx] + [
random.randint(0,
len(self.image_ids) - 1) for _ in range(3)
]
final_annots = []
# combined image by 4 images
combined_img = np.full((self.resize * 2, self.resize * 2, 3),
114,
dtype=np.uint8)
for i, index in enumerate(imgs_indices):
img = self.load_image(index)
annot = self.load_annotations(index)
origin_height, origin_width, _ = img.shape
resize_factor = self.resize / max(origin_height, origin_width)
resize_height, resize_width = int(
origin_height * resize_factor), int(origin_width *
resize_factor)
img = cv2.resize(img, (resize_width, resize_height))
annot[:, :4] *= resize_factor
# top left img
if i == 0:
# combined image coordinates
x1a, y1a, x2a, y2a = max(x_ctr - resize_width,
0), max(y_ctr - resize_height,
0), x_ctr, y_ctr
# single img choosen area
x1b, y1b, x2b, y2b = resize_width - (
x2a - x1a), resize_height - (
y2a - y1a), resize_width, resize_height
# top right img
elif i == 1:
x1a, y1a, x2a, y2a = x_ctr, max(y_ctr - resize_height,
0), min(
x_ctr + resize_width,
self.resize * 2), y_ctr
x1b, y1b, x2b, y2b = 0, resize_height - (y2a - y1a), min(
resize_width, x2a - x1a), resize_height
# bottom left img
elif i == 2:
x1a, y1a, x2a, y2a = max(x_ctr - resize_width,
0), y_ctr, x_ctr, min(
self.resize * 2,
y_ctr + resize_height)
x1b, y1b, x2b, y2b = resize_width - (x2a - x1a), 0, max(
x_ctr, resize_width), min(y2a - y1a, resize_height)
# bottom right img
elif i == 3:
x1a, y1a, x2a, y2a = x_ctr, y_ctr, min(
x_ctr + resize_width,
self.resize * 2), min(self.resize * 2,
y_ctr + resize_height)
x1b, y1b, x2b, y2b = 0, 0, min(
resize_width, x2a - x1a), min(y2a - y1a, resize_height)
# combined_img[ymin:ymax, xmin:xmax]
combined_img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]
padw, padh = x1a - x1b, y1a - y1b
# annot coordinates transform
if annot.shape[0] > 0:
annot[:, 0] = annot[:, 0] + padw
annot[:, 1] = annot[:, 1] + padh
annot[:, 2] = annot[:, 2] + padw
annot[:, 3] = annot[:, 3] + padh
final_annots.append(annot)
final_annots = np.concatenate(final_annots, axis=0)
final_annots[:, 0:4] = np.clip(final_annots[:, 0:4], 0,
self.resize * 2)
final_annots = final_annots[final_annots[:, 2] -
final_annots[:, 0] > 1]
final_annots = final_annots[final_annots[:, 3] -
final_annots[:, 1] > 1]
sample = {'img': combined_img, 'annot': final_annots, 'scale': 1.}
else:
img = self.load_image(idx)
annot = self.load_annotations(idx)
scale = 1.
origin_height, origin_width, _ = img.shape
resize_factor = self.resize / max(origin_height, origin_width)
resize_height, resize_width = int(
origin_height * resize_factor), int(origin_width *
resize_factor)
img = cv2.resize(img, (resize_width, resize_height))
annot[:, :4] *= resize_factor
scale *= resize_factor
sample = {'img': img, 'annot': annot, 'scale': scale}
if self.transform:
sample = self.transform(sample)
return sample
def load_image(self, image_index):
image_info = self.coco.loadImgs(self.image_ids[image_index])[0]
path = os.path.join(self.image_root_dir, image_info['file_name'])
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img.astype(np.float32) / 255.
def load_annotations(self, image_index):
# get ground truth annotations
annotations_ids = self.coco.getAnnIds(
imgIds=self.image_ids[image_index], iscrowd=False)
annotations = np.zeros((0, 5))
# some images appear to miss annotations
if len(annotations_ids) == 0:
return annotations
# parse annotations
coco_annotations = self.coco.loadAnns(annotations_ids)
for _, a in enumerate(coco_annotations):
# some annotations have basically no width / height, skip them
if a['bbox'][2] < 1 or a['bbox'][3] < 1:
continue
annotation = np.zeros((1, 5))
if a['bbox'][2] > 0 and a['bbox'][3] > 0:
annotation[0, :4] = a['bbox']
annotation[0, 4] = self.find_coco_label_from_category_id(
a['category_id'])
annotations = np.append(annotations, annotation, axis=0)
# transform from [x_min, y_min, w, h] to [x_min, y_min, x_max, y_max]
annotations[:, 2] = annotations[:, 0] + annotations[:, 2]
annotations[:, 3] = annotations[:, 1] + annotations[:, 3]
return annotations
def find_coco_label_from_category_id(self, category_id):
return self.category_id_to_coco_label[category_id]
def find_category_id_from_coco_label(self, coco_label):
return self.coco_label_to_category_id[coco_label]
def num_classes(self):
return 80
def image_aspect_ratio(self, image_index):
image = self.coco.loadImgs(self.image_ids[image_index])[0]
return float(image['width']) / float(image['height'])