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))