不乱于心,不困于情。
不畏将来,不念过往。如此,安好。

torch如何指定多GPU训练

训练的时候可以使用 .to(device)和.cuda(),注意两者的区别。

一、.to(device),可以指定使用CPU或GPU
#使用单GPU或者CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#也可以手动指定使用GPU或CPU,如下:
# device=torch.device("CPU")
# device=torch.device("GPU")
model.to(device)#将模型加载到相应的设备中

#使用多GPU
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model,device_ids=[0,1,2])
model.to(device)
#【注】:这里必须使用0块显卡,不然会报错

#如果使用所有的GPU,那么:
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.to(device)

二、cuda(),只能指定GPU
#指定单个GPU
import os
os.environ['CUDA_VISIBLE_DEVICES']='2'
model.cuda()

#指定多个GPU
from torch.nn import DataParallel
#上面这一行一定要加,因为没加一直报错都开始怀疑torch版本是不是有问题了哈哈
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,2,3,4'#代表使用的第一块,第三块,第四块和第五块卡
device_ids = [0,2,3,4]
model = torch.nn.Dataparallel(model, device_ids =device_ids)
model = model.cuda()

#如果默认使用所有的GPU,那么:
model = torch.nn.Dataparallel(model )

#也可以在终端指定GPU
CUDA_VISIBLE_DEVICES=0,2,3,4 python ***.py

三、分布式训练torch.distributed.launch
-m torch.distributed.launch --nproc_per_node=NGPUS --master_port=$PORT
#NGPUS代表使用了几块GPU卡
#上面那条语句是按顺序选择卡,如果想要指定固定的卡,可以在跑程序时先在终端指定卡号然后再输入上面的语句:
#终端指定卡号:
export CUDA_VISIBLE_DEVICES=0,2

赞(0) 打赏
未经允许不得转载:seo优化_前端开发_渗透技术 » torch如何指定多GPU训练

相关推荐

  • 暂无文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

微信扫一扫打赏