参数服务器#

小技巧

对于生产级的分布式训练实现，请使用 Ray Train。

参数服务器是一个用于分布式机器学习训练的框架。

在参数服务器框架中，中央服务器（或一组服务器节点）维护机器学习模型（例如，神经网络）的全局共享参数，而计算更新（即梯度下降更新）所需的数据和计算则分布在工作节点上。

../../_images/param_actor.png

参数服务器是许多机器学习应用程序的核心部分。本文件将介绍如何使用Ray actor实现简单的同步和异步参数服务器。

要运行该应用程序，首先安装一些依赖项。

pip install torch torchvision filelock

我们首先定义一些辅助函数并导入一些依赖项。

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from filelock import FileLock
import numpy as np

import ray


def get_data_loader():
    """安全下载数据。返回训练/验证集数据加载器。"""
    mnist_transforms = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    # 我们在此处添加 FileLock，因为多个工作进程会希望
    # 下载数据，这可能会导致覆盖，因为
    # DataLoader 不是线程安全的。
    with FileLock(os.path.expanduser("~/data.lock")):
        train_loader = torch.utils.data.DataLoader(
            datasets.MNIST(
                "~/data", train=True, download=True, transform=mnist_transforms
            ),
            batch_size=128,
            shuffle=True,
        )
        test_loader = torch.utils.data.DataLoader(
            datasets.MNIST("~/data", train=False, transform=mnist_transforms),
            batch_size=128,
            shuffle=True,
        )
    return train_loader, test_loader


def evaluate(model, test_loader):
    """评估模型在验证数据集上的准确性。"""
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            # 这仅用于加快评估速度。
            if batch_idx * len(data) > 1024:
                break
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    return 100.0 * correct / total

设置：定义神经网络#

我们定义一个小型神经网络用于训练。我们提供一些辅助函数来获取数据，包括梯度和权重的获取/设置方法。

class ConvNet(nn.Module):
    """用于MNIST数据集的小型卷积神经网络。"""

    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
        self.fc = nn.Linear(192, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 3))
        x = x.view(-1, 192)
        x = self.fc(x)
        return F.log_softmax(x, dim=1)

    def get_weights(self):
        return {k: v.cpu() for k, v in self.state_dict().items()}

    def set_weights(self, weights):
        self.load_state_dict(weights)

    def get_gradients(self):
        grads = []
        for p in self.parameters():
            grad = None if p.grad is None else p.grad.data.cpu().numpy()
            grads.append(grad)
        return grads

    def set_gradients(self, gradients):
        for g, p in zip(gradients, self.parameters()):
            if g is not None:
                p.grad = torch.from_numpy(g)

定义参数服务器#

参数服务器将保存模型的副本。 在训练过程中，它将：

  1. 接收梯度并将其应用到模型中。

  2. 将更新后的模型发送回工作节点。

@ray.remote 装饰器定义了一个远程进程。它封装了 ParameterServer 类，并允许用户将其实例化为一个远程演员。

@ray.remote
class ParameterServer(object):
    def __init__(self, lr):
        self.model = ConvNet()
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr)

    def apply_gradients(self, *gradients):
        summed_gradients = [
            np.stack(gradient_zip).sum(axis=0) for gradient_zip in zip(*gradients)
        ]
        self.optimizer.zero_grad()
        self.model.set_gradients(summed_gradients)
        self.optimizer.step()
        return self.model.get_weights()

    def get_weights(self):
        return self.model.get_weights()

定义工作节点#

工作节点还会保存模型的副本。在训练过程中，它将 持续评估数据并将梯度发送 到参数服务器。工作节点将其模型与 参数服务器的模型权重进行同步。

@ray.remote
class DataWorker(object):
    def __init__(self):
        self.model = ConvNet()
        self.data_iterator = iter(get_data_loader()[0])

    def compute_gradients(self, weights):
        self.model.set_weights(weights)
        try:
            data, target = next(self.data_iterator)
        except StopIteration:  # 当一个时期结束时，开始一个新的时期。
            self.data_iterator = iter(get_data_loader()[0])
            data, target = next(self.data_iterator)
        self.model.zero_grad()
        output = self.model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        return self.model.get_gradients()

同步参数服务器训练#

现在我们将创建一个同步参数服务器训练方案。我们将首先实例化一个参数服务器的进程，并伴随多个工作节点。

iterations = 200
num_workers = 2

ray.init(ignore_reinit_error=True)
ps = ParameterServer.remote(1e-2)
workers = [DataWorker.remote() for i in range(num_workers)]

我们还将在驱动进程上实例化一个模型，以便在训练期间评估测试准确性。

model = ConvNet()
test_loader = get_data_loader()[1]

训练交替进行：

  1. 在当前权重的基础上计算梯度

  2. 使用梯度更新参数服务器的权重。

print("Running synchronous parameter server training.")
current_weights = ps.get_weights.remote()
for i in range(iterations):
    gradients = [worker.compute_gradients.remote(current_weights) for worker in workers]
    # 在所有梯度可用后计算更新。
    current_weights = ps.apply_gradients.remote(*gradients)

    if i % 10 == 0:
        # 评估当前模型。
        model.set_weights(ray.get(current_weights))
        accuracy = evaluate(model, test_loader)
        print("Iter {}: \taccuracy is {:.1f}".format(i, accuracy))

print("Final accuracy is {:.1f}.".format(accuracy))
# 在下一个示例之前，清理 Ray 资源和进程。
ray.shutdown()

异步参数服务器训练#

我们现在将创建一个异步参数服务器训练方案。我们将首先实例化一个参数服务器进程，以及多个工作进程。

print("Running Asynchronous Parameter Server Training.")

ray.init(ignore_reinit_error=True)
ps = ParameterServer.remote(1e-2)
workers = [DataWorker.remote() for i in range(num_workers)]

在这里，工作节点将异步计算给定当前权重的梯度，并在准备好后立即将这些梯度发送到参数服务器。当参数服务器完成应用新的梯度后，服务器将向工作节点发送当前权重的副本。然后，工作节点将更新权重并重复该过程。

current_weights = ps.get_weights.remote()

gradients = {}
for worker in workers:
    gradients[worker.compute_gradients.remote(current_weights)] = worker

for i in range(iterations * num_workers):
    ready_gradient_list, _ = ray.wait(list(gradients))
    ready_gradient_id = ready_gradient_list[0]
    worker = gradients.pop(ready_gradient_id)

    # 计算并应用梯度。
    current_weights = ps.apply_gradients.remote(*[ready_gradient_id])
    gradients[worker.compute_gradients.remote(current_weights)] = worker

    if i % 10 == 0:
        # 每更新10次后评估当前模型。
        model.set_weights(ray.get(current_weights))
        accuracy = evaluate(model, test_loader)
        print("Iter {}: \taccuracy is {:.1f}".format(i, accuracy))

print("Final accuracy is {:.1f}.".format(accuracy))

最终想法#

这种方法是强大的，因为它使您能够仅通过几行代码在 Python 应用程序中实现参数服务器。因此，这简化了使用参数服务器的应用程序的部署，并修改参数服务器的行为。

例如，分片参数服务器、更改更新规则、在异步和同步更新之间切换、忽略滞后工作节点，或任何其他自定义，仅需几行额外的代码。