目录
- 目录
- pytorch多gpu并行训练
- 1.单机多卡并行训练
- 1.1.torch.nn.DataParallel
- 1.2.如何平衡DataParallel带来的显存使用不平衡的问题
- 1.3.torch.nn.parallel.DistributedDataParallel
- 2.多机多gpu训练
- 2.1.初始化
- 2.1.1.初始化backend
- 2.1.2.初始化init_method
- 2.1.2.1.使用TCP初始化
- 2.1.2.2.使用共享文件系统初始化
- 2.1.3.初始化rank和world_size
- 2.1.4.初始化中一些需要注意的地方
- 2.2.数据的处理-DataLoader
- 2.3.模型的处理
- 2.4.模型的保存与加载
- 2.1.初始化
- 1.单机多卡并行训练
pytorch多gpu并行训练
注: 以下都在Ubuntu上面进行的调试, 使用的Ubuntu版本包括14, 18LST
参考文档:
data_parallel_tutorial
distributeddataparallel
environment-variable-initialization
PYTORCH 1.0 DISTRIBUTED TRAINER WITH AMAZON AWS
pytorch/examples/imagenet/main.py
Distributed-VGG-F
Getting Started with Distributed Data Parallel
1.单机多卡并行训练
1.1.torch.nn.DataParallel
我一般在使用多GPU的时候, 会喜欢使用os.environ['CUDA_VISIBLE_DEVICES']
来限制使用的GPU个数, 例如我要使用第0和第3编号的GPU, 那么只需要在程序中设置:
os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
但是要注意的是, 这个参数的设定要保证在模型加载到gpu上之前, 我一般都是在程序开始的时候就设定好这个参数, 之后如何将模型加载到多GPU上面呢?
如果是模型, 那么需要执行下面的这几句代码:
model = nn.DataParallel(model)
model = model.cuda()
如果是数据, 那么直接执行下面这几句代码就可以了:
inputs = inputs.cuda()
labels = labels.cuda()
其实如果看pytorch官网给的示例代码,我们可以看到下面这样的代码
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
model.to(device)
这个和我上面写的好像有点不太一样, 但是如果看一下DataParallel
的内部代码, 我们就可以发现, 其实是一样的:
class DataParallel(Module):
def __init__(self, module, device_ids=None, output_device=None, dim=0):
super(DataParallel, self).__init__()
if not torch.cuda.is_available():
self.module = module
self.device_ids = []
return
if device_ids is None:
device_ids = list(range(torch.cuda.device_count()))
if output_device is None:
output_device = device_ids[0]
我截取了其中一部分代码, 我们可以看到如果我们不设定好要使用的device_ids的
话, 程序会自动找到这个机器上面可以用的所有的显卡, 然后用于训练. 但是因为我们前面使用os.environ['CUDA_VISIBLE_DEVICES']
限定了这个程序可以使用的显卡, 所以这个地方程序如果自己获取的话, 获取到的其实就是我们上面设定的那几个显卡.
我没有进行深入得到考究, 但是我感觉使用os.environ['CUDA_VISIBLE_DEVICES']
对可以使用的显卡进行限定之后, 显卡的实际编号和程序看到的编号应该是不一样的, 例如上面我们设定的是os.environ['CUDA_VISIBLE_DEVICES&#