what.models.detection.ssd.ssd.mobilenet_v3_ssd_lite_create

  1import torch
  2from torch.nn import Conv2d, Sequential, ModuleList, BatchNorm2d
  3from torch import nn
  4from ..nn.mobilenet_v3 import MobileNetV3_Large, MobileNetV3_Small, Block, hswish
  5
  6from .ssd import SSD
  7from .predictor import Predictor
  8from . import mobilenet_ssd_config as config
  9
 10
 11def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
 12    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
 13    """
 14    ReLU = nn.ReLU if onnx_compatible else nn.ReLU6
 15    return Sequential(
 16        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
 17               groups=in_channels, stride=stride, padding=padding),
 18        BatchNorm2d(in_channels),
 19        ReLU(),
 20        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
 21    )
 22
 23
 24def create_mobilenetv3_large_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
 25    base_net = MobileNetV3_Large().features
 26
 27    source_layer_indexes = [ 15, 21 ]
 28    extras = ModuleList([
 29        Block(3, 960, 256, 512, hswish(), None, stride=2),
 30        Block(3, 512, 128, 256, hswish(), None, stride=2),
 31        Block(3, 256, 128, 256, hswish(), None, stride=2),
 32        Block(3, 256, 64, 64, hswish(), None, stride=2)
 33    ])
 34
 35    regression_headers = ModuleList([
 36        SeperableConv2d(in_channels=round(112 * width_mult), out_channels=6 * 4,
 37                        kernel_size=3, padding=1, onnx_compatible=False),
 38        SeperableConv2d(in_channels=960, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 39        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 40        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 41        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 42        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
 43    ])
 44
 45    classification_headers = ModuleList([
 46        SeperableConv2d(in_channels=round(112 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
 47        SeperableConv2d(in_channels=960, out_channels=6 * num_classes, kernel_size=3, padding=1),
 48        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
 49        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
 50        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
 51        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
 52    ])
 53
 54    return SSD(num_classes, base_net, source_layer_indexes,
 55               extras, classification_headers, regression_headers, is_test=is_test, config=config)
 56
 57
 58def create_mobilenetv3_small_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
 59    base_net = MobileNetV3_Small().features
 60
 61    source_layer_indexes = [ 11, 17 ]
 62    extras = ModuleList([
 63        Block(3, 576, 256, 512, hswish(), None, stride=2),
 64        Block(3, 512, 128, 256, hswish(), None, stride=2),
 65        Block(3, 256, 128, 256, hswish(), None, stride=2),
 66        Block(3, 256, 64, 64, hswish(), None, stride=2)
 67    ])
 68
 69    regression_headers = ModuleList([
 70        SeperableConv2d(in_channels=round(48 * width_mult), out_channels=6 * 4,
 71                        kernel_size=3, padding=1, onnx_compatible=False),
 72        SeperableConv2d(in_channels=576, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 73        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 74        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 75        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
 76        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
 77    ])
 78
 79    classification_headers = ModuleList([
 80        SeperableConv2d(in_channels=round(48 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
 81        SeperableConv2d(in_channels=576, out_channels=6 * num_classes, kernel_size=3, padding=1),
 82        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
 83        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
 84        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
 85        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
 86    ])
 87
 88    return SSD(num_classes, base_net, source_layer_indexes,
 89               extras, classification_headers, regression_headers, is_test=is_test, config=config)
 90
 91
 92def create_mobilenetv3_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
 93    predictor = Predictor(net, config.image_size, config.image_mean,
 94                          config.image_std,
 95                          nms_method=nms_method,
 96                          iou_threshold=config.iou_threshold,
 97                          candidate_size=candidate_size,
 98                          sigma=sigma,
 99                          device=device)
100    return predictor
def SeperableConv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
12def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
13    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
14    """
15    ReLU = nn.ReLU if onnx_compatible else nn.ReLU6
16    return Sequential(
17        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
18               groups=in_channels, stride=stride, padding=padding),
19        BatchNorm2d(in_channels),
20        ReLU(),
21        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
22    )

Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.

def create_mobilenetv3_large_ssd_lite( num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
25def create_mobilenetv3_large_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
26    base_net = MobileNetV3_Large().features
27
28    source_layer_indexes = [ 15, 21 ]
29    extras = ModuleList([
30        Block(3, 960, 256, 512, hswish(), None, stride=2),
31        Block(3, 512, 128, 256, hswish(), None, stride=2),
32        Block(3, 256, 128, 256, hswish(), None, stride=2),
33        Block(3, 256, 64, 64, hswish(), None, stride=2)
34    ])
35
36    regression_headers = ModuleList([
37        SeperableConv2d(in_channels=round(112 * width_mult), out_channels=6 * 4,
38                        kernel_size=3, padding=1, onnx_compatible=False),
39        SeperableConv2d(in_channels=960, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
40        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
41        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
42        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
43        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
44    ])
45
46    classification_headers = ModuleList([
47        SeperableConv2d(in_channels=round(112 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
48        SeperableConv2d(in_channels=960, out_channels=6 * num_classes, kernel_size=3, padding=1),
49        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
50        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
51        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
52        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
53    ])
54
55    return SSD(num_classes, base_net, source_layer_indexes,
56               extras, classification_headers, regression_headers, is_test=is_test, config=config)
def create_mobilenetv3_small_ssd_lite( num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
59def create_mobilenetv3_small_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
60    base_net = MobileNetV3_Small().features
61
62    source_layer_indexes = [ 11, 17 ]
63    extras = ModuleList([
64        Block(3, 576, 256, 512, hswish(), None, stride=2),
65        Block(3, 512, 128, 256, hswish(), None, stride=2),
66        Block(3, 256, 128, 256, hswish(), None, stride=2),
67        Block(3, 256, 64, 64, hswish(), None, stride=2)
68    ])
69
70    regression_headers = ModuleList([
71        SeperableConv2d(in_channels=round(48 * width_mult), out_channels=6 * 4,
72                        kernel_size=3, padding=1, onnx_compatible=False),
73        SeperableConv2d(in_channels=576, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
74        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
75        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
76        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
77        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
78    ])
79
80    classification_headers = ModuleList([
81        SeperableConv2d(in_channels=round(48 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
82        SeperableConv2d(in_channels=576, out_channels=6 * num_classes, kernel_size=3, padding=1),
83        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
84        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
85        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
86        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
87    ])
88
89    return SSD(num_classes, base_net, source_layer_indexes,
90               extras, classification_headers, regression_headers, is_test=is_test, config=config)
def create_mobilenetv3_ssd_lite_predictor( net, candidate_size=200, nms_method=None, sigma=0.5, device=device(type='cpu')):
 93def create_mobilenetv3_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
 94    predictor = Predictor(net, config.image_size, config.image_mean,
 95                          config.image_std,
 96                          nms_method=nms_method,
 97                          iou_threshold=config.iou_threshold,
 98                          candidate_size=candidate_size,
 99                          sigma=sigma,
100                          device=device)
101    return predictor