【庖丁解牛】从零实现RetinaNet（二）：RetinaNet中的resize、数据预读取、数据增强、collater处理
文章目录
- RetinaNet中的resize
- 数据预读取
- 数据增强
- collater处理
所有代码已上传到本人github repository：https://github.com/zgcr/pytorch-ImageNet-CIFAR-COCO-VOC-training
如果觉得有用，请点个star哟！
下列代码均在pytorch1.4版本中测试过，确认正确无误。
RetinaNet中的resize
在以往的分类任务中，对于图片resize我们不会保持其长宽比，而是长和宽直接resize到指定尺寸。在RetinaNet和其他大多数目标检测网络的训练中，resize必须要保持图片原来的长宽比。
RetinaNet的resize遵循以下原则：
- 首先确定两个resize后短长边的阈值，短边为800，长边为1333。
- 然后用800除以原始短边长度得到scale，然后用scale乘以原始长边长度得到缩放后长边长度。
- 检测缩放后长边的长度，如果缩放到长边长度不超过1333，则就用这个scale乘以原始长短边程度进行resize；如果缩放后长边的长度超过1333，则用1333除以原始长边长度得到新的scale，用新的scale乘以原始长短边长度进行resize。
800和1333的这两个阈值是固定值，如果我们想resize到其他尺寸，比如600，那么长边就不能超过600除以800乘以1333，即1000。在RetinaNet论文中给出了各个resize尺寸下的模型点数，这个resize尺寸实际上指的就是上面resize后短边的阈值。
各个resize尺寸下短长边的阈值：
min_length=400,max_length=667
min_length=500,max_length=833
min_length=600,max_length=1000
min_length=700,max_length=1166
min_length=800,max_length=1333
对于一个batch的图片，resize后可能会有两种情况：短边为800，长边小于等于1333；短边小于等于800，长边为1333。在进行训练时，我们必须要保证每个batch中图片的shape完全一致。注意每张图片的短边可能都不一样，有的可能是图片的宽，有的可能是图片的宽。因此，我们将所有图片都填充到1333x1333也就是resize后长边的尺寸，空白的地方全部用0值填充。
在github上的许多RetinaNet实现中，并没有采用上面的标准resize做法。而是直接用resize后尺寸除以长边来求得scale，这样长边resize后始终等于resize后尺寸。我对这两种方法进行了实验，发现直接用resize后尺寸除以长边求得scale的resize方法训练出来的模型效果更好。为了探究原因，我计算了这两种resize方法resize后的图片尺寸和图片面积，并把整个COCO数据集的图片面积求平均值，结果如下：
# retinanet resize method即上面一开始提到的RetinaNet标准Resize方法，my resize method即直接用resize后尺寸除以长边求得scale然后resize的方法。
# retinanet resize method，resize为短边阈值，per_image_average_area为图片resize后的平均面积，input shape为填充0值使所有图片长宽相等后输入网络时的尺寸
# resize=400,per_image_average_area=223743,input shape=[667,667]
# resize=500,per_image_average_area=347964,input shape=[833,833]
# resize=600,per_image_average_area=502820,input shape=[1000,1000]
# resize=700,per_image_average_area=682333,input shape=[1166,1166]
# resize=800,per_image_average_area=891169,input shape=[1333,1333]
# my resize method，resize为reisze后尺寸，per_image_average_area为图片resize后的平均面积，input shape为填充0值使所有图片长宽相等后输入网络时的尺寸
# resize=600,per_image_average_area=258182,input shape=[600,600]
# resize=667,per_image_average_area=318986,input shape=[667,667]
# resize=700,per_image_average_area=351427,input shape=[700,700]
# resize=800,per_image_average_area=459021,input shape=[800,800]
# resize=833,per_image_average_area=497426,input shape=[833,833]
# resize=900,per_image_average_area=580988,input shape=[900,900]
# resize=1000,per_image_average_area=717349,input shape=[1000,1000]
# resize=1166,per_image_average_area=974939,input shape=[1166,1166]
# resize=1333,per_image_average_area=1274284,input shape=[1333,1333]
通过上面的计算可以发现，如果以最终输入网络的input shape为基准，相同input shape下后一种resize方法得到的图片平均面积更大，换句话来说，就是resize后图片的清晰度更高。显然，网络学习的图片清晰度越高，则最终模型表现效果越好，这也与我前面的实验结论吻合。另外，相同图片清晰度下后一种resize方法需要的input shape更小，那么网络的前向计算flops就会更小，占用显存也会变少，在同一张显卡上进行训练时batchsize可以调整的更大，训练速度也更快，这对训练和推理都非常有利。因此，在后面的实验中，我使用后一种resize方法。
使用后一种resize方法后，如何与RetinaNet论文中报告的点数进行对点？
由于我只有2张2080ti显卡，在使用后一种resize方法时，resize=600时刚好一张2080ti的batchsize可以调到8，两张2080ti显卡的总batchsize为16，这个batchsize也是原论文中训练时采用的batchsize，方便对点。由于目标检测中使用apex自动混合精度训练时有时会使得loss变nan，因此在后面的实验中我们不使用apex。
在上面的计算结果中，第一种resize方法分辨率为400和500是图片的平均面积是223743和347964，两者尺寸差100，面积差124221。而后一种resize方法分辨率为600时图片的平均面积是258182，比第一种resize方法分辨率为400的图片平均面积大34439，是面积差124221的0.277倍。我们假设所有图片平均的宽高比是1比1，把0.277开根号再乘以100加上400，则可以得到后一种resize方法在与第一种resize方法保持图片平均面积一致的情况下的估算分辨率为450。在RetinaNet论文中，分辨率400和500报告的mAP分别为30.5和32.5，我们假设mAP会随着面积增大而线性增大。那么分辨率450时的估计mAP为2乘以0.277再加上30.5，大概31.1左右。后面的复现我会将模型表现与这个31.1相比，如果达到31.1，说明复现的模型达到了论文中报告的性能。
数据预读取
正常情况下，pytorch训练时先加载本次batch的数据，然后再进行本次batch的前向计算，最后反向传播。所谓数据与读取就是模型在进行本次batch的前向计算和反向传播时就预先加载下一个batch的数据，这样就节省了下加载数据的时间（相当于加载数据与前向计算和反向传播并行了）。
数据预读取代码如下：
class COCODataPrefetcher():
def __init__(self, loader):
self.loader = iter(loader)
self.stream = torch.cuda.Stream()
self.preload()
def preload(self):
try:
sample = next(self.loader)
self.next_input, self.next_annot = sample['img'], sample['annot']
except StopIteration:
self.next_input = None
self.next_annot = None
return
with torch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True)
self.next_annot = self.next_annot.cuda(non_blocking=True)
self.next_input = self.next_input.float()
def next(self):
torch.cuda.current_stream().wait_stream(self.stream)
input = self.next_input
annot = self.next_annot
self.preload()
return input, annot
下面提供了一个训练中使用数据预读取类的例子：
def train(train_loader, model, criterion, optimizer, scheduler, epoch, logger,
args):
cls_losses, reg_losses, losses = [], [], []
# switch to train mode
model.train()
iters = len(train_loader.dataset) // args.batch_size
prefetcher = COCODataPrefetcher(train_loader)
images, annotations = prefetcher.next()
iter_index = 1
while images is not None:
images, annotations = images.cuda().float(), annotations.cuda()
cls_heads, reg_heads, batch_anchors = model(images)
cls_loss, reg_loss = criterion(cls_heads, reg_heads, batch_anchors,
annotations)
loss = cls_loss + reg_loss
if cls_loss == 0.0 or reg_loss == 0.0:
optimizer.zero_grad()
continue
if args.apex:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
optimizer.step()
optimizer.zero_grad()
cls_losses.append(cls_loss.item())
reg_losses.append(reg_loss.item())
losses.append(loss.item())
images, annotations = prefetcher.next()
if iter_index % args.print_interval == 0:
logger.info(
f"train: epoch {epoch:0>3d}, iter [{iter_index:0>5d}, {iters:0>5d}], cls_loss: {cls_loss.item():.2f}, reg_loss: {reg_loss.item():.2f}, loss_total: {loss.item():.2f}"
)
iter_index += 1
scheduler.step(np.mean(losses))
return np.mean(cls_losses), np.mean(reg_losses), np.mean(losses)
train函数代表一个epoch内的训练过程。
数据增强
在分类任务中，我们直接调用torchvision.transform中的各个数据增强函数即可实现数据增强。在目标检测任务中，由于数据增强后图片上目标的位置可能发生变化，因此我们必须自己定义数据增强函数同时处理图片和目标的坐标。对于RetinaNet，训练只需要做randomflip数据增强，然后resize即可。测试时则直接resize。除此之外，我还实现了RandomCrop和RandomTranslate数据增强。
数据增强代码如下：
class RandomFlip(object):
def __init__(self, flip_prob=0.5):
self.flip_prob = flip_prob
def __call__(self, sample):
if np.random.uniform(0, 1) < self.flip_prob:
image, annots, scale = sample['img'], sample['annot'], sample[
'scale']
image = image[:, ::-1, :]
_, width, _ = image.shape
x1 = annots[:, 0].copy()
x2 = annots[:, 2].copy()
annots[:, 0] = width - x2
annots[:, 2] = width - x1
sample = {'img': image, 'annot': annots, 'scale': scale}
return sample
class RandomCrop(object):
def __init__(self, crop_prob=0.5):
self.crop_prob = crop_prob
def __call__(self, sample):
image, annots, scale = sample['img'], sample['annot'], sample['scale']
if annots.shape[0] == 0:
return sample
if np.random.uniform(0, 1) < self.crop_prob:
h, w, _ = image.shape
max_bbox = np.concatenate([
np.min(annots[:, 0:2], axis=0),
np.max(annots[:, 2:4], axis=0)
],
axis=-1)
max_left_trans, max_up_trans = max_bbox[0], max_bbox[1]
max_right_trans, max_down_trans = w - max_bbox[2], h - max_bbox[3]
crop_xmin = max(
0, int(max_bbox[0] - random.uniform(0, max_left_trans)))
crop_ymin = max(0,
int(max_bbox[1] - random.uniform(0, max_up_trans)))
crop_xmax = max(
w, int(max_bbox[2] + random.uniform(0, max_right_trans)))
crop_ymax = max(
h, int(max_bbox[3] + random.uniform(0, max_down_trans)))
image = image[crop_ymin:crop_ymax, crop_xmin:crop_xmax]
annots[:, [0, 2]] = annots[:, [0, 2]] - crop_xmin
annots[:, [1, 3]] = annots[:, [1, 3]] - crop_ymin
sample = {'img': image, 'annot': annots, 'scale': scale}
return sample
class RandomTranslate(object):
def __init__(self, translate_prob=0.5):
self.translate_prob = translate_prob
def __call__(self, sample):
image, annots, scale = sample['img'], sample['annot'], sample['scale']
if annots.shape[0] == 0:
return sample
if np.random.uniform(0, 1) < self.translate_prob:
h, w, _ = image.shape
max_bbox = np.concatenate([
np.min(annots[:, 0:2], axis=0),
np.max(annots[:, 2:4], axis=0)
],
axis=-1)
max_left_trans, max_up_trans = max_bbox[0], max_bbox[1]
max_right_trans, max_down_trans = w - max_bbox[2], h - max_bbox[3]
tx = random.uniform(-(max_left_trans - 1), (max_right_trans - 1))
ty = random.uniform(-(max_up_trans - 1), (max_down_trans - 1))
M = np.array([[1, 0, tx], [0, 1, ty]])
image = cv2.warpAffine(image, M, (w, h))
annots[:, [0, 2]] = annots[:, [0, 2]] + tx
annots[:, [1, 3]] = annots[:, [1, 3]] + ty
sample = {'img': image, 'annot': annots, 'scale': scale}
return sample
collater处理
对于一个batch的images和annotations，我们最后还需要用collater函数将images和annotations的shape全部对齐后才能输入模型进行训练。
collater函数代码如下：
def collater(data):
imgs = [s['img'] for s in data]
annots = [s['annot'] for s in data]
scales = [s['scale'] for s in data]
imgs = torch.from_numpy(np.stack(imgs, axis=0))
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], :] = annot
else:
annot_padded = torch.ones((len(annots), 1, 5)) * (-1)
imgs = imgs.permute(0, 3, 1, 2)
return {'img': imgs, 'annot': annot_padded, 'scale': scales}
对于images，由于我们前面的Resize类已经将其shape对齐了，所以这里不再做处理。对于annotations，由于每张图片标注的object数量都不一样，还有可能出现某张图上没有标注object的情况。我们取一个batch中所有图片里单张图片中标注object数量的最大值，然后用值-1填充其他图片的annotations，使得所有图片的annotations中object数量都等于这个最大值。在进行训练时，我们会在loss部分处理掉这部分值-1的annotations。