add squeezenet ssd lite

train_ssd.py

@@ -13,7 +13,7 @@
 from vision.ssd.vgg_ssd import create_vgg_ssd
 from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
 from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
-from vision.ssd.fpn_mobilenetv1_ssd import create_fpn_mobilenetv1_ssd
+from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
 from vision.datasets.voc_dataset import VOCDataset
 from vision.datasets.open_images import OpenImagesDataset
 from vision.nn.multibox_loss import MultiboxLoss
@@ -92,10 +92,14 @@
                     help='Directory for saving checkpoint models')
 
 
+logging.basicConfig(stream=sys.stdout, level=logging.INFO,
+                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 args = parser.parse_args()
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
+
 if args.use_cuda and torch.cuda.is_available():
     torch.backends.cudnn.benchmark = True
+    logging.info("Use Cuda.")
 
 
 def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1):
@@ -159,9 +163,6 @@ def test(loader, net, criterion, device):
 
 
 if __name__ == '__main__':
-    logging.basicConfig(stream=sys.stdout, level=logging.INFO,
-                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-
     timer = Timer()
 
     logging.info(args)
@@ -174,6 +175,9 @@ def test(loader, net, criterion, device):
     elif args.net == 'mb1-ssd-lite':
         create_net = create_mobilenetv1_ssd_lite
         config = mobilenetv1_ssd_config
+    elif args.net == 'sq-ssd-lite':
+        create_net = create_squeezenet_ssd_lite
+        config = mobilenetv1_ssd_config
     else:
         logging.fatal("The net type is wrong.")
         parser.print_help(sys.stderr)

vision/ssd/squeezenet_ssd_lite.py

@@ -0,0 +1,85 @@
+import torch
+from torch.nn import Conv2d, Sequential, ModuleList, ReLU, BatchNorm2d
+from torchvision import models
+
+from .ssd import SSD
+from .predictor import Predictor
+from .config import mobilenetv1_ssd_config as config
+
+
+def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
+    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
+    """
+    return Sequential(
+        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
+               groups=in_channels, stride=stride, padding=padding),
+        ReLU(),
+        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
+    )
+
+
+def create_squeezenet_ssd_lite(num_classes, is_test=False):
+    base_net = models.squeezenet1_1(False).features  # disable dropout layer
+
+    source_layer_indexes = [
+        12
+    ]
+    extras = ModuleList([
+        Sequential(
+            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
+            ReLU(),
+            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
+        ),
+        Sequential(
+            Conv2d(in_channels=512, out_channels=256, kernel_size=1),
+            ReLU(),
+            SeperableConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
+        ),
+        Sequential(
+            Conv2d(in_channels=512, out_channels=128, kernel_size=1),
+            ReLU(),
+            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
+        ),
+        Sequential(
+            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
+            ReLU(),
+            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
+        ),
+        Sequential(
+            Conv2d(in_channels=256, out_channels=128, kernel_size=1),
+            ReLU(),
+            SeperableConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
+        )
+    ])
+
+    regression_headers = ModuleList([
+        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
+        Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
+    ])
+
+    classification_headers = ModuleList([
+        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=512, out_channels=6 * num_classes, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
+        SeperableConv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=3, padding=1),
+        Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
+    ])
+
+    return SSD(num_classes, base_net, source_layer_indexes,
+               extras, classification_headers, regression_headers, is_test=is_test, config=config)
+
+
+def create_squeezenet_ssd_lite_predictor(net, candidate_size=200, nms_method=None, sigma=0.5, device=torch.device('cpu')):
+    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