PyTorch & Distributed Framework Ray: tutorial introdutório para o nível de babá

Fonte:
tradução oficial do blog : comunidade de desenvolvedores PyTorch (conta oficial do WeChat)

O aprendizado de máquina de hoje requer computação distribuída. Seja treinando redes, ajustando hiperparâmetros, servindo modelos ou processando dados, o aprendizado de máquina é computacionalmente intensivo. Se você não acessar o cluster, a velocidade será muito lenta.

Ray é uma estrutura Python distribuída popular que pode ser combinada com o PyTorch para expandir rapidamente os aplicativos de aprendizado de máquina.

Este artigo apresenta os vários elementos do ecossistema Ray e como usá-lo com o PyTorch!

O que é Ray?

Ray é uma biblioteca de código aberto para Python paralelo e distribuído.

De uma perspectiva de alto nível, o ecossistema Ray consiste em três partes: o sistema Ray central, bibliotecas extensíveis para aprendizado de máquina (incluindo bibliotecas nativas e bibliotecas de terceiros) e um sistema para iniciar clusters em qualquer cluster ou ferramenta de provedor de nuvem.

Sistema central de Ray

Ray pode ser usado para estender aplicativos Python em vários núcleos ou máquinas. Tem várias vantagens principais, incluindo:

Simplicidade: você pode estender seu aplicativo Python sem reescrever, e o mesmo código pode ser executado em uma ou várias máquinas.

Robustez: o aplicativo pode lidar com as falhas da máquina e processar a preempção normalmente.

Desempenho: as tarefas são executadas com atrasos de milissegundos, podem ser estendidas a dezenas de milhares de núcleos e processam dados numéricos com sobrecarga de serialização mínima.

Ecologia da biblioteca de Ray

Como Ray é uma estrutura geral, a comunidade estabeleceu muitas bibliotecas e estruturas para realizar diferentes tarefas com base nela.

A maioria dessas bibliotecas e estruturas oferece suporte a PyTorch, com modificação mínima de código e integração perfeita entre si. A seguir estão algumas das muitas bibliotecas do ecossistema.

RaySGD

Legenda: Comparação de DataParallel e Ray de PyTorch na instância p3dn.24xlarge (Ray usa
DataParallel distribuído de PyTorch abaixo).

RaySGD é uma biblioteca que fornece pacotes de treinamento distribuído para treinamento paralelo de dados. Por exemplo, RaySGD TorchTrainer (https://docs.ray.io/en/master/raysgd/raysgd_pytorch.html) é um wrapper em torno de torch.distributed.launch. Ele fornece uma API Python que facilita a incorporação de treinamento distribuído em um aplicativo Python maior, em vez de envolver seu código de treinamento em um script bash.

Algumas outras vantagens da biblioteca são:

Fácil de usar: você pode estender o DistributedDataParallel nativo do PyTorch sem monitorar nós individuais.

Escalabilidade: você pode expandir os dados nativos distribuídos do PyTorch em paralelo sem monitorar um único nó. Você pode expandir para cima e para baixo. Comece com uma única CPU. Só precisa alterar duas linhas de código para expandir para clusters com vários nós, várias CPUs ou várias GPUs.

Treinamento acelerado: NVIDIA Apex possui suporte integrado para treinamento de precisão mista.

Função tolerante a falhas: oferece suporte à recuperação automática quando a máquina em nuvem é interrompida.

Compatibilidade: suporta integração perfeita com outras bibliotecas, como NVIDIA Apex. Pode ser perfeitamente integrado com outras bibliotecas, como Ray Tune e Ray Serve.

Você pode começar a usar o TorchTrainer instalando o Ray (pip install -U ray torch) e executando o seguinte código:

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
​
import ray
from ray.util.sgd.torch import TorchTrainer
# https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py
from ray.util.sgd.torch.resnet import ResNet18
​
​
def cifar_creator(config):
    """Returns dataloaders to be used in `train` and `validate`."""
    tfms = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2023, 0.1994, 0.2010)),
    ])  # meanstd transformation
    train_loader = DataLoader(
        CIFAR10(root="~/data", download=True, transform=tfms), batch_size=config["batch"])
    validation_loader = DataLoader(
        CIFAR10(root="~/data", download=True, transform=tfms), batch_size=config["batch"])
    return train_loader, validation_loader
​
​
def optimizer_creator(model, config):
    """Returns an optimizer (or multiple)"""
    return torch.optim.SGD(model.parameters(), lr=config["lr"])
​
ray.init()
​
trainer = TorchTrainer(
    model_creator=ResNet18,  # A function that returns a nn.Module
    data_creator=cifar_creator,  # A function that returns dataloaders
    optimizer_creator=optimizer_creator,  # A function that returns an optimizer
    loss_creator=torch.nn.CrossEntropyLoss,  # A loss function
    config={"lr": 0.01, "batch": 64},  # parameters
    num_workers=2,  # amount of parallelism
    use_gpu=torch.cuda.is_available(),
    use_tqdm=True)
​
stats = trainer.train()
print(trainer.validate())
​
torch.save(trainer.state_dict(), "checkpoint.pt")
trainer.shutdown()
print("success!")

Observação: este script fará o download do CIFAR10 e usará o modelo ResNet18 para classificação de imagens. Você só precisa alterar um parâmetro (num_workers = N) para utilizar várias GPUs.

Ray Tune

Legenda: O algoritmo de otimização implementado pelo Ray Tune, como o treinamento baseado em grupo (conforme mostrado na figura acima), pode ser usado com o PyTorch para obter um modelo de alto desempenho.

Ray Tune é uma biblioteca Python para execução de experimentos e ajuste de hiperparâmetros de qualquer escala. Algumas das vantagens desta biblioteca são:

