pytorch调用多GPU及GPU选择
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
多机多卡任务还未试验过，暂不书写，详细方法可以参考
https://zhuanlan.zhihu.com/p/86441879 博主写的很详细。
上一篇：
通过多线程提高代码的执行效率例子