what.models.detection.ssd.ssd.mobilenet_v1_ssd_create

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