Question

在我的程序中，我需要共同训练共享单词嵌入的两个不同模型。我必须使用多GPU进行训练，因为我的数据太多了。如果以单个GPU进行训练，则将花费7天以上。因此，我的想法描述如下图。

mdoel-0在GPU-0和GPU-1之间分布训练，模型1在GPU-2和GPU-3上的工作方式相同。此外，两个模型都共享单词嵌入参数，它们将由它们同步更新。

我已就此写了一份草稿。该草案利用了在一台计算机上针对多GPU的多处理和分布式培训的优势。 run_0和run_1的功能分别引用model-0和model-1。

import torch
args.distributed_word_size=2
mp = torch.multiprocessing.get_context('spawn')
args.distributed_init_method_run_0 = 'tcp://localhost:00000'
args.distributed_init_method_run_1 = 'tcp://localhost:11111'

procs = []
for i in range(args.distributed_word_size):
    args.distributed_rank = i
    args.device_id = i
    procs.append(mp.Process(target=run_0, args=(args, ), daemon=True))
    procs[i].start()
for i in range(args.distributed_word_size):
    args.distributed_rank = i
    args.device_id = i + args.distributed_word_size
    procs.append(mp.Process(target=run_1, args=(args, ), daemon=True))
    procs[i+args.distributed_word_size].start()
for p in procs:
    p.joint()

def run_0(args):
    torch.distributed.init_process_group(
            backend=args.distributed_backend,
            init_method=args.distributed_init_method_run_0,
            world_size=args.distributed_world_size,
            rank=args.distributed_rank)
    main_0(args)
    ....

def run_1(args):
    torch.distributed.init_process_group(
            backend=args.distributed_backend,
            init_method=args.distributed_init_method_run_1,
            world_size=args.distributed_world_size,
            rank=args.distributed_rank)
    main_1(args)
    ....

我想知道上面的代码段是否正确，并且有更好的解决方案。特别是，如何在两个模型中实现同步更新的共享词嵌入？

使用多GPU共同训练两个模型

0 个答案: