what.models.detection.ssd.ssd.mobilenet_v2_ssd_lite_create API documentation

# what.models.detection.ssd.ssd.mobilenet_v2_ssd_lite_create

View Source

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

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

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

    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