From 7ee480fc48ddf36570287a7109feb2e103201cab Mon Sep 17 00:00:00 2001
From: tangwei <tangwei12@baidu.com>
Date: Tue, 28 Apr 2020 14:51:10 +0800
Subject: [PATCH] code reformat

---
 fleetrec/run.py |  6 ++++--
 readme.md       | 13 +++++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/fleetrec/run.py b/fleetrec/run.py
index 641b2ce4..9a4645a7 100644
--- a/fleetrec/run.py
+++ b/fleetrec/run.py
@@ -164,10 +164,12 @@ def local_mpi_engine(args):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='fleet-rec run')
     parser.add_argument("-m", "--model", type=str)
-    parser.add_argument("-e", "--engine", type=str, choices=clusters)
-    parser.add_argument("-d", "--device", type=str, choices=["CPU", "GPU"], default="CPU")
+    parser.add_argument("-e", "--engine", type=str, choices=["single", "local_cluster", "cluster"])
+    parser.add_argument("-d", "--device", type=str, choices=["cpu", "gpu"], default="cpu")
 
     args = parser.parse_args()
+    args.engine = args.engine.upper()
+    args.device = args.device.upper()
 
     if not os.path.isfile(args.model):
         raise FileNotFoundError("argument model: {} do not exist".format(args.model))
diff --git a/readme.md b/readme.md
index 5beb6ea6..1565b367 100644
--- a/readme.md
+++ b/readme.md
@@ -28,17 +28,25 @@ cd FleetRec
 
 python -m fleetrec.run \
        -m fleetrec/examples/ctr-dnn_train.yaml \
+       -d cpu \
        -e single 
+
+# 使用GPU资源进行训练
+python -m fleetrec.run \
+       -m fleetrec/examples/ctr-dnn_train.yaml \
+       -d gpu \
+       -e single
 ```
 
 ### 本地模拟分布式训练
 
 ```bash
 cd FleetRec
-
+# 使用CPU资源进行训练
 python -m fleetrec.run \
        -m fleetrec/examples/ctr-dnn_train.yaml \
-       -e local_cluster 
+       -d cpu \
+       -e local_cluster
 ```
 
 ### 集群提交分布式训练<需要用户预先配置好集群环境，本提交命令不包含提交客户端>
@@ -48,6 +56,7 @@ cd FleetRec
 
 python -m fleetrec.run \
        -m fleetrec/examples/ctr-dnn_train.yaml \
+       -d cpu \
        -e cluster
 ```
 
-- 
GitLab