Accelerate documentation

Iniciando o treinamento distribuído a partir de Jupyter Notebooks

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Iniciando o treinamento distribuído a partir de Jupyter Notebooks

Este tutorial ensina você a fazer o ajuste fino de um modelo de visão computacional com 🤗 Accelerate a partir de um Jupyter Notebook em um sistema distribuído. Você também aprenderá como configurar alguns requisitos necessários para garantir que seu ambiente esteja configurado corretamente, que seus dados tenham sido preparados adequadamente e, por fim, como iniciar o treinamento.

Este tutorial também está disponível como um Jupyter Notebook aqui

Configurando o Ambiente

Antes que qualquer treinamento possa ser realizado, um arquivo de configuração do Accelerate deve existir no sistema. Normalmente, isso pode ser feito executando o seguinte em um terminal e respondendo aos prompts:

accelerate config

No entanto, se os padrões gerais forem suficientes e você não estiver executando em um TPU, o Accelerate possui uma utilidade para escrever rapidamente sua configuração de dispositivo em um arquivo de configuração por meio de utils.write_basic_config().

O código a seguir reiniciará o Jupyter após gravar a configuração, pois o runtime CUDA ou o runtime XPU foi chamado para executar essa operação.

CUDA e XPU não podem ser inicializados mais de uma vez em um sistema multi-dispositivo. Não há problema em depurar no notebook e ter chamadas para CUDA/XPU, mas, para finalmente treinar, será necessário realizar uma limpeza completa e reiniciar.

import os
from accelerate.utils import write_basic_config

write_basic_config()  # Write a config file
os._exit(00)  # Restart the notebook

Preparando o Conjunto de Dados e o Modelo

A seguir, você deve preparar seu conjunto de dados. Como mencionado anteriormente, é preciso ter muito cuidado ao preparar os DataLoaders e o modelo para garantir que nada seja colocado em nenhuma GPU.

Se for o caso, recomenda-se colocar esse código específico em uma função e chamá-la a partir da interface de inicialização do notebook, que será mostrada mais adiante.

Certifique-se de que o conjunto de dados foi baixado com base nas instruções aqui

import os, re, torch, PIL
import numpy as np

from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator
from accelerate.utils import set_seed
from timm import create_model

Primeiro, você precisa criar uma função para extrair o nome da classe com base em um nome de arquivo:

import os

data_dir = "../../images"
fnames = os.listdir(data_dir)
fname = fnames[0]
print(fname)
beagle_32.jpg

Neste caso, o rótulo é beagle. Usando regex, você pode extrair o rótulo do nome do arquivo:

import re


def extract_label(fname):
    stem = fname.split(os.path.sep)[-1]
    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
extract_label(fname)

E você pode ver que ele retornou corretamente o nome certo para o nosso arquivo:

"beagle"

A seguir, uma classe Dataset deve ser criada para lidar com a obtenção da imagem e do rótulo:

class PetsDataset(Dataset):
    def __init__(self, file_names, image_transform=None, label_to_id=None):
        self.file_names = file_names
        self.image_transform = image_transform
        self.label_to_id = label_to_id

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx):
        fname = self.file_names[idx]
        raw_image = PIL.Image.open(fname)
        image = raw_image.convert("RGB")
        if self.image_transform is not None:
            image = self.image_transform(image)
        label = extract_label(fname)
        if self.label_to_id is not None:
            label = self.label_to_id[label]
        return {"image": image, "label": label}

Agora, para construir o conjunto de dados. Fora da função de treinamento, você pode encontrar e declarar todos os nomes de arquivos e rótulos e usá-los como referências dentro da função iniciada:

fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]

Agora, colete todos os rótulos:

all_labels = [extract_label(fname) for fname in fnames]
id_to_label = list(set(all_labels))
id_to_label.sort()
label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}

Em seguida, você deve criar uma função get_dataloaders que retornará os dataloaders que você construiu. Conforme mencionado anteriormente, se os dados forem enviados automaticamente para um dispositivo GPU ou TPU ao construir seus DataLoaders, eles devem ser construídos usando este método.

def get_dataloaders(batch_size: int = 64):
    "Builds a set of dataloaders with a batch_size"
    random_perm = np.random.permutation(len(fnames))
    cut = int(0.8 * len(fnames))
    train_split = random_perm[:cut]
    eval_split = random_perm[cut:]

    # For training a simple RandomResizedCrop will be used
    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)

    # For evaluation a deterministic Resize will be used
    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)

    # Instantiate dataloaders
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
    return train_dataloader, eval_dataloader

Por fim, você deve importar o agendador para ser usado mais tarde:

from torch.optim.lr_scheduler import CosineAnnealingLR

Escrevendo a Função de Treinamento

Agora você pode construir o loop de treinamento. notebook_launcher() funciona passando uma função para ser chamada que será executada em todo o sistema distribuído.

Aqui está um loop básico de treinamento para o problema de classificação de animais:

O código foi dividido para permitir explicações em cada seção. Uma versão completa, que pode ser copiada e colada, estará disponível no final

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    accelerator = Accelerator(mixed_precision=mixed_precision)

Primeiro, você deve definir a semente e criar um objeto Accelerator o mais cedo possível no loop de treinamento.

