PyTorch 多进程|极客笔记

PyTorch 多进程

在深度学习的训练过程中，往往需要大量的计算资源来加速模型训练。PyTorch 提供了多进程的能力，可以充分利用多核CPU或多GPU来并行计算，从而提高训练速度。本文将详细介绍 PyTorch 中多进程的使用方法。

多进程概述

多进程是指在一个程序中同时运行多个独立的进程。在深度学习中，可以利用多进程来加速数据加载、模型训练等计算密集型任务。PyTorch 提供了 torch.multiprocessing 模块来实现多进程功能。

创建多进程

在 PyTorch 中创建一个多进程需要借助 torch.multiprocessing.Process 类。以下是一个简单的示例代码，演示了如何创建一个多进程并运行一个简单的任务：

import torch
import torch.multiprocessing as mp

def my_task(rank):
    print(f"Running task in process {rank}")

if __name__ == "__main__":
    processes = []
    for rank in range(4):
        p = mp.Process(target=my_task, args=(rank,))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

在上面的代码中，定义了一个 my_task 函数，接受一个 rank 参数，并打印当前进程的 rank。在 main 函数中，创建了 4 个进程分别执行 my_task 函数，并通过 join 方法等待所有进程执行完毕。

运行以上代码，可以看到输出如下：

Running task in process 0
Running task in process 1
Running task in process 2
Running task in process 3

共享内存

在多进程中，每个进程都有自己独立的内存空间，无法直接共享变量。但是在 PyTorch 中提供了 torch.multiprocessing.Array 和 torch.multiprocessing.Value 来实现共享内存。

torch.multiprocessing.Array 可以创建一个共享数组，可以在多个进程之间共享数据。下面是一个示例代码：

import torch
import torch.multiprocessing as mp

def shared_task(arr, val):
    arr[val] *= 2

if __name__ == "__main__":
    arr = mp.Array('d', [1.0, 2.0, 3.0, 4.0])
    val = mp.Value('i', 1)

    processes = []
    for _ in range(4):
        p = mp.Process(target=shared_task, args=(arr, val))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    print(arr[:])

在上面的代码中，创建了一个含有四个元素的共享数组 arr 和一个共享变量 val。多个进程可以通过共享数组 arr 和共享变量 val 来共享数据。运行以上代码，可以看到输出如下：

[1.0, 4.0, 3.0, 4.0]

多GPU加速

在深度学习中，通常会使用多个GPU来加速训练。PyTorch 提供了 torch.nn.DataParallel 类来实现模型的多GPU并行计算。以下是一个简单的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

def train(rank):
    model = Net()
    model = nn.DataParallel(model)
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    for _ in range(10):
        data = torch.randn(2, 10)
        target = torch.randn(2, 1)

        optimizer.zero_grad()
        output = model(data)
        loss = nn.MSELoss()(output, target)
        loss.backward()
        optimizer.step()

if __name__ == "__main__":
    processes = []
    for rank in range(4):
        p = mp.Process(target=train, args=(rank,))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

在上面的代码中，定义了一个简单的神经网络模型 Net，并使用 nn.DataParallel 实现了模型的多GPU并行计算。在 train 函数中，每个进程训练模型 10 次。运行以上代码，可以看到模型在多个GPU上并行训练。