Accelerate documentation

Iniciar el entrenamiento distribuido desde Jupyter Notebooks

Accelerate

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).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Iniciar el entrenamiento distribuido desde Jupyter Notebooks

En este tutorial aprenderás a realizar el ajuste fino de un modelo de visión por computadora con 🤗 Accelerate desde un Jupyter Notebook en un sistema distribuido. También aprenderás a configurar algunos requisitos necesarios para garantizar que tu entorno esté configurado correctamente, que tus datos estén preparados adecuadamente y, por último, cómo iniciar el entrenamiento.

Este tutorial también está disponible como cuaderno de Jupyter aquí

Configuración del entorno

Antes de poder realizar cualquier entrenamiento, debe existir un archivo de configuración de Accelerate en el sistema. Normalmente, esto se puede hacer ejecutando lo siguiente en una terminal y respondiendo a las indicaciones:

accelerate config

Sin embargo, si los valores predeterminados generales son adecuados y no estás ejecutando en un TPU, Accelerate tiene una utilidad para escribir rápidamente tu configuración de dispositivo en un archivo de configuración mediante utils.write_basic_config().

El siguiente código reiniciará Jupyter después de escribir la configuración, ya que se invocó el runtime de CUDA o el runtime de XPU para realizar esta operación.

No se pueden inicializar CUDA y XPU más de una vez en un sistema de varios dispositivos. No hay problema con depurar en el cuaderno y realizar llamadas a CUDA/XPU, pero para entrenar finalmente se deberá realizar una limpieza completa y un reinicio.

import os
from accelerate.utils import write_basic_config

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

Preparación del conjunto de datos y el modelo

A continuación, debe preparar su conjunto de datos. Como se mencionó anteriormente, debe tener mucho cuidado al preparar los DataLoaders y el modelo para asegurarse de que nada se coloque en ninguna GPU.

Si es así, se recomienda colocar ese código específico en una función y llamarla desde la interfaz de inicio del notebook, que se mostrará más adelante.

Asegúrate de que el conjunto de datos se descargue siguiendo las instrucciones aquí

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

Primero, necesitas crear una función para extraer el nombre de la clase en función de un nombre de archivo:

import os

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

beagle_32.jpg

En este caso, la etiqueta es beagle. Usando expresiones regulares, puedes extraer la etiqueta del nombre del archivo:

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)

Y se puede ver que devolvió correctamente el nombre adecuado para nuestro archivo:

"beagle"

A continuación, se debe crear una clase Dataset para gestionar la obtención de la imagen y la etiqueta:

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}

Ahora, para construir el conjunto de datos. Fuera de la función de entrenamiento, puedes buscar y declarar todos los nombres de archivos y etiquetas, y usarlos como referencias dentro de la función lanzada:

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

A continuación, recolecta todas las etiquetas:

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)}

A continuación, debe crear una función get_dataloaders que devolverá los dataloaders que haya construido. Como se mencionó anteriormente, si los datos se envían automáticamente a la GPU o a un dispositivo TPU al construir sus DataLoaders, estos deben crearse mediante 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 último, debe importar el programador que se utilizará más adelante:

from torch.optim.lr_scheduler import CosineAnnealingLR

Escribiendo la función de entrenamiento

Ahora puedes construir el bucle de entrenamiento. notebook_launcher() funciona pasando una función que se llamará y se ejecutará en todo el sistema distribuido.

Aquí hay un bucle de entrenamiento básico para el problema de clasificación de animales:

El código se ha dividido para permitir explicaciones en cada sección. Una versión completa que se puede copiar y pegar estará disponible al final

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

Primero, debe establecer la semilla y crear un objeto Accelerator lo antes posible en el bucle de entrenamiento.

Si el entrenamiento se realiza en un TPU, el bucle de entrenamiento debe recibir el modelo como parámetro y este debe instanciarse fuera de la función del bucle de entrenamiento. Consulta las mejores prácticas para TPU para saber por qué.

A continuación, debe construir sus dataloaders y crear su modelo:

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

Construye el modelo aquí para que la semilla también controle la inicialización de los nuevos pesos

