pytorch训练时gpu利用率低_Pytorch多GPU训练

以下都在Ubuntu上面进行的调试, 使用的Ubuntu版本包括14, 18LST

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']="0,2", 但是程序看到的显卡编号应该被改成了 '0,1', 也就是说程序所使用的显卡编号实际上是经过了一次映射之后才会映射到真正的显卡编号上面的, 例如这里的程序看到的1对应实际的2

1.2.如何平衡DataParallel带来的显存使用不平衡的问题

这个问题其实讨论的也比较多了, 官方给的解决方案就是使用  DistributedDataParallel来代替  DataParallel(实际上 DistributedDataParallel显存分配的也不是很平衡), 但是从某些角度来说,  DataParallel使用起来确实比较方便, 而且最近使用  DistributedDataParallel 遇到一些小问题. 所以这里提供一个解决显存使用不平衡问题的方案:首先这次的解决方案来自transformer-XL的官方代码: https://github.com/kimiyoung/transformer-xl然后我将其中的平衡GPU显存的代码提取了出来(原代码好像有点小问题)放到了github上面:https://github.com/Link-Li/Balanced-DataParallel这里的代码是原作者继承了  DataParallel 类之后进行了改写:

class BalancedDataParallel(DataParallel):    def __init__(self, gpu0_bsz, *args, **kwargs