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