diff --git a/official/cv/mobilenetv3/scripts/run_train.sh b/official/cv/mobilenetv3/scripts/run_train.sh index 075b3721cac91d742f62144da52936ccaff3a2d4..11126508a48ab568292e00ec91310b3e43683e47 100644 --- a/official/cv/mobilenetv3/scripts/run_train.sh +++ b/official/cv/mobilenetv3/scripts/run_train.sh @@ -46,11 +46,20 @@ run_gpu() cd ../train || exit export CUDA_VISIBLE_DEVICES="$3" - mpirun -n $2 --allow-run-as-root --output-filename log_output --merge-stderr-to-stdout \ - python ${BASEPATH}/../train.py \ + if [ $2 -eq 1 ] ; then + python ${BASEPATH}/../train.py \ --dataset_path=$4 \ --device_target=$1 \ - &> ../train.log & # dataset train folder + --run_distribute=False \ + &> ../train.log & + else + mpirun -n $2 --allow-run-as-root --output-filename log_output --merge-stderr-to-stdout \ + python ${BASEPATH}/../train.py \ + --dataset_path=$4 \ + --device_target=$1 \ + --run_distribute=True \ + &> ../train.log & # dataset train folder + fi; } run_cpu() diff --git a/official/cv/mobilenetv3/train.py b/official/cv/mobilenetv3/train.py index 01b2ebdbe26b90186373e4d89e1ed873adfb2359..9d7f14df4886fcd3a991efa2211212fbec4e0334 100644 --- a/official/cv/mobilenetv3/train.py +++ b/official/cv/mobilenetv3/train.py @@ -49,7 +49,7 @@ parser = argparse.ArgumentParser(description='Image classification') parser.add_argument('--dataset_path', type=str, default=None, help='Dataset path') parser.add_argument('--pre_trained', type=str, default=None, help='Pretrained checkpoint path') parser.add_argument('--device_target', type=str, default="GPU", help='run device_target') -parser.add_argument('--run_distribute', type=ast.literal_eval, default=True, help='Run distribute') +parser.add_argument('--run_distribute', type=ast.literal_eval, default=False, help='Run distribute') args_opt = parser.parse_args() if args_opt.device_target == "GPU":