what.models.detection.ssd.ssd.mobilenet_v2_ssd_lite_create

 1import torch
 2from torch import nn
 3from torch.nn import Conv2d, Sequential, ModuleList, BatchNorm2d
 4
 5
 6from what.models.detection.ssd.ssd import SSD, GraphPath
 7from what.models.detection.ssd.ssd import mobilenet_ssd_config as config
 8from what.models.detection.ssd.ssd.predictor import Predictor
 9from what.models.detection.ssd.nn.mobilenet_v2 import MobileNetV2, InvertedResidual
10
11
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    )
23
24
25def create_mobilenet_v2_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
26    base_net = MobileNetV2(width_mult=width_mult, use_batch_norm=use_batch_norm,
27                           onnx_compatible=onnx_compatible).features
28
29    source_layer_indexes = [
30        GraphPath(14, 'conv', 3),
31        19,
32    ]
33    extras = ModuleList([
34        InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
35        InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
36        InvertedResidual(256, 256, stride=2, expand_ratio=0.5),
37        InvertedResidual(256, 64, stride=2, expand_ratio=0.25)
38    ])
39
40    regression_headers = ModuleList([
41        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * 4,
42                        kernel_size=3, padding=1, onnx_compatible=False),
43        SeperableConv2d(in_channels=1280, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
44        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
45        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
46        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
47        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
48    ])
49
50    classification_headers = ModuleList([
51        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
52        SeperableConv2d(in_channels=1280, out_channels=6 * num_classes, kernel_size=3, padding=1),
53        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
54        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
55        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
56        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
57    ])
58
59    return SSD(num_classes, base_net, source_layer_indexes,
60               extras, classification_headers, regression_headers, is_test=is_test, config=config)
61
62
63def create_mobilenet_v2_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
64    predictor = Predictor(net, config.image_size, config.image_mean,
65                          config.image_std,
66                          nms_method=nms_method,
67                          iou_threshold=config.iou_threshold,
68                          candidate_size=candidate_size,
69                          sigma=sigma,
70                          device=device)
71    return predictor
def SeperableConv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
13def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
14    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
15    """
16    ReLU = nn.ReLU if onnx_compatible else nn.ReLU6
17    return Sequential(
18        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
19               groups=in_channels, stride=stride, padding=padding),
20        BatchNorm2d(in_channels),
21        ReLU(),
22        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
23    )

Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.

def create_mobilenet_v2_ssd_lite( num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
26def create_mobilenet_v2_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True, onnx_compatible=False, is_test=False):
27    base_net = MobileNetV2(width_mult=width_mult, use_batch_norm=use_batch_norm,
28                           onnx_compatible=onnx_compatible).features
29
30    source_layer_indexes = [
31        GraphPath(14, 'conv', 3),
32        19,
33    ]
34    extras = ModuleList([
35        InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
36        InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
37        InvertedResidual(256, 256, stride=2, expand_ratio=0.5),
38        InvertedResidual(256, 64, stride=2, expand_ratio=0.25)
39    ])
40
41    regression_headers = ModuleList([
42        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * 4,
43                        kernel_size=3, padding=1, onnx_compatible=False),
44        SeperableConv2d(in_channels=1280, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
45        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
46        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
47        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
48        Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
49    ])
50
51    classification_headers = ModuleList([
52        SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * num_classes, kernel_size=3, padding=1),
53        SeperableConv2d(in_channels=1280, out_channels=6 * num_classes, kernel_size=3, padding=1),
54        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
55        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
56        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
57        Conv2d(in_channels=64, out_channels=6 * num_classes, kernel_size=1),
58    ])
59
60    return SSD(num_classes, base_net, source_layer_indexes,
61               extras, classification_headers, regression_headers, is_test=is_test, config=config)
def create_mobilenet_v2_ssd_lite_predictor( net, candidate_size=200, nms_method=None, sigma=0.5, device=device(type='cpu')):
64def create_mobilenet_v2_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
65    predictor = Predictor(net, config.image_size, config.image_mean,
66                          config.image_std,
67                          nms_method=nms_method,
68                          iou_threshold=config.iou_threshold,
69                          candidate_size=candidate_size,
70                          sigma=sigma,
71                          device=device)
72    return predictor