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