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
-
GraphPath(s0=14, name='conv', s1=3)
path = GraphPath(s0=14, name='conv', s1=3)，对应着我想在basenet的第14层找到名为conv的层的集合， 该集合中的第3层我想用来做分类与回归
end_layer_index = 14
-
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)的格式
