- 单机多卡:只需运行一份代码,由该代码分配该台机器上GPU资源的使用
- 多机多卡:每台机器上都需要运行一份代码,机器之间需要互相通信传递梯度,并且模型参数的更新也存在同步训练模式和异步训练模式的区别
一、单机多gpu运行原理
1.1 单机多卡并行机制:
首先是将模型加载到主GPU,然后将模型复制到其他指定的GPU,接着把输入数据按batch维度划分给各个GPU去计算(batch_size_per_gpu=total_batch_size / num_gpu),各个GPU上的模型和数据独立进行前向推理计算,得到结果(loss)后汇总到主gpu上反向传播更新模型权重,再将主GPU模型更新复制到其他GPU上。
1.2 具体实现
# torch.cuda.set_device(origin_device) # 设置.cuda()时默认使用的gpu_id
os.environ["CUDA_VISIBLE_DEVICS"] = "0,1