服务器使用指定的 GPU 训练神经网络(笔记)_torch.distributed.launch 指定gpu

简单记个笔记

当我们用实验室服务器的多个 GPU 训练神经网络并指定特定的 GPU 时可以使用以下命令:

1、torch 1.9.0 之前的版本(不包括 1.9.0)

不指定使用哪块 GPU

python -m torch.distributed.launch --nproc_per_node=4 --use_env train_multi_GPU.py

说明:

nproc_per_node 参数为使用 GPU 数量，

train_multi_GPU.py 是你要训练的 python 脚本文件。

指定使用特定的 GPU

CUDA_VISIBLE_DEVICES=0,2,4,6 python -m torch.distributed.launch --nproc_per_node=4 --use_env train_multi_GPU.py

说明:

CUDA_VISIBLE_DEVICES=0,2,4,6 指定使用这几个 GPU(例如我只要使用设备中的第 0, 2, 4, 6 块 GPU 设备)

2、torch 1.9.0 之后的版本(包括 1.9.0)

不指定使用哪块 GPU

torchrun --nproc_per_node=4 train_multi_GPU.py

指定使用特定的 GPU

CUDA_VISIBLE_DEVICES=0,2,4,6 torchrun --nproc_per_node=4 train_multi_GPU.py