所有代碼已上傳到本人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。