1. 概述
本文介绍下 SSD 网络。
2. SSD 网络
2.1. SSD 网络数据处理
图像数据在输入网络之前和之后都需要对图像进行处理才能传入网络进行预测或训练。
首先看一下数据的预处理，这里以训练模式为例。
data_transform = {
"train": transform.Compose([transform.SSDCropping(),
transform.Resize(),
transform.ColorJitter(),
transform.ToTensor(),
transform.RandomHorizontalFlip(),
transform.Normalization(),
transform.AssignGTtoDefaultBox()]),
}
第一步对图像进行随机裁剪，看 SSDCropping 类。
- 初始化两个变量 self.sample_options 和 self.dboxes
- self.sample_options: 随机裁剪的选项
- dboxes300_coco 下文会介绍，该函数的功能是生成 Default Box 。
class SSDCropping(object): def __init__(self): self.sample_options = ( # Do nothing None, # min IoU, max IoU (0.1, None), (0.3, None), (0.5, None), (0.7, None), (0.9, None), # no IoU requirements (None, None), ) self.dboxes = dboxes300_coco()
SSDCropping 类的前向传播函数如下：
- 选中 self.sample_options 中的 None ，则不对图像进行随机裁剪
- htot, wtot: 原始图像的高宽
- min_iou, max_iou: IOU 最小值和最大值，如果是 None 的话，最小值设为负无穷，最大值设为正无穷。
- w, h: 随机生成裁剪后的 w 和 h ，保证宽高比例在 0.5-2 之间，注意这里的 w 和 h 是归一化后的。
- left, top: 随机生成左上角坐标， left 范围为 [0, wtot - w]; top 范围为 [0, htot - h]
- right, bottom: right = left + w; bottom = top + h
- ious: 计算 bboxes 与 [left, top, right, bottom] 的 iou ，如果 iou 值不是所有的都大于 min_iou 或者都小于 max_iou 的化则重新生成随机数。注意 bboxes 在加载数据时已经归一化了。
- masks: 查找所有的 gt box 的中心点有没有在随机采样的范围内，如果没有则重新采样
- bboxes[bboxes[:, 0] < left, 0]: 将 gt box 的坐标限制在裁剪图像的范围内
- bboxes = bboxes[masks, :]: 虑除不在采样范围中的 gt box
- left_idx, top_idx, right_idx, bottom_idx: 确定裁剪范围，相对于原图
- image.crop: 在原图上进行裁剪
- bboxes[:, 0] = (bboxes[:, 0] - left) / w: 调整 bbox 到裁剪后图像中的坐标，并归一化
def __call__(self, image, target): # Ensure always return cropped image while True: mode = random.choice(self.sample_options) if mode is None: # 不做随机裁剪处理 return image, target htot, wtot = target['height_width'] min_iou, max_iou = mode min_iou = float('-inf') if min_iou is None else min_iou max_iou = float('+inf') if max_iou is None else max_iou # Implementation use 5 iteration to find possible candidate for _ in range(5): # 0.3*0.3 approx. 0.1 w = random.uniform(0.3, 1.0) h = random.uniform(0.3, 1.0) if w/h < 0.5 or w/h > 2: # 保证宽高比例在0.5-2之间 continue # left 0 ~ wtot - w, top 0 ~ htot - h left = random.uniform(0, 1.0 - w) top = random.uniform(0, 1.0 - h) right = left + w bottom = top + h # boxes的坐标是在0-1之间的 bboxes = target["boxes"] ious = calc_iou_tensor(bboxes, torch.tensor([[left, top, right, bottom]])) # tailor all the bboxes and return # all(): Returns True if all elements in the tensor are True, False otherwise. if not ((ious > min_iou) & (ious < max_iou)).all(): continue # discard any bboxes whose center not in the cropped image xc = 0.5 * (bboxes[:, 0] + bboxes[:, 2]) yc = 0.5 * (bboxes[:, 1] + bboxes[:, 3]) # 查找所有的gt box的中心点有没有在采样patch中的 masks = (xc > left) & (xc < right) & (yc > top) & (yc < bottom) # if no such boxes, continue searching again # 如果所有的gt box的中心点都不在采样的patch中，则重新找 if not masks.any(): continue # 修改采样patch中的所有gt box的坐标（防止出现越界的情况） bboxes[bboxes[:, 0] < left, 0] = left bboxes[bboxes[:, 1] < top, 1] = top bboxes[bboxes[:, 2] > right, 2] = right bboxes[bboxes[:, 3] > bottom, 3] = bottom # 虑除不在采样patch中的gt box bboxes = bboxes[masks, :] # 获取在采样patch中的gt box的标签 labels = target['labels'] labels = labels[masks] # 裁剪patch left_idx = int(left * wtot) top_idx = int(top * htot) right_idx = int(right * wtot) bottom_idx = int(bottom * htot) image = image.crop((left_idx, top_idx, right_idx, bottom_idx)) # 调整裁剪后的bboxes坐标信息 bboxes[:, 0] = (bboxes[:, 0] - left) / w bboxes[:, 1] = (bboxes[:, 1] - top) / h bboxes[:, 2] = (bboxes[:, 2] - left) / w bboxes[:, 3] = (bboxes[:, 3] - top) / h # 更新crop后的gt box坐标信息以及标签信息 target['boxes'] = bboxes target['labels'] = labels return image, target
Resize 是将输入图像固定到 300*300 大小。
ColorJitter 是对图像颜色信息进行随机调整。
ToTensor 是将 PIL 图像转为 Tensor 格式，将数据格式变为 chw 后，除以 255 将像素值归一化到 [0,1] 之间。
RandomHorizontalFlip 是将图像随机水平翻转，同时翻转 bboxes 。
Normalization 是对数据进行标准化，相素值减去均值后除以方差。
AssignGTtoDefaultBox 是为 gt box 分配 Default Box 。
- self.default_box: 生成每个预测特征层各个尺度下的 Default Box
- self.encoder: Encoder 实现对 self.default_box 编码。
class AssignGTtoDefaultBox(object): def __init__(self): self.default_box = dboxes300_coco() self.encoder = Encoder(self.default_box) def __call__(self, image, target): boxes = target['boxes'] labels = target["labels"] # bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732) bboxes_out, labels_out = self.encoder.encode(boxes, labels) target['boxes'] = bboxes_out target['labels'] = labels_out return image, target
下面看 Encoder 类里做了什么。
- self.dboxes: Default Box 坐标，格式为: xmin ymin xmax ymax
- self.dboxes_xywh: Default Box 坐标，格式为: x y w h
- self.nboxes: Default Box 数量
class Encoder(object): def __init__(self, dboxes): self.dboxes = dboxes(order='ltrb') self.dboxes_xywh = dboxes(order='xywh').unsqueeze(dim=0) self.nboxes = self.dboxes.size(0) # default boxes的数量 self.scale_xy = dboxes.scale_xy self.scale_wh = dboxes.scale_wh
看一下 encode 函数。该函数负责划分正负样本， SSD 网络选择正样本的策略有两条，第一， 与 gt box 有最大 iou 的 Default Box 被划分为正样本；第二，与 gt box 的 iou > 0.5 的 Default Box 被划分为正样本。为什么这样选择正样本，首先对于第一条，是为了保证每一个 gt box 至少被分配一个 Default box 进行预测，这会存在另一个问题，一个 Default Box 可能负责预测多个 gt box ，这中情况要排除， encode 函数中考虑到了这种情况，对于第二条，是为了保证尽可能多的 Default Box 去预测某一个 gt box ，因为一个 Default Box 就是一个预测器，这里可能理解为使用了集成学习的思想。
- 传入参数为从 VOC 数据集中读取的 bbox 信息和 label 信息，以及阈值 criteria=0.5
- ious: 计算每一个 gt box 与 Default Box 的 IOU
- best_dbox_ious, best_dbox_idx: 计算 Default Box 对应的最大 IOU 的 gt box ，返回最大 IOU 值和索引
- best_bbox_ious, best_bbox_idx: 计算 gt box 对应的最大 IOU 的 Default Box ，返回最大 IOU 值和索引
- best_dbox_ious.index_fill_: 将每个 gt box 匹配到的最佳 Default Box 设置为正样本，值设为 2.0 ，大于 0.5 的都将设为正样本 。 index_fill_(dim, index, val) 根据 index 中指定的顺序索引, 用值 val 填充 self tensor 中的元素。
- idx: 列表，0 ~ 每张图片的 bbox 个数
- best_dbox_idx[best_bbox_idx[idx]] = idx: 之前已经计算了 gt box 与 Default Box 之间的 iou ，并且找到了每个 Default Box 对应的最大 iou 的 gt box ，以及每个 gt box 对应的最大 iou 的 Default Box ，但是不同的 gt box 可能与同一个 Default Box 有最大的 iou ，这句代码的作用就是让一个 Default Box 只对应一个 gt box 。为了方便训练，一个 gt box 可以对应多个 Default Box ，但一个 Default Box 只能对应一个 gt box ，这很好理解，如果一个 Default Box 对应多个 gt box ，那么我们网络输出的参数该去拟合哪一个 gt box 呢？
- masks: 保存与 Box iou 大于 0.5 的 Default Box ，这些 Default Box 将作为正样本
- 下面的代码是生成最终的标签和边界框，用来计算损失，这里一张图像一共生成 8732 个 Default Box ，所以最终标签的维度是 [8732] ，边界框的维度是 [8732, 4]
- labels_out: 创建一个 8732 大小的 Tensor ，保存所有 Default Box 的标签
- labels_out[masks]: 为正样本设置为对应的类别信息
- bboxes_out: 创建一个 8732 大小的边界框参数
- bboxes_out[masks, :]: 为正样本设置为对应的边界框信息
- 最后将 bboxes_out 转换成中心点坐标和宽高模式
def encode(self, bboxes_in, labels_in, criteria=0.5): """ encode: input : bboxes_in (Tensor nboxes x 4), labels_in (Tensor nboxes) output : bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732) criteria : IoU threshold of bboexes """ ious = calc_iou_tensor(bboxes_in, self.dboxes) # [nboxes, 8732] # [8732,] best_dbox_ious, best_dbox_idx = ious.max(dim=0) # 寻找每个default box匹配到的最大IoU bboxes_in # [nboxes,] best_bbox_ious, best_bbox_idx = ious.max(dim=1) # 寻找每个bboxes_in匹配到的最大IoU default box # set best ious 2.0 # 将每个bboxes_in匹配到的最佳default box设置为正样本（对应论文中Matching strategy的第一条） best_dbox_ious.index_fill_(0, best_bbox_idx, 2.0) # 将相应default box的匹配最大IoU bboxes_in信息替换成best_bbox_idx idx = torch.arange(0, best_bbox_idx.size(0), dtype=torch.int64) best_dbox_idx[best_bbox_idx[idx]] = idx # filter IoU > 0.5 # 寻找与bbox_in iou大于0.5的default box,对应论文中Matching strategy的第二条(这里包括了第一条匹配到的信息) masks = best_dbox_ious > criteria # [8732,] labels_out = torch.zeros(self.nboxes, dtype=torch.int64) labels_out[masks] = labels_in[best_dbox_idx[masks]] bboxes_out = self.dboxes.clone() # 将default box匹配到正样本的地方设置成对应正样本box信息 bboxes_out[masks, :] = bboxes_in[best_dbox_idx[masks], :] # Transform format to xywh format x = 0.5 * (bboxes_out[:, 0] + bboxes_out[:, 2]) # x y = 0.5 * (bboxes_out[:, 1] + bboxes_out[:, 3]) # y w = bboxes_out[:, 2] - bboxes_out[:, 0] # w h = bboxes_out[:, 3] - bboxes_out[:, 1] # h bboxes_out[:, 0] = x bboxes_out[:, 1] = y bboxes_out[:, 2] = w bboxes_out[:, 3] = h return bboxes_out, labels_out
2.2. SSD 网络结构
SSD 网络代码在 SSD300 这个类中实现。
class SSD300(nn.Module):
def __init__(self, backbone=None, num_classes=21):
super(SSD300, self).__init__()
if backbone is None:
raise Exception("backbone is None")
if not hasattr(backbone, "out_channels"):
raise Exception("the backbone not has attribute: out_channel")
self.feature_extractor = backbone
self.num_classes = num_classes
# out_channels = [1024, 512, 512, 256, 256, 256] for resnet50
self._build_additional_features(self.feature_extractor.out_channels)
self.num_defaults = [4, 6, 6, 6, 4, 4]
location_extractors = []
confidence_extractors = []
# out_channels = [1024, 512, 512, 256, 256, 256] for resnet50
for nd, oc in zip(self.num_defaults, self.feature_extractor.out_channels):
# nd is number_default_boxes, oc is output_channel
location_extractors.append(nn.Conv2d(oc, nd * 4, kernel_size=3, padding=1))
confidence_extractors.append(nn.Conv2d(oc, nd * self.num_classes, kernel_size=3, padding=1))
self.loc = nn.ModuleList(location_extractors)
self.conf = nn.ModuleList(confidence_extractors)
self._init_weights()
default_box = dboxes300_coco()
self.compute_loss = Loss(default_box)
self.encoder = Encoder(default_box)
self.postprocess = PostProcess(default_box)
SSD300 在上一篇文章中已经介绍了一部分，这篇文章从 Default Box 的生成开始。 Default Box 与 Faster-RCNN 网络中的 Anchor 类似，但是生成规则有些不同，下面我们介绍如何生成 Default Box 。
- figsize: 输入网络的图像的大小
- feat_size: 每个预测层的feature map尺寸
- step: 每个特征层上的一个 cell 在原图上的跨度
- scales: 每个特征层上预测的 Default Box 的 scale ，当前层的 Default Box 的尺度由 sacles[l] 和 scales[l+1] 的值共同决定
- aspect_ratios: 每个预测特征层上预测的 Default Box 的 ratios
- dboxes: 保存生成的 dboxes
def dboxes300_coco(): figsize = 300 feat_size = [38, 19, 10, 5, 3, 1] steps = [8, 16, 32, 64, 100, 300] scales = [21, 45, 99, 153, 207, 261, 315] aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]] dboxes = DefaultBoxes(figsize, feat_size, steps, scales, aspect_ratios) return dboxes
生成 Default Box 的实现在 DefaultBoxes 类中。
- 遍历每一个预测特征层，首先计算归一化后的 Default Box ，相当于 Faster-RCNN 的 anchros base
- Default Box 的计算公式为：
- sk1 = scales[idx] / fig_size
- sk2 = scales[idx + 1] / fig_size
- sk3 = sqrt(sk1 * sk2)
- 比例为 1:1 的 Default Box 使用 (sk1, sk1), (sk3, sk3)
- 遍历 aspect_ratios ，计算不同比例的 Default Box
- w, h = sk1 * sqrt(alpha), sk1 / sqrt(alpha): w, h 存储 Default Box 的宽和高
- 以上生成了归一化的 Default Box ，接下来将归一化的 Default Box 映射到原图上，这里在实现上与 Faster-RCNN 也有区别， Faster-RCNN 使用了 meshgrid 函数实现，这里是使用 itertools.product 函数实现，原理是一样的
- self.default_boxes: 存储的是宽高和中心点坐标的形式的 Default Box
- self.dboxes_ltrb: 存储的是左上角和右下角坐标的形式的 Default Box
class DefaultBoxes(object): def __init__(self, fig_size, feat_size, steps, scales, aspect_ratios, scale_xy=0.1, scale_wh=0.2): self.fig_size = fig_size # 输入网络的图像大小 300 # [38, 19, 10, 5, 3, 1] self.feat_size = feat_size # 每个预测层的feature map尺寸 self.scale_xy_ = scale_xy self.scale_wh_ = scale_wh # According to https://github.com/weiliu89/caffe # Calculation method slightly different from paper # [8, 16, 32, 64, 100, 300] self.steps = steps # 每个特征层上的一个cell在原图上的跨度 # [21, 45, 99, 153, 207, 261, 315] self.scales = scales # 每个特征层上预测的default box的scale fk = fig_size / np.array(steps) # 计算每层特征层的fk # [[2], [2, 3], [2, 3], [2, 3], [2], [2]] self.aspect_ratios = aspect_ratios # 每个预测特征层上预测的default box的ratios self.default_boxes = [] # size of feature and number of feature # 遍历每层特征层，计算default box for idx, sfeat in enumerate(self.feat_size): sk1 = scales[idx] / fig_size # scale转为相对值[0-1] sk2 = scales[idx + 1] / fig_size # scale转为相对值[0-1] sk3 = sqrt(sk1 * sk2) # 先添加两个1:1比例的default box宽和高 all_sizes = [(sk1, sk1), (sk3, sk3)] # 再将剩下不同比例的default box宽和高添加到all_sizes中 for alpha in aspect_ratios[idx]: w, h = sk1 * sqrt(alpha), sk1 / sqrt(alpha) all_sizes.append((w, h)) all_sizes.append((h, w)) # 计算当前特征层对应原图上的所有default box for w, h in all_sizes: for i, j in itertools.product(range(sfeat), repeat=2): # i -> 行（y）， j -> 列（x） # 计算每个default box的中心坐标（范围是在0-1之间） cx, cy = (j + 0.5) / fk[idx], (i + 0.5) / fk[idx] self.default_boxes.append((cx, cy, w, h)) # 将default_boxes转为tensor格式 self.dboxes = torch.as_tensor(self.default_boxes, dtype=torch.float32) # 这里不转类型会报错 self.dboxes.clamp_(min=0, max=1) # 将坐标（x, y, w, h）都限制在0-1之间 # For IoU calculation # ltrb is left top coordinate and right bottom coordinate # 将(x, y, w, h)转换成(xmin, ymin, xmax, ymax)，方便后续计算IoU(匹配正负样本时) self.dboxes_ltrb = self.dboxes.clone() self.dboxes_ltrb[:, 0] = self.dboxes[:, 0] - 0.5 * self.dboxes[:, 2] # xmin self.dboxes_ltrb[:, 1] = self.dboxes[:, 1] - 0.5 * self.dboxes[:, 3] # ymin self.dboxes_ltrb[:, 2] = self.dboxes[:, 0] + 0.5 * self.dboxes[:, 2] # xmax self.dboxes_ltrb[:, 3] = self.dboxes[:, 1] + 0.5 * self.dboxes[:, 3] # ymax
接下来看一下 SSD 网络中的损失函数，实现在类 Loss 中。 SSD 网络的损失函数包括类别损失和定位损失。
- self.location_loss: 定位损失使用 SmoothL1Loss 损失函数
- self.dboxes: 将 dboxes 转换成可以训练的类型 Parameter ，维度变化 [num_anchors, 4] -> [1, 4, num_anchors]
- self.confidence_loss: 类别损失使用 CrossEntropyLoss 交叉熵损失函数
- self.scale_xy, self.scale_wh: 边界框回归参数的缩放因子，计算损失时，将缩放因子放大，预测时将缩放因子缩小回放大前的尺寸。
class Loss(nn.Module): def __init__(self, dboxes): super(Loss, self).__init__() # Two factor are from following links # http://jany.st/post/2017-11-05-single-shot-detector-ssd-from-scratch-in-tensorflow.html self.scale_xy = 1.0 / dboxes.scale_xy # 10 self.scale_wh = 1.0 / dboxes.scale_wh # 5 self.location_loss = nn.SmoothL1Loss(reduction='none') # [num_anchors, 4] -> [4, num_anchors] -> [1, 4, num_anchors] self.dboxes = nn.Parameter(dboxes(order="xywh").transpose(0, 1).unsqueeze(dim=0), requires_grad=False) self.confidence_loss = nn.CrossEntropyLoss(reduction='none')
Loss 类的前向传播函数如下所示：
- Loss 类的前向传播有 4 个参数，分别预测边界框、预测类别、真实边界框、真实类别
- mask: 保存正样本的 mask ， torch.gt(a, b) ，如果 a > b 则为 True ，不大于 b 则为 False ，预处理时，已将正样本对应的标签赋值
- pos_num: 计算每个 batch 中每张图片的正样本个数
- self._location_vec: 计算 GT box 相对 Default Box 的回归参数，也就是边界框的正样本
- self.location_loss: 计算定位损失，只有正样本有定位损失
- self.confidence_loss: 计算类别损失
- con_neg: 拷贝一份类别损失，用于计算负样本，负样本的选取原则是选取类别损失最高的样本作为负样本，并且负样本的选取比例是正样本的 3 倍
- con_neg[mask] = 0.0: mask 保存着正样本的损失，将正样本的值全部设为 0 ，从剩下的样本中选取负样本
- con_idx: 类别损失降序排列后的索引，表示排序后的 tensor 的值在原 tensor 上的索引。我们的目的是选取损失最大的 topk
- con_rank: 对 con_idx 进行升序排序后的索引。将损失值按降序的顺序分配索引，比如最大损失值的索引为 2 ，那么对应 con_ran 的索引为 2 的位置为 0 。
- neg_num: 用于损失计算的负样本数是正样本的 3 倍， 但不能超过总样本数 8732
- neg_mask: 负样本的索引，这里获得负样本的索引的方法虽然很巧妙，但是不太好理解，建议通过实验理解
- con_loss: 类别损失的 loss 是选取的正样本 loss 加上选取的负样本 loss 之和
- total_loss: 网络的损失为定位损失和类别损失的和
- num_mask: 统计一个 batch 中每张图像中是否存在正样本
- pos_num.float().clamp(min=1e-6): 如果图片中没有正样本，则给对应的 pos_num 设置为一个较小的值，后面防止被除数为 0
- (total_loss * num_mask / pos_num).mean(dim=0): 只计算存在正样本的图像损失
def forward(self, ploc, plabel, gloc, glabel): # type: (Tensor, Tensor, Tensor, Tensor) -> Tensor # 获取正样本的mask Tensor: [N, 8732] mask = torch.gt(glabel, 0) # (gt: >) # mask1 = torch.nonzero(glabel) # 计算一个batch中的每张图片的正样本个数 Tensor: [N] pos_num = mask.sum(dim=1) # 计算gt的location回归参数 Tensor: [N, 4, 8732] vec_gd = self._location_vec(gloc) # sum on four coordinates, and mask # 计算定位损失(只有正样本) loc_loss = self.location_loss(ploc, vec_gd).sum(dim=1) # Tensor: [N, 8732] loc_loss = (mask.float() * loc_loss).sum(dim=1) # Tenosr: [N] # hard negative mining Tenosr: [N, 8732] con = self.confidence_loss(plabel, glabel) # positive mask will never selected # 获取负样本 con_neg = con.clone() con_neg[mask] = 0.0 # 按照confidence_loss降序排列 con_idx(Tensor: [N, 8732]) _, con_idx = con_neg.sort(dim=1, descending=True) _, con_rank = con_idx.sort(dim=1) # 这个步骤比较巧妙 # number of negative three times positive # 用于损失计算的负样本数是正样本的3倍（在原论文Hard negative mining部分）， # 但不能超过总样本数8732 neg_num = torch.clamp(3 * pos_num, max=mask.size(1)).unsqueeze(-1) neg_mask = torch.lt(con_rank, neg_num) # (lt: <) Tensor [N, 8732] # confidence最终loss使用选取的正样本loss+选取的负样本loss con_loss = (con * (mask.float() + neg_mask.float())).sum(dim=1) # Tensor [N] # avoid no object detected # 避免出现图像中没有GTBOX的情况 total_loss = loc_loss + con_loss # eg. [15, 3, 5, 0] -> [1.0, 1.0, 1.0, 0.0] num_mask = torch.gt(pos_num, 0).float() # 统计一个batch中的每张图像中是否存在正样本 pos_num = pos_num.float().clamp(min=1e-6) # 防止出现分母为零的情况 ret = (total_loss * num_mask / pos_num).mean(dim=0) # 只计算存在正样本的图像损失 return ret
边界框回归参数计算实现如下：
- 定位损失包括中心点坐标的偏移和宽高的缩放。
- tx = (x-xa)/wa
- ty = (y-ya)/ha
- tw = log(w/wa)
- th = log(h/ha)
def _location_vec(self, loc): # type: (Tensor) -> Tensor gxy = self.scale_xy * (loc[:, :2, :] - self.dboxes[:, :2, :]) / self.dboxes[:, 2:, :] # Nx2x8732 gwh = self.scale_wh * (loc[:, 2:, :] / self.dboxes[:, 2:, :]).log() # Nx2x8732 return torch.cat((gxy, gwh), dim=1).contiguous()
SSD 网络中损失函数与 Faster-RCNN 网络相同，分别使用 SmoothL1 和交叉熵损失计算定位损失和类别损失。
SSD 网络的前向传播，网络参数是 VOC 图像信息和标签。
- detection_features: 保存 6 层预测特征层，包括修改后的主干网络输出的预测特征层和额外添加的 5 层预测特征层。
- self.bbox_view: 计算边界回归参数和类别置信度
- 如果是训练模式计算损失 self.compute_loss ，并返回损失值
- 如果是预测模式则将计算得到的类别信息和位置信息输入给 postprocess 模块处理。
- self.postprocess: 将预测回归参数叠加到 Default Box 上得到最终预测 box ，并执行非极大值抑制虑除重叠框
def forward(self, image, targets=None): x = self.feature_extractor(image) # Feature Map 38x38x1024, 19x19x512, 10x10x512, 5x5x256, 3x3x256, 1x1x256 detection_features = torch.jit.annotate(List[Tensor], []) # [x] detection_features.append(x) for layer in self.additional_blocks: x = layer(x) detection_features.append(x) # Feature Map 38x38x4, 19x19x6, 10x10x6, 5x5x6, 3x3x4, 1x1x4 locs, confs = self.bbox_view(detection_features, self.loc, self.conf) # For SSD 300, shall return nbatch x 8732 x {nlabels, nlocs} results # 38x38x4 + 19x19x6 + 10x10x6 + 5x5x6 + 3x3x4 + 1x1x4 = 8732 if self.training: if targets is None: raise ValueError("In training mode, targets should be passed") # bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732) bboxes_out = targets['boxes'] bboxes_out = bboxes_out.transpose(1, 2).contiguous() # print(bboxes_out.is_contiguous()) labels_out = targets['labels'] # print(labels_out.is_contiguous()) # ploc, plabel, gloc, glabel loss = self.compute_loss(locs, confs, bboxes_out, labels_out) return {"total_losses": loss} # 将预测回归参数叠加到default box上得到最终预测box，并执行非极大值抑制虑除重叠框 # results = self.encoder.decode_batch(locs, confs) results = self.postprocess(locs, confs) return results
下面看一下 PostProcess 类。该类的作用的是将网络输出的预测信息进行处理得到原图上对应的预测信息。
- 首先将传入的 dboxes 信息的维度由 [num_anchors, 4] 转换成 [1, num_anchors, 4]
- iou 阈值设为 0.5
- 最大输出边界框个数为 100
class PostProcess(nn.Module): def __init__(self, dboxes): super(PostProcess, self).__init__() # [num_anchors, 4] -> [1, num_anchors, 4] self.dboxes_xywh = nn.Parameter(dboxes(order='xywh').unsqueeze(dim=0), requires_grad=False) self.scale_xy = dboxes.scale_xy # 0.1 self.scale_wh = dboxes.scale_wh # 0.2 self.criteria = 0.5 self.max_output = 100
PostProcess 类的前向传播函数。
- 输入参数是网络预测得到的边界框回归参数和类别预测值
- self.scale_back_batch: 通过预测的边界框回归参数得到最终预测坐标, 将类别预测结果通过 softmax 处理
- self.decode_single_new: 对预测得到的信息进行处理，对越界的 bbox 进行裁剪；移除归为背景类别的概率信息；移除低概率目标；移除控的边界框；进行非极大值抑制处理；保留 100 个 bbox 输出
def forward(self, bboxes_in, scores_in): # 通过预测的boxes回归参数得到最终预测坐标, 将预测目标score通过softmax处理 bboxes, probs = self.scale_back_batch(bboxes_in, scores_in) outputs = torch.jit.annotate(List[Tuple[Tensor, Tensor, Tensor]], []) # 遍历一个batch中的每张image数据 # bboxes: [batch, 8732, 4] for bbox, prob in zip(bboxes.split(1, 0), probs.split(1, 0)): # split_size, split_dim # bbox: [1, 8732, 4] bbox = bbox.squeeze(0) prob = prob.squeeze(0) outputs.append(self.decode_single_new(bbox, prob, self.criteria, self.max_output)) return outputs
数据解码的函数这里就不再进一步介绍了，跟 Faster-RCNN 基本是一样的。