• Capaz de iniciar uma varredura de hiperparâmetros distribuídos de vários nós em menos de 10 linhas de código.
• Suporta todas as principais estruturas de aprendizado de máquina, incluindo PyTorch.
• Suporte de primeira classe para GPU.
• Gerencie automaticamente os pontos de verificação e registre no TensorBoard.
• Acesse os algoritmos mais avançados, como treinamento baseado em grupo (PBT), BayesOptSearch, HyperBand / ASHA.

Você pode instalar o Ray (pip install ray torch torchvision) e executar o código a seguir para começar a usar o Ray Tune.

import numpy as np
import torch
import torch.optim as optim
​
from ray import tune
from ray.tune.examples.mnist_pytorch import get_data_loaders, train, test
import ray
import sys
​
if len(sys.argv) > 1:
    ray.init(redis_address=sys.argv[1])
​
import torch.nn as nn
import torch.nn.functional as F
​
class ConvNet(nn.Module):
    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 train_mnist(config):
    model = ConvNet()
    train_loader, test_loader = get_data_loaders()
    optimizer = optim.SGD(
        model.parameters(), lr=config["lr"], momentum=config["momentum"])
    for i in range(10):
        train(model, optimizer, train_loader, torch.device("cpu"))
        acc = test(model, test_loader, torch.device("cpu"))
        tune.track.log(mean_accuracy=acc)
        if i % 5 == 0:
            # This saves the model to the trial directory
            torch.save(model.state_dict(), "./model.pth")
​
from ray.tune.schedulers import ASHAScheduler
​
search_space = {
    "lr": tune.choice([0.001, 0.01, 0.1]),
    "momentum": tune.uniform(0.1, 0.9)
}
​
analysis = tune.run(
    train_mnist,
    num_samples=30,
    scheduler=ASHAScheduler(metric="mean_accuracy", mode="max", grace_period=1),
    config=search_space)

Observação: este script mostra como usar o algoritmo de parada antecipada mais avançado AHSA, que pode encerrar experimentos menos promissores e alocar mais tempo e recursos para os mais promissores.

Ray Serve

Legenda: Ray Serve pode ser usado não apenas para servir o modelo sozinho, mas também para estender outras ferramentas de serviço, como FastAPI.

Ray Serve é uma biblioteca fácil de usar de serviços de modelo extensível. Algumas das vantagens desta biblioteca são:

Um kit de ferramentas pode ser usado para servir de tudo, desde modelos de aprendizado profundo (PyTorch, TensorFlow etc.) a modelos scikit-learn e lógica de negócios arbitrária de Python.
Escalável para muitas máquinas, seja em seu data center ou na nuvem.
Compatível com muitas outras bibliotecas, como Ray Tune e FastAPI.

Se você deseja saber como integrar Ray Serve e Ray Tune em seu fluxo de trabalho PyTorch, você deve verificar a documentação para exemplos de código completos.

RLlib

Legenda: RLlib fornece métodos de personalização para quase todos os aspectos de treinamento, incluindo modelos de rede neural, distribuição de ação, definição de estratégia, ambiente e processo de coleta de amostra.

RLlib fornece métodos personalizados para quase todos os aspectos de treinamento, incluindo modelos de rede neural, distribuição de ação, definição de estratégia, ambiente e processo de coleta de amostra.

RLlib é uma biblioteca para aprendizado por reforço. Ela não apenas fornece alta escalabilidade, mas também fornece uma API unificada para vários aplicativos. As vantagens incluem:

• Compatível nativamente com PyTorch, TensorFlow Eager e TensorFlow (1.xe 2.x).
• Suporta algoritmos sem modelo, baseados em modelo, de evolução, planejamento e multiagentes.
• Suporta tipos de modelo complexos, como redes de atenção e pilhas LSTM, por meio de sinalizadores de configuração simples e wrappers automáticos.
• Compatibilidade com outras bibliotecas, como Ray Tune.

Cluster Launcher

Legenda: Ray Cluster Launcher simplifica o processo de inicialização e dimensionamento em qualquer cluster ou provedor de nuvem.

Depois de desenvolver um aplicativo em seu laptop e desejar estendê-lo para a nuvem (talvez com mais dados ou mais GPU), as próximas etapas nem sempre são claras. Esse processo será configurado pela equipe de infraestrutura para você ou percorrerá as seguintes etapas:

1. Escolha um provedor de nuvem (AWS, GCP ou Azure).
2. Navegue no console de gerenciamento para definir tipos de instância, políticas de segurança, nós, restrições de instância, etc.
3. Descubra como distribuir seus scripts Python no cluster.

Uma maneira mais fácil é usar o Ray Cluster Launcher para iniciar e expandir máquinas em qualquer cluster ou provedor de nuvem. O Cluster Launcher permite que você dimensione, sincronize arquivos, envie scripts, encaminhamento de porta, etc. Isso significa que você pode executar seu cluster Ray em Kubernetes, AWS, GCP, Azure ou clusters privados sem ter que entender os detalhes de baixo nível de gerenciamento de cluster.

Resumindo

Legenda: Ray fornece uma base de computação distribuída para o Fusion Engine do Ant Financial Group.

Este artigo contém alguns dos benefícios do Ray no ecossistema PyTorch. Ray é amplamente utilizado em várias aplicações, desde Ant Financial Group usando Ray para apoiar seus negócios financeiros, LinkedIn executando Ray on Yarn e Pathmind usando Ray para conectar aprendizagem por reforço com software de simulação e assim por diante.

Se você tiver dúvidas ou ideias sobre Ray, ou quiser aprender mais sobre Python paralelo e distribuído, junte-se à comunidade do projeto via Discourse, Slack ou GitHub.