what.models.detection.ssd.ssd.mobilenet_v1_ssd_create API documentation

# what.models.detection.ssd.ssd.mobilenet_v1_ssd_create

View Source

     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): View Source

     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): View Source

    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