Se o treinamento for realizado em um TPU, o seu loop de treinamento deve receber o modelo como parâmetro e ele deve ser instanciado fora da função do loop de treinamento. Consulte as melhores práticas para TPU para entender o porquê

A seguir, você deve criar seus dataloaders e seu modelo:

    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

Você constrói o modelo aqui para que a semente também controle a inicialização dos novos pesos

Como você está executando transfer learning neste exemplo, o encoder do modelo começa congelado para que a head do modelo possa ser treinada apenas inicialmente:

    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

Normalizar os lotes de imagens tornará o treinamento um pouco mais rápido:

    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

Para tornar essas constantes disponíveis no dispositivo ativo, você deve configurá-lo como o dispositivo do Accelerator:

    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

Agora instancie o restante das classes do PyTorch usadas para o treinamento:

    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

Antes de passar tudo para prepare().

Não há uma ordem específica para memorizar; você só precisa descompactar os objetos na mesma ordem em que os passou para o método prepare.

    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

Agora treine o modelo:

    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

O loop de avaliação terá uma aparência um pouco diferente em comparação ao loop de treinamento. O número de elementos passados, bem como a precisão total de cada lote, será adicionado a duas constantes:

        model.eval()
        accurate = 0
        num_elems = 0

Em seguida, você tem o restante do seu loop padrão do PyTorch:

        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)

Antes de chegarmos à última grande diferença.

Ao realizar uma avaliação distribuída, as previsões e os rótulos precisam ser passados por gather() para que todos os dados estejam disponíveis no dispositivo atual e uma métrica corretamente calculada possa ser obtida:

            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

Agora você só precisa calcular a métrica real para este problema, e pode imprimi-la no processo principal usando print():

        eval_metric = accurate.item() / num_elems
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

Uma versão completa deste loop de treinamento está disponível abaixo:

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    # Initialize accelerator
    accelerator = Accelerator(mixed_precision=mixed_precision)
    # Build dataloaders
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)

    # Instantiate the model (you build the model here so that the seed also controls new weight initializations)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

    # Freeze the base model
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

    # You can normalize the batches of images to be a bit faster
    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

    # To make these constants available on the active device, set it to the accelerator device
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

    # Instantiate the optimizer
    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)

    # Instantiate the learning rate scheduler
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

    # Prepare everything
    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    # Now you train the model
    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

        model.eval()
        accurate = 0
        num_elems = 0
        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)
            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

        eval_metric = accurate.item() / num_elems
        # Use accelerator.print to print only on the main process.
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

Usando o notebook_launcher

O que resta é usar o notebook_launcher().

Você passa a função, os argumentos (como uma tupla) e o número de processos para o treinamento. (Veja a documentação para mais informações)

from accelerate import notebook_launcher
args = ("fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=2)

No caso de execução em vários nós, é necessário configurar uma sessão Jupyter em cada nó e executar a célula de lançamento ao mesmo tempo.

Para um ambiente contendo 2 nós (computadores) com 8 GPUs cada e o computador principal com endereço IP “172.31.43.8”, ficaria assim:

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)

E na segunda sessão do Jupyter na outra máquina:

Observe como o node_rank foi alterado

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)

No caso de execução no TPU, ficaria assim:

model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

args = (model, "fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=8)

Para iniciar o processo de treinamento com elasticidade, habilitando a tolerância a falhas, você pode usar o recurso elastic_launch fornecido pelo PyTorch. Isso requer a configuração de parâmetros adicionais, como rdzv_backend e max_restarts. Aqui está um exemplo de como usar o notebook_launcher com capacidades elásticas:

notebook_launcher(
    training_loop,
    args,
    num_processes=2,
    max_restarts=3
)

Durante a execução, ele exibirá o progresso e também indicará em quantos dispositivos a execução foi realizada. Este tutorial foi executado com duas GPUs:

Launching training on 2 GPUs.
epoch 0: 88.12
epoch 1: 91.73
epoch 2: 92.58
epoch 3: 93.90
epoch 4: 94.71

E é isso!

Observe que notebook_launcher() ignora o arquivo de configuração do Accelerate; para iniciar com base na configuração, use:

accelerate launch

Depuração

Um problema comum ao executar o notebook_launcher é receber um erro de que a CUDA/XPU já foi inicializada. Isso geralmente se deve a uma importação ou código anterior no notebook que faz uma chamada para a subbiblioteca torch.cuda ou torch.xpu do PyTorch. Para ajudar a identificar o que deu errado, você pode iniciar o notebook_launcher com ACCELERATE_DEBUG_MODE=yes no seu ambiente e uma verificação adicional será feita ao iniciar, garantindo que um processo regular possa ser criado e utilizar a CUDA/XPU sem problemas. (Seu código de CUDA/XPU ainda pode ser executado em seguida).

Conclusão

Este notebook mostrou como executar treinamento distribuído a partir de um Jupyter Notebook. Algumas observações importantes para lembrar:

  • Certifique-se de salvar qualquer código que use CUDA/XPU (ou imports de CUDA/XPU) para a função passada para notebook_launcher()
  • Defina o num_processes como o número de dispositivos usados para o treinamento (como o número de GPUs, XPUs, CPUs, TPUs, etc.)
  • Ao usar o TPU, declare o modelo fora da função do loop de treinamento
Update on GitHub