solve Warning

convert_to_caffe2_models.py

@@ -20,12 +20,15 @@
 
 class_names = [name.strip() for name in open(label_path).readlines()]
 num_classes = len(class_names)
+IMAGE_SIZE=300
 
 if net_type == 'vgg16-ssd':
     net = create_vgg_ssd(len(class_names), is_test=True)
 elif net_type == 'mb1-ssd':
     net = create_mobilenetv1_ssd(len(class_names), is_test=True)
 elif net_type == 'mb1-ssd-lite':
+    from vision.ssd.config.mobilenetv1_ssd_config import image_size
+    IMAGE_SIZE=image_size
     net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
 elif net_type == 'mb2-ssd-lite':
     net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
@@ -34,7 +37,11 @@
 else:
     print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
     sys.exit(1)
-net.load(model_path)
+
+device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+net.load_state_dict(torch.load(model_path, map_location=device))
+net.to(device)
 net.eval()
 
 model_path = f"models/{net_type}.onnx"
@@ -43,7 +50,7 @@
 predict_net_path = f"models/{net_type}_predict_net.pb"
 predict_net_txt_path = f"models/{net_type}_predict_net.pbtxt"
 
-dummy_input = torch.randn(1, 3, 300, 300)
+dummy_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE,dtype=torch.float32, device=device.type)
 torch.onnx.export(net, dummy_input, model_path, verbose=False, output_names=['scores', 'boxes'])
 
 model = onnx.load(model_path)

train_ssd.py

@@ -3,7 +3,7 @@
 import logging
 import sys
 import itertools
-
+from torchsummary import summary
 import torch
 from torch.utils.data import DataLoader, ConcatDataset
 from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
@@ -37,7 +37,7 @@
 
 
 parser.add_argument('--net', default="vgg16-ssd",
-                    help="The network architecture, it can be mb1-ssd, mb1-lite-ssd, mb2-ssd-lite, mb3-large-ssd-lite, mb3-small-ssd-lite or vgg16-ssd.")
+                    help="The network architecture, it can be mb1-ssd, mb1-ssd-lite, mb2-ssd-lite, mb3-large-ssd-lite, mb3-small-ssd-lite or vgg16-ssd.")
 parser.add_argument('--freeze_base_net', action='store_true',
                     help="Freeze base net layers.")
 parser.add_argument('--freeze_net', action='store_true',
@@ -246,6 +246,7 @@ def test(loader, net, criterion, device):
                             shuffle=False)
     logging.info("Build network.")
     net = create_net(num_classes)
+    summary(net.to("cuda"),(3,config.image_size,config.image_size),device='cuda')
     min_loss = -10000.0
     last_epoch = -1
 

vision/ssd/mobilenet_v2_ssd_lite.py

@@ -39,10 +39,10 @@ def create_mobilenetv2_ssd_lite(num_classes, width_mult=1.0, use_batch_norm=True
     regression_headers = ModuleList([
         SeperableConv2d(in_channels=round(576 * width_mult), out_channels=6 * 4,
                         kernel_size=3, padding=1, onnx_compatible=False),
-        SeperableConv2d(in_channels=1280, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
-        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
-        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
-        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=False),
+        SeperableConv2d(in_channels=1280, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=onnx_compatible),
+        SeperableConv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=onnx_compatible),
+        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=onnx_compatible),
+        SeperableConv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1, onnx_compatible=onnx_compatible),
         Conv2d(in_channels=64, out_channels=6 * 4, kernel_size=1),
     ])
 

vision/transforms/transforms.py

@@ -244,7 +244,7 @@ def __call__(self, image, boxes=None, labels=None):
         height, width, _ = image.shape
         while True:
             # randomly choose a mode
-            mode = random.choice(self.sample_options)
+            mode = random.choice(np.array(self.sample_options,dtype=object))
             if mode is None:
                 return image, boxes, labels