1、shell 脚本中

  1. # jy: 设置可运行的显卡编号(编号即 nvidia-smi 看到的显卡编号);
  2. export CUDA_VISIBLE_DEVICES=2
  3. #export CUDA_VISIBLE_DEVICES=0,1
  6. # jy: Allow multiple threads(允许多线程)
  7. export OMP_NUM_THREADS=8
  10. # jy: Use distributed data parallel (使用分布式数据并行; 多卡处理, 每张卡占用内存并不
  11. #     会因为卡的增多而下降)
  12. # jy: 只使用一张 GPU 卡(如果 train.py 脚本逻辑中会使用 GPU, 则会自动使用以上设置的 GPU 卡号)
  13. python train.py \
  14.     --para_name xxx \
  15.     --para_name2 \
  16.     "$@"
  19. # jy: 使用多张 GPU 卡
  20. python -m torch.distributed.launch --nproc_per_node $NUM_GPU --master_port $PORT_ID train.py \
  21.     --para_name xxx \
  22.     --para_name2 \
  23.     "$@"

2、python 脚本中

import os

os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
#os.environ['CUDA_VISIBLE_DEVICES'] = "5,6"

3、命令行中指定

  • 单 GPU
    • CUDA_VISIBLE_DEVICES=2 python train.py
  • 多 GPU
    • CUDA_VISIBLE_DEVICES=2,3,4 python train.py