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)

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为空

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
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)的格式

MobileNetV2-SSDLite代码分析-4 ssd

你可能感兴趣的:(MobileNetV2-SSDLite代码分析-4 ssd)