pytorch调用多GPU及GPU选择

最新推荐文章于 2024-05-20 05:57:09 发布

giganticpower

最新推荐文章于 2024-05-20 05:57:09 发布

阅读量5.5k

点赞数

文章标签：深度学习神经网络 pytorch 机器学习 python

本文链接：https://blog.csdn.net/giganticpower/article/details/113636023

版权

pytorch多卡调用与选择

第一步，查看gpu环境：
执行：nvidia-smi
第二步，gpu选择：
执行：os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
将gpu环境设定为0卡和3卡。
使用显卡训练有两个部分需要进行挂gpu。
1数据：直接

inputs = inputs.cuda() 
labels = labels.cuda()

在模型上并行我们需要：

model = nn.DataParallel(model) # or torch.nn.DataParallel(net,device_ids=[0,1]).cuda()
model = model.cuda()

平衡DataParallel带来的显存使用不平衡：
使用DistributedDataParallel来代替DataParallel
初始化：

torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)

分发：

torch.distributed.init_process_group(backend="nccl")
model = DistributedDataParallel(model) # device_ids will include all GPU devices by default

但是这里需要注意的是, 如果使用这句代码, 直接在pycharm或者别的编辑器中,是没法正常运行的, 因为这个需要在shell的命令行中运行, 如果想要正确执行这段代码, 假设这段代码的名字是main.py, 可以使用如下的方法进行:

python -m torch.distributed.launch main.py

多机多卡任务还未试验过，暂不书写，详细方法可以参考多GPU使用

https://zhuanlan.zhihu.com/p/86441879
博主写的很详细。

关注