CV-目标检测-炼丹

模型训练

前面的章节，我们已经对目标检测训练的各个重要的知识点进行了讲解，下面我们需要将整个流程串起来，对模型进行训练。

目标检测网络的训练大致是如下的流程：

• 设置各种超参数
• 定义数据加载模块 dataloader
• 定义网络 model
• 定义损失函数 loss
• 定义优化器 optimizer
• 遍历训练数据，预测-计算loss-反向传播

import time                                                                                                                                    
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
from model import tiny_detector, MultiBoxLoss
from datasets import PascalVOCDataset
from utils import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cudnn.benchmark = True

# Data parameters
data_folder = '../../../dataset/VOCdevkit'  # data files root path
keep_difficult = True  # use objects considered difficult to detect?
n_classes = len(label_map)  # number of different types of objects

# Learning parameters
total_epochs = 230 # number of epochs to train
batch_size = 32  # batch size
workers = 4  # number of workers for loading data in the DataLoader
print_freq = 100  # print training status every __ batches
lr = 1e-3  # learning rate
decay_lr_at = [150, 190]  # decay learning rate after these many epochs
decay_lr_to = 0.1  # decay learning rate to this fraction of the existing learning rate
momentum = 0.9  # momentum
weight_decay = 5e-4  # weight decay

def main():
    """
    Training.
    """
    # Initialize model and optimizer
    model = tiny_detector(n_classes=n_classes)
    criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy)
    optimizer = torch.optim.SGD(params=model.parameters(),
                                lr=lr, 
                                momentum=momentum,
                                weight_decay=weight_decay)

    # Move to default device
    model = model.to(device)
    criterion = criterion.to(device)

    # Custom dataloaders
    train_dataset = PascalVOCDataset(data_folder,
                                     split='train',
                                     keep_difficult=keep_difficult)
    train_loader = torch.utils.data.DataLoader(train_dataset,   
                                    batch_size=batch_size,
                                    shuffle=True,
                                    collate_fn=train_dataset.collate_fn, 
                                    num_workers=workers,
                                    pin_memory=True) 

    # Epochs
    for epoch in range(total_epochs):
        # Decay learning rate at particular epochs
        if epoch in decay_lr_at:
            adjust_learning_rate(optimizer, decay_lr_to)

        # One epoch's training                                                                                                                 
        train(train_loader=train_loader,
              model=model,
              criterion=criterion,
              optimizer=optimizer,
              epoch=epoch)

        # Save checkpoint
        save_checkpoint(epoch, model, optimizer)

其中，我们对单个epoch的训练逻辑进行了封装，其具体实现如下：

def train(train_loader, model, criterion, optimizer, epoch):
    """
    One epoch's training.

    :param train_loader: DataLoader for training data
    :param model: model
    :param criterion: MultiBox loss
    :param optimizer: optimizer
    :param epoch: epoch number
    """
    model.train()  # training mode enables dropout

    batch_time = AverageMeter()  # forward prop. + back prop. time
    data_time = AverageMeter()  # data loading time
    losses = AverageMeter()  # loss

    start = time.time()

    # Batches
    for i, (images, boxes, labels, _) in enumerate(train_loader):
        data_time.update(time.time() - start)

        # Move to default device
        images = images.to(device)  # (batch_size (N), 3, 224, 224)
        boxes = [b.to(device) for b in boxes]
        labels = [l.to(device) for l in labels]

        # Forward prop.
        predicted_locs, predicted_scores = model(images)  # (N, 441, 4), (N, 441, n_classes)

        # Loss
        loss = criterion(predicted_locs, predicted_scores, boxes, labels)  # scalar

        # Backward prop.
        optimizer.zero_grad()
        loss.backward()

        # Update model
        optimizer.step()

        losses.update(loss.item(), images.size(0))
        batch_time.update(time.time() - start)

        start = time.time()

        # Print status
        if i % print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch,
                                                                  i, 
                                                                  len(train_loader),
                                                                  batch_time=batch_time,
                                                                  data_time=data_time, 
                                                                  loss=losses))
    del predicted_locs, predicted_scores, images, boxes, labels  # free some memory since their histories may be stored

后处理

之前我们的提到过，模型不是直接预测的目标框信息，而是预测的基于anchor的偏移，且经过了编码。因此后处理的第一步，就是对模型的回归头的输出进行解码，拿到真正意义上的目标框的预测结果。

后处理还需要做什么呢？由于我们预设了大量的先验框，因此预测时在目标周围会形成大量高度重合的检测框，而我们目标检测的结果只希望保留一个足够准确的预测框，所以就需要使用某些算法对检测框去重。这个去重算法叫做NMS，下面我们详细来讲一讲。
在这里插入图片描述

单图预测推理

当模型已经训练完成后，下面我们来看下如何对单张图片进行推理，得到目标检测结果。

首先我们需要导入必要的python包，然后加载训练好的模型权重。

随后我们需要定义预处理函数。为了达到最好的预测效果，测试环节的预处理方案需要和训练时保持一致，仅去除掉数据增强相关的变换即可。

因此，这里我们需要进行的预处理为：

• 将图片缩放为 224 * 224 的大小
• 转换为 Tensor 并除 255
• 进行减均值除方差的归一化

VOC测试集评测

以分类模型中最简单的二分类为例，对于这种问题，我们的模型最终需要判断样本的结果是0还是1，或者说是positive还是negative。我们通过样本的采集，能够直接知道真实情况下，哪些数据结果是positive，哪些结果是negative。同时，我们通过用样本数据跑出分类模型的结果，也可以知道模型认为这些数据哪些是positive，哪些是negative。因此，我们就能得到这样四个基础指标，称他们是一级指标（最底层的）：

AP指标即Average Precision 即平均精确度。

mAP即Mean Average Precision即平均AP值，是对多个验证集个体求平均AP值，作为object detection中衡量检测精度的指标。

参考材料

[1] Datawhale：动手学CV之目标检测

[2] CSDN：检测学习笔记

[3] CSDN：CV-目标检测-不讲武德-炼丹与品尝