what.models.detection.ssd.ssd.squeezenet_ssd_lite_create

 1import torch
 2from torch.nn import Conv2d, Sequential, ModuleList, ReLU
 3from ..nn.squeezenet import squeezenet1_1
 4
 5from .ssd import SSD
 6from .predictor import Predictor
 7from . import squeezenet_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_squeezenet_ssd_lite(num_classes, is_test=False):
22    base_net = squeezenet1_1(False).features  # disable dropout layer
23
24    source_layer_indexes = [
25        12
26    ]
27    extras = ModuleList([
28        Sequential(
29            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
30            ReLU(),
31            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=2),
32        ),
33        Sequential(
34            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
35            ReLU(),
36            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
37        ),
38        Sequential(
39            Conv2d(in_channels=512, out_channels=128, kernel_size=1),
40            ReLU(),
41            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
42        ),
43        Sequential(
44            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
45            ReLU(),
46            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
47        ),
48        Sequential(
49            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
50            ReLU(),
51            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
52        )
53    ])
54
55    regression_headers = ModuleList([
56        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
57        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
58        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
59        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
60        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
61        Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
62    ])
63
64    classification_headers = ModuleList([
65        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
66        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
67        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
68        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
69        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
70        Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
71    ])
72
73    return SSD(num_classes, base_net, source_layer_indexes,
74               extras, classification_headers, regression_headers, is_test=is_test, config=config)
75
76
77def create_squeezenet_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
78    predictor = Predictor(net, config.image_size, config.image_mean,
79                          config.image_std,
80                          nms_method=nms_method,
81                          iou_threshold=config.iou_threshold,
82                          candidate_size=candidate_size,
83                          sigma=sigma,
84                          device=device)
85    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_squeezenet_ssd_lite(num_classes, is_test=False):
22def create_squeezenet_ssd_lite(num_classes, is_test=False):
23    base_net = squeezenet1_1(False).features  # disable dropout layer
24
25    source_layer_indexes = [
26        12
27    ]
28    extras = ModuleList([
29        Sequential(
30            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
31            ReLU(),
32            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=2),
33        ),
34        Sequential(
35            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
36            ReLU(),
37            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
38        ),
39        Sequential(
40            Conv2d(in_channels=512, out_channels=128, kernel_size=1),
41            ReLU(),
42            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
43        ),
44        Sequential(
45            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
46            ReLU(),
47            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
48        ),
49        Sequential(
50            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
51            ReLU(),
52            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
53        )
54    ])
55
56    regression_headers = ModuleList([
57        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
58        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
59        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
60        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
61        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
62        Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
63    ])
64
65    classification_headers = ModuleList([
66        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
67        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
68        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
69        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
70        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
71        Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
72    ])
73
74    return SSD(num_classes, base_net, source_layer_indexes,
75               extras, classification_headers, regression_headers, is_test=is_test, config=config)
def create_squeezenet_ssd_lite_predictor( net, candidate_size=200, nms_method=None, sigma=0.5, device=device(type='cpu')):
78def create_squeezenet_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
79    predictor = Predictor(net, config.image_size, config.image_mean,
80                          config.image_std,
81                          nms_method=nms_method,
82                          iou_threshold=config.iou_threshold,
83                          candidate_size=candidate_size,
84                          sigma=sigma,
85                          device=device)
86    return predictor