Como en este ejemplo se realiza transfer learning, el codificador del modelo comienza congelado para que la cabeza del modelo pueda entrenarse solo al inicio:

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

Normalizar los lotes de imágenes hará que el entrenamiento sea un poco más rápido:

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

Para hacer estas constantes disponibles en el dispositivo activo, debe establecerlo como el dispositivo del Accelerator:

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

A continuación, instancie el resto de las clases de PyTorch utilizadas para el entrenamiento:

    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 pasar todo a prepare().

No hay un orden específico que deba recordarse; solo necesita desempacar los objetos en el mismo orden en que los pasó al método prepare.

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

Ahora entrena el 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()

El bucle de evaluación se verá ligeramente diferente en comparación con el bucle de entrenamiento. El número de elementos pasados, así como la precisión total de cada lote, se añadirán a dos constantes:

        model.eval()
        accurate = 0
        num_elems = 0

A continuación tienes el resto de tu bucle estándar de 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 pasar a la última gran diferencia.

Al realizar una evaluación distribuida, las predicciones y las etiquetas deben pasar por gather() para que todos los datos estén disponibles en el dispositivo actual y se pueda obtener una métrica calculada correctamente:

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

Ahora solo necesitas calcular la métrica real para este problema, y puedes imprimirla en el proceso principal usando print():

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

A continuación se encuentra una versión completa de este bucle de entrenamiento:

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 el notebook_launcher

Lo único que queda es usar notebook_launcher().

Pasa la función, los argumentos (como una tupla) y el número de procesos en los que entrenar. (Consulta la documentación para obtener más información)

from accelerate import notebook_launcher

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

En el caso de ejecutar en varios nodos, es necesario configurar una sesión de Jupyter en cada nodo y ejecutar la celda de lanzamiento al mismo tiempo.

Para un entorno que contiene 2 nodos (computadoras) con 8 GPUs cada uno y la computadora principal con una dirección IP de “172.31.43.8”, se vería así:

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

Y en la segunda sesión de Jupyter en la otra máquina:

Fíjate en cómo ha cambiado node_rank

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

En el caso de ejecutarlo en la TPU, se vería así:

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 el proceso de entrenamiento con elasticidad, habilitando la tolerancia a fallos, puede utilizar la función elastic_launch proporcionada por PyTorch. Esto requiere establecer parámetros adicionales como rdzv_backend y max_restarts. A continuación se muestra un ejemplo de cómo usar notebook_launcher con capacidades elásticas:

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

Mientras se ejecuta, mostrará el progreso e indicará en cuántos dispositivos se ejecutó. Este tutorial se ejecutó con dos 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

¡Y eso es todo!

Tenga en cuenta que notebook_launcher() ignora el archivo de configuración de Accelerate; para lanzar basándose en la configuración, use:

accelerate launch

Depuración

Un problema común al ejecutar notebook_launcher es recibir un error de que CUDA/XPU ya ha sido inicializado. Esto suele deberse a una importación o a código previo en el notebook que realiza una llamada a la sublibrería de PyTorch torch.cuda o torch.xpu. Para ayudar a identificar qué salió mal, puedes lanzar notebook_launcher con ACCELERATE_DEBUG_MODE=yes en tu entorno y se realizará una comprobación adicional al generar el proceso para verificar que se pueda crear un proceso regular y utilizar CUDA/XPU sin problemas. (Tu código de CUDA/XPU aún puede ejecutarse después).

Conclusión

Este cuaderno mostró cómo realizar entrenamiento distribuido desde dentro de un Jupyter Notebook. Algunas notas clave para recordar:

Asegúrate de guardar cualquier código que use CUDA/XPU (o importaciones de CUDA/XPU) para la función pasada a notebook_launcher()
Establezca num_processes como el número de dispositivos utilizados para el entrenamiento (como el número de GPUs, XPUs, CPUs, TPUs, etc.)
Si se utiliza la TPU, declare su modelo fuera de la función del bucle de entrenamiento

Update on GitHub

←Lanzamiento de scripts de Accelerate ¡Empieza aquí!→