单机多GPU训练

最近跑程序老被killed 掉,然后就考虑使用两个GPU,但是,我不知道该怎么将单GPU改成多GPU?解决办法如下:

  可以在终端直接用以下命令:

CUDA_VISIBLE_DEVICES=0,1  python -u -m torch.distributed.launch --nproc_per_node 2 --nnodes 1 run_rank.py

CUDA_VISIBLE_DEVICES=0,1 : 定义GPU设备。根据gpu的数量来设定,初始gpu为0,这里我的gpu数量为2

torch.distributed.launch:启动方式

nproc_per_node :表示GPU数量(我这里用了两个GPU)

nnodes : 由于是用单机多卡训练,所以设置为1

你可能感兴趣的:(python,深度学习,人工智能,python)