Python单程序 多进程 多显卡同时独立训练多个模型

一、问题描述

        在一个python程序中实现多个模型（使用Pytorch实现）独立地训练，每个模型可指定的显卡，模型之间没有共享变量、参数也完全独立。如下图：

        类似的问题（但不是上面描述的问题）：模型多显卡训练，这应该是指在多张显卡训练一个模型。而我们上面描述的是指在多张显卡中分别训练多个独立的模型，并且是并行地训练。

二、实现方式

        思路：使用python自带的并行库multiprocessing，添加多个子进程，模型在子进程（可指定不同的显卡）中进行训练。备注：各个子进程的内存空间一般是独立的，当然可以自由添加各个子进程的共享变量，但这些共享变量必须要加锁更新才行，否则会出现各种问题。而且，共享变量越多，多进程并行就越复杂。

        示例代码

import os
import time
import torch
import torch.nn as nn
import multiprocessing as mp

# 控制模型选择的函数
def train(is_end, gpu_id, model_name):
    if model_name == 'TransE':
        transe_run(gpu_id, model_name)
        transe_run(gpu_id, model_name)
        with is_end.get_lock():
            is_end.value += 1
    elif model_name == 'TransH':
        transh_run(gpu_id, model_name)
        transh_run(gpu_id, model_name)
        transh_run(gpu_id, model_name)
        with is_end.get_lock():
            is_end.value += 1
    else:
        sacn_run(gpu_id, model_name)
        sacn_run(gpu_id, model_name)
        sacn_run(gpu_id, model_name)
        sacn_run(gpu_id, model_name)
        with is_end.get_lock():
            is_end.value += 1

# 父进程中的主函数
def main(os_context):
    # 设置一个共享变量，用来记录完成训练的子进程个数
    is_end = os_context.Value("i", 0)
    # 创建多个子进程
    transe_trainer = os_context.Process(target=train, args=(is_end, 0, 'TransE'))
    transh_trainer = os_context.Process(target=train, args=(is_end, 1, 'TransH'))
    sacn_trainer = os_context.Process(target=train, args=(is_end, 2, 'SACN'))
    # 启动子进程
    transe_trainer.start()
    transh_trainer.start()
    sacn_trainer.start()
    # 下面的代码使得共享变量能够同步更新
    transe_trainer.join()
    transh_trainer.join()
    sacn_trainer.join()
    # 输出完成训练的子进程个数
    print('finish count: %d' % is_end.value)


# 程序入口
if __name__ == '__main__':
    # 输出CPU个数和GPU个数
    print('The machine has %d CPUs.' % os.cpu_count())
    print('The machine has %d GPUs.' % torch.cuda.device_count())
    # 获取操作系统上下文
    context = mp.get_context('spawn')
    # 运行主函数
    main(context)


# 模拟TransE训练
def transe_run(gpu_id, model_name):
    # 设置模型训练所在GPU
    torch.cuda.set_device(gpu_id)
    torch.cuda.is_available()
    print('b train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))
    # 模型的初始化
    entity_embedding = nn.Parameter(torch.zeros(2000, 50))
    nn.init.uniform_(tensor=entity_embedding, a=-20, b=20)
    relation_embedding = nn.Parameter(torch.zeros(16, 50))
    nn.init.uniform_(tensor=relation_embedding, a=-20, b=20)
    entity_embedding.cuda(gpu_id)
    relation_embedding.cuda(gpu_id)
    # 模拟模型的训练
    for i in range(100):
        entity_embedding = entity_embedding + torch.tensor([0.1])
        relation_embedding = relation_embedding + torch.tensor([-0.1])
        time.sleep(1.5)
    print('a train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))


# 模拟TransH训练
def transh_run(gpu_id, model_name):
    # 设置模型训练所在GPU
    torch.cuda.set_device(gpu_id)
    torch.cuda.is_available()
    print('b train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))
    # 模型的初始化
    entity_embedding = nn.Parameter(torch.zeros(2000, 100))
    nn.init.uniform_(tensor=entity_embedding, a=-20, b=20)
    relation_embedding = nn.Parameter(torch.zeros(16, 100))
    nn.init.uniform_(tensor=relation_embedding, a=-20, b=20)
    entity_embedding.cuda(gpu_id)
    relation_embedding.cuda(gpu_id)
    # 模拟模型的训练
    for i in range(100):
        entity_embedding = entity_embedding + torch.tensor([0.1])
        relation_embedding = relation_embedding + torch.tensor([-0.1])
        time.sleep(1.5)
    print('a train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))


# 模拟SACN训练
def sacn_run(gpu_id, model_name):
    # 设置模型训练所在GPU
    torch.cuda.set_device(gpu_id)
    torch.cuda.is_available()
    print('b train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))
    # 模型的初始化
    entity_embedding = nn.Parameter(torch.zeros(2000, 150))
    nn.init.uniform_(tensor=entity_embedding, a=-20, b=20)
    relation_embedding = nn.Parameter(torch.zeros(16, 150))
    nn.init.uniform_(tensor=relation_embedding, a=-20, b=20)
    entity_embedding.cuda(gpu_id)
    relation_embedding.cuda(gpu_id)
    # 模拟模型的训练
    for i in range(100):
        entity_embedding = entity_embedding + torch.tensor([0.1])
        relation_embedding = relation_embedding + torch.tensor([-0.1])
        time.sleep(1.5)
    print('a train [%s] on gpu[%d]' % (model_name, torch.cuda.current_device()))

        效果（模拟程序输出）

The machine has 88 CPUs.
The machine has 4 GPUs.
b train [TransH] on gpu[1]
b train [TransE] on gpu[0]
b train [SACN] on gpu[2]
a train [TransH] on gpu[1]
b train [TransH] on gpu[1]
a train [SACN] on gpu[2]
b train [SACN] on gpu[2]
a train [TransE] on gpu[0]
b train [TransE] on gpu[0]
a train [TransH] on gpu[1]
b train [TransH] on gpu[1]
a train [SACN] on gpu[2]
b train [SACN] on gpu[2]
a train [TransE] on gpu[0]
a train [TransH] on gpu[1]
a train [SACN] on gpu[2]
b train [SACN] on gpu[2]
a train [SACN] on gpu[2]
finish count: 3

总结

        实现上述的方式不仅仅可以通过multiprocessing，也可以通过shell的方式（一个shell脚本对应一个进程，在命令行窗口执行多个脚本就能开多个进程），并且这种方式更加简单。然而，使用shell的方式多进程训练也有一定的限制，例如它只能在命令行窗口中执行。不过，按道理来说python应该也是可以调用shell脚本来执行的（猜的，待以后有空再做调研）。