!9768 mode_gpu_export

From: @bai-yangfan
Reviewed-by: @zhunaipan
Signed-off-by:

model_zoo/official/cv/mobilenetv2_quant/export.py

@@ -22,7 +22,7 @@ from mindspore import Tensor, context, load_checkpoint, load_param_into_net, exp
 from mindspore.compression.quant import QuantizationAwareTraining
 
 from src.mobilenetV2 import mobilenetV2
-from src.config import config_ascend_quant
+from src.config import config_quant
 
 parser = argparse.ArgumentParser(description='Image classification')
 parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
@@ -30,13 +30,8 @@ parser.add_argument('--device_target', type=str, default=None, help='Run device
 args_opt = parser.parse_args()
 
 if __name__ == '__main__':
-    cfg = None
-    if args_opt.device_target == "Ascend":
-        cfg = config_ascend_quant
-        context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=False)
-    else:
-        raise ValueError("Unsupported device target: {}.".format(args_opt.device_target))
-
+    cfg = config_quant(args_opt.device_target)
+    context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target, save_graphs=False)
     # define fusion network
     network = mobilenetV2(num_classes=cfg.num_classes)
     # convert fusion network to quantization aware network

model_zoo/official/cv/mobilenetv2_quant/src/config.py

@@ -38,6 +38,8 @@ config_ascend_quant = ed({
 
 config_gpu_quant = ed({
     "num_classes": 1000,
+    "image_height": 224,
+    "image_width": 224,
     "batch_size": 300,
     "epoch_size": 60,
     "start_epoch": 200,
@@ -52,3 +54,14 @@ config_gpu_quant = ed({
     "keep_checkpoint_max": 300,
     "save_checkpoint_path": "./checkpoint",
 })
+
+def config_quant(device_target):
+    if device_target not in ["Ascend", "GPU"]:
+        raise ValueError("Unsupported device target: {}.".format(device_target))
+    configs = ed({
+        "Ascend": config_ascend_quant,
+        "GPU": config_gpu_quant
+    })
+    config = configs.Ascend if device_target == "Ascend" else configs.GPU
+    config["device_target"] = device_target
+    return config