diff --git a/fluid/recommendation/gru4rec/README.md b/fluid/recommendation/gru4rec/README.md index 81137405052f08a625ed73e35d30846330c86e7e..d4489bfbbfb92f15a43d0384bf24b353d107c6fc 100644 --- a/fluid/recommendation/gru4rec/README.md +++ b/fluid/recommendation/gru4rec/README.md @@ -74,17 +74,16 @@ python convert_format.py ``` ## 训练 ---use_cuda 表示使用gpu --parallel 表示多卡 +'--use_cuda 1' 表示使用gpu --parallel 1 表示多卡 GPU 环境 -运行命令 `CUDA_VISIBLE_DEVICES=0 python train.py train_file test_file --use_cuda` 开始训练模型。 - -```python -CUDA_VISIBLE_DEVICES=0 python train.py small_train.txt small_test.txt --use_cuda +运行命令 `CUDA_VISIBLE_DEVICES=0 python train.py train_file test_file --use_cuda 1` 开始训练模型。 +``` +CUDA_VISIBLE_DEVICES=0 python train.py small_train.txt small_test.txt --use_cuda 1 ``` CPU 环境 运行命令 `python train.py train_file test_file` 开始训练模型。 -```python +``` python train.py small_train.txt small_test.txt ``` diff --git a/fluid/recommendation/gru4rec/train.py b/fluid/recommendation/gru4rec/train.py index 5474c82800373e52dc4a4537ea978972c0d4edbb..12bc1bc20ed93e1376e97ef4c8e76af51a529cc8 100644 --- a/fluid/recommendation/gru4rec/train.py +++ b/fluid/recommendation/gru4rec/train.py @@ -185,6 +185,7 @@ def train_net(): test_file = args.test_file use_cuda = True if args.use_cuda else False parallel = True if args.parallel else False + print("use_cuda:", use_cuda, "parallel:", parallel) batch_size = 50 vocab, train_reader, test_reader = utils.prepare_data( train_file, test_file,batch_size=batch_size * get_cards(args),\