选择 GPU / CPU

选择 GPU / CPU#

要在特定的 GPU 上运行程序,您可以使用 CUDA_VISIBLE_DEVICES 环境变量。例如,要在第二个 GPU 上运行程序,您可以使用:

CUDA_VISIBLE_DEVICES=1 python my_program.py

要在多个GPU上跑你的程序,你可以用:

CUDA_VISIBLE_DEVICES=0,1 python my_program.py

如果要禁用GPU(即使用CPU),你可以用:

CUDA_VISIBLE_DEVICES="" python my_program.py