what.models.detection.ssd.ssd.mobilenet_v1_ssd_lite_create

 1import torch
 2from torch.nn import Conv2d, Sequential, ModuleList, ReLU, BatchNorm2d
 3from ..nn.mobilenet_v1 import MobileNetV1
 4
 5from .ssd import SSD
 6from .predictor import Predictor
 7from . import mobilenet_ssd_config as config
 8
 9
10def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
11    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
12    """
13    return Sequential(
14        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
15               groups=in_channels, stride=stride, padding=padding),
16        ReLU(),
17        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
18    )
19
20
21def create_mobilenetv1_ssd_lite(num_classes, is_test=False):
22    base_net = MobileNetV1(1001).model  # disable dropout layer
23
24    source_layer_indexes = [
25        12,
26        14,
27    ]
28
29    extras = ModuleList([
30        Sequential(
31            Conv2d(in_channels=1024, out_channels=256, kernel_size=1),
32            ReLU(),
33            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
34        ),
35        Sequential(
36            Conv2d(in_channels=512, out_channels=128, kernel_size=1),
37            ReLU(),
38            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
39        ),
40        Sequential(
41            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
42            ReLU(),
43            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
44        ),
45        Sequential(
46            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
47            ReLU(),
48            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
49        )
50    ])
51
52    regression_headers = ModuleList([
53        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
54        SeperableConv2d(in_channels=1024, out_channels=6 * 4, kernel_size=3, padding=1),
55        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
56        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
57        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
58        Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
59    ])
60
61    classification_headers = ModuleList([
62        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
63        SeperableConv2d(in_channels=1024, out_channels=6 * num_classes, kernel_size=3, padding=1),
64        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
65        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
66        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
67        Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
68    ])
69
70    return SSD(num_classes, base_net, source_layer_indexes,
71               extras, classification_headers, regression_headers, is_test=is_test, config=config)
72
73
74def create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=None):
75    predictor = Predictor(net, config.image_size, config.image_mean,
76                          config.image_std,
77                          nms_method=nms_method,
78                          iou_threshold=config.iou_threshold,
79                          candidate_size=candidate_size,
80                          sigma=sigma,
81                          device=device)
82    return predictor
def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
11def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
12    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
13    """
14    return Sequential(
15        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
16               groups=in_channels, stride=stride, padding=padding),
17        ReLU(),
18        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
19    )

Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.

def create_mobilenetv1_ssd_lite(num_classes, is_test=False):
22def create_mobilenetv1_ssd_lite(num_classes, is_test=False):
23    base_net = MobileNetV1(1001).model  # disable dropout layer
24
25    source_layer_indexes = [
26        12,
27        14,
28    ]
29
30    extras = ModuleList([
31        Sequential(
32            Conv2d(in_channels=1024, out_channels=256, kernel_size=1),
33            ReLU(),
34            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
35        ),
36        Sequential(
37            Conv2d(in_channels=512, out_channels=128, kernel_size=1),
38            ReLU(),
39            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
40        ),
41        Sequential(
42            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
43            ReLU(),
44            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
45        ),
46        Sequential(
47            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
48            ReLU(),
49            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
50        )
51    ])
52
53    regression_headers = ModuleList([
54        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
55        SeperableConv2d(in_channels=1024, out_channels=6 * 4, kernel_size=3, padding=1),
56        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
57        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
58        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
59        Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
60    ])
61
62    classification_headers = ModuleList([
63        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
64        SeperableConv2d(in_channels=1024, out_channels=6 * num_classes, kernel_size=3, padding=1),
65        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
66        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
67        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
68        Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
69    ])
70
71    return SSD(num_classes, base_net, source_layer_indexes,
72               extras, classification_headers, regression_headers, is_test=is_test, config=config)
def create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=None):
75def create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=None):
76    predictor = Predictor(net, config.image_size, config.image_mean,
77                          config.image_std,
78                          nms_method=nms_method,
79                          iou_threshold=config.iou_threshold,
80                          candidate_size=candidate_size,
81                          sigma=sigma,
82                          device=device)
83    return predictor