MobileNetV2-SSDLite代码分析-4 ssd

Github-pytorch-ssd
vision/ssd
ssd.py

import torch.nn as nn
import torch
import numpy as np
from typing import List, Tuple
import torch.nn.functional as F

from ..utils import box_utils
from collections import namedtuple
GraphPath = namedtuple("GraphPath", ['s0', 'name', 's1']) 

namedtuple: 有字典功能的tuple

def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    ReLU = nn.ReLU if onnx_compatible else nn.ReLU6
    return Sequential(
        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
               groups=in_channels, stride=stride, padding=padding),
        BatchNorm2d(in_channels),
        ReLU(),
        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
    )

SSD_lite的lite就体现在上述这段代码中,它将SSD中的conv用深度可分离卷积来代替了.

关于SSD的定义，先放一段mobilenetv2_ssd_lite.py中的代码在这里，可以大致了解每个参数的含义

def create_mobilenetv2_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
    base_net = MobileNetV2(width_mult=width_mult, use_batch_norm=use_batch_norm,
                           onnx_compatible=onnx_compatible).features

    source_layer_indexes = [
        GraphPath(14, 'conv', 3),
        19,
    ]
    extras = ModuleList([
        InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
        InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
        InvertedResidual(256, 256, stride=2, expand_ratio=0.5),
        InvertedResidual(256, 64, stride=2, expand_ratio=0.25)
    ])

    regression_headers = ModuleList([
        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * 4,
                        kernel_size=3, padding=1, onnx_compatible=False),
        SeperableConv2d(in_channels=1280, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
    ])

    classification_headers = ModuleList([
        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
        SeperableConv2d(in_channels=1280, out_channels=6 * num_classes, kernel_size=3, padding=1),
        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
    ])

    return SSD(num_classes, base_net, source_layer_indexes,
               extras, classification_headers, regression_headers, is_test=is_test, config=config)

 在上面的代码中, create_mobilenetv2_ssd_lite定义了source_layer_indexes, extras, regression_headers, classification_headers

最后有个config属性, 这是预先定义好的关于生成anchor的设置, 包括: 图像尺寸、平均值、方差、iou阈值，anchor属性，和生成priors的函数。

其中anchor属性：'SSDSpec', ['feature_map_size', 'shrinkage', 'box_sizes', 'aspect_ratios']

来看看上面定义的四个东西具体是怎么在网络中使用的

class SSD(nn.Module):
    def __init__(self, num_classes: int, base_net: nn.ModuleList, source_layer_indexes: List[int],
                 extras: nn.ModuleList, classification_headers: nn.ModuleList,
                 regression_headers: nn.ModuleList, is_test=False, config=None, device=None):
        """Compose a SSD model using the given components.
        """
        super(SSD, self).__init__()

        self.num_classes = num_classes
        self.base_net = base_net
        self.source_layer_indexes = source_layer_indexes
        self.extras = extras
        self.classification_headers = classification_headers
        self.regression_headers = regression_headers
        self.is_test = is_test
        self.config = config

        # register layers in source_layer_indexes by adding them to a module list
        self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes if isinstance(t, tuple) and not isinstance(t, GraphPath)])
        if device:
            self.device = device
        else:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if is_test:
            self.config = config
            self.priors = config.priors.to(self.device)
            
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        confidences = []
        locations = []
        start_layer_index = 0
        header_index = 0
        for end_layer_index in self.source_layer_indexes:
            if isinstance(end_layer_index, GraphPath):
                path = end_layer_index
                end_layer_index = end_layer_index.s0
                added_layer = None
            elif isinstance(end_layer_index, tuple):
                added_layer = end_layer_index[1]
                end_layer_index = end_layer_index[0]
                path = None
            else:
                added_layer = None
                path = None
            for layer in self.base_net[start_layer_index: end_layer_index]:
                x = layer(x)
            if added_layer:
                y = added_layer(x)
            else:
                y = x
            if path:
                sub = getattr(self.base_net[end_layer_index], path.name)
                for layer in sub[:path.s1]:
                    x = layer(x)
                y = x
                for layer in sub[path.s1:]:
                    x = layer(x)
                end_layer_index += 1
            start_layer_index = end_layer_index
            confidence, location = self.compute_header(header_index, y)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)

        for layer in self.base_net[end_layer_index:]:
            x = layer(x)

        for layer in self.extras:
            x = layer(x)
            confidence, location = self.compute_header(header_index, x)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)

        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)
        
        if self.is_test:
            confidences = F.softmax(confidences, dim=2)
            boxes = box_utils.convert_locations_to_boxes(
                locations, self.priors, self.config.center_variance, self.config.size_variance
            )
            boxes = box_utils.center_form_to_corner_form(boxes)
            return confidences, boxes
        else:
            return confidences, locations

    def compute_header(self, i, x):
        confidence = self.classification_headers[i](x)
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)

        location = self.regression_headers[i](x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)

1. source_layer_add_ones

   GraphPath = namedtuple("GraphPath", ['s0', 'name', 's1'])
   source_layer_indexes = [
           GraphPath(14, 'conv', 3),
           19,
       ]
   self.source_layer_add_ons = nn.ModuleList([t[1] for t in source_layer_indexes if isinstance(t, tuple) and not isinstance(t, GraphPath)])
   # 这里实际sorce_layer_add_ons为空
   

2. forward过程

   针对source_layer_indexes:

   • 找到start_layer_index和end_layer_index

     1. GraphPath(s0=14, name='conv', s1=3)

        path = GraphPath(s0=14, name='conv', s1=3)，对应着我想在basenet的第14层找到名为conv的层的集合， 该集合中的第3层我想用来做分类与回归

        end_layer_index = 14

     2. 19

        则直接end_layer_index = 19

   • 针对base_net[start_layer_index, end_layer_index]中的每一个layer

     • x = layer(x)
     • 将最后一层的结果保存成y（y就是拿来做检测和回归的特征）
     • if path:
       • sub = mobilenet v2中第14层的conv（nn.module，包含了dwconv, relu, conv, relu等等）
       • sub[1: 3]：正常的正向传播
       • 将sub中的第3层的结果取出，保存为y
       • sub[3: ]：正常的正向传播
       • end_layer_index += 1
     • start_layer_index = end_layer_index
     • 根据y的结果去计算confidence和location， 并保存到结果中
     • header_index += 1

   将base_net[end_layer_index: ]中的层正常的正向传播

   针对extras:(这个就是相对mobilenet V2 额外添加的四个层用来做检测的)

   • x = layer(x)
   • 对每一个x计算confidence和location，将它放到结果中去
   • header_index += 1

   处理confidences和location，在第1维度做拼接

   如果不是测试阶段，直接返回confidences和location

   如果是测试阶段，还需要多做如下的处理：

   • 对confidences做softmax操作
   • 将location转成我们需要的boxes值
   • return confidences和boxes

3. compute_header操作

   处理过程是这样的

   针对特征x,

   分别用两个卷积去得到confidence和location

   confidence需要permute(0,2,3,1).contiguous再进行view，使得符合(img, -1, classes)的格式

   location需要permute(0,2,3,1).contiguous再进行view，使得符合(img,-1, 4)的格式