如何在pytorch中正确使用分布式数据并行

时间:2019-07-06 00:35:59

标签: deep-learning pytorch lstm

我是PyTorch的新手,但我在Keras和TensorFlow方面拥有丰富的经验。我已遵循此article在自己的训练脚本上使用DDP。但是,由于某种原因,我总是最终得到:

  

进程0以退出状态1终止。

我尝试在多个GPU平台(Google Colab,Kaggle,Floyd Hub)上运行相同的代码,但是它们几乎都引发了相同的错误。

我还尝试过禁用join=True选项,但是培训过程甚至没有开始。

与DDP相关的代码:

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # Explicitly setting seed to make sure that models created in two processes
    # start from same random weights and biases.
    torch.manual_seed(42)


def cleanup():
    dist.destroy_process_group()


def run_demo(fn, *args):
    mp.spawn(
        fn,
        args = (args[0], args[1], args[2], args[3], args[4]),
        nprocs = 1, # Also tried 2 , but no difference
        join = True
    )

我的训练代码:

def train(model, X, batch_size = 32, epochs = 75, gradient_acc = 0):
    setup(1, 2)
    device = model.get_default_device()
    model = model.to(device, non_blocking = True)
    ddp_model = DDP(model, device_ids = [0]) # Only one GPU

    # ...

    ddp_model.hidden_enc = ddp_model.init_hidden_enc()
    ddp_model.hidden_dec = ddp_model.init_hidden_dec()

    ddp_model.train()

    for ep in range(epochs):
        loss_br = 0; nb_batch_steps = 0
        for step, batch in enumerate( data_loader ):
            batch = batch.to(device, non_blocking = True)
            nb_batch_steps += 1
            loss = ddp_model(batch)
            # ...
    cleanup()

致电培训代码:

if __name__ == "__main__":
    run_demo(
        train,
        model, 
        holder[:], # X
        32, 
        75,
        3
    )

我希望模型可以使用分布式并行程序包在多个进程上运行。有趣的是,有时在不使用DDP的情况下运行CUDA时会遇到Out of Memory异常。我知道spawn.py会终止所有进程,如果status code > 1存在任何可用进程,但是我似乎还无法弄清楚如何避免此问题。我们非常感谢您的帮助。

1 个答案:

答案 0 :(得分:0)

根据pytorch DDP tutorial

  

在整个过程中,DDP在正向传递中插入必要的参数同步,在向后传递中插入渐变同步。用户可以将进程映射到可用资源,只要进程不共享GPU设备即可。

我认为当每个进程需要共享此GPU设备时,您无法在一个GPU上的DDP中初始化模型。