From 8e9505d613d2fa6bc87ed73e1e0f0c8e2b1a1843 Mon Sep 17 00:00:00 2001 From: Kaustubh Pandey Date: Fri, 5 Jul 2019 11:48:35 +0530 Subject: [PATCH] Added mb2-ssd-lite Fixed missing mb2-ssd-lite --- convert_to_caffe2_models.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/convert_to_caffe2_models.py b/convert_to_caffe2_models.py index a9ed34ea..25d978d2 100644 --- a/convert_to_caffe2_models.py +++ b/convert_to_caffe2_models.py @@ -2,6 +2,7 @@ from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite +from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite import sys import torch.onnx @@ -26,6 +27,8 @@ net = create_mobilenetv1_ssd(len(class_names), is_test=True) elif net_type == 'mb1-ssd-lite': 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) elif net_type == 'sq-ssd-lite': net = create_squeezenet_ssd_lite(len(class_names), is_test=True) else: @@ -58,4 +61,4 @@ f.write(str(init_net)) with open(predict_net_txt_path, 'w') as f: - f.write(str(predict_net)) \ No newline at end of file + f.write(str(predict_net))