Create backbones.py

model/backbones.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+import torchvision.models as models
+
+
+SUPPORTED_BACKBONES = [
+    'mobilenetv2'
+]
+
+
+def build_backbone(backbone_name):
+    if backbone_name == 'mobilenetv2':
+        backbone = models.mobilenet_v2(pretrained=True).features
+        backbone.out_channels = [24, 32, 96, 320]
+        return backbone
+    else:
+        raise NotImplementedError('Backbone only supports mobilenetv2 as now')