Aprendizaje Federado utilizando Hugging Face y Flower

'Aprendizaje Federado con Hugging Face y Flower'

Este tutorial mostrará cómo aprovechar Hugging Face para federar el entrenamiento de modelos de lenguaje en múltiples clientes utilizando Flower. Más específicamente, ajustaremos finamente un modelo Transformer preentrenado (distilBERT) para la clasificación de secuencias sobre un conjunto de datos de calificaciones de IMDB. El objetivo final es detectar si una calificación de película es positiva o negativa.

Un cuaderno también está disponible aquí, pero en lugar de ejecutarse en múltiples clientes separados, utiliza la funcionalidad de simulación de Flower (usando flwr['simulation']) para emular una configuración federada dentro de Google Colab (esto también significa que en lugar de llamar a start_server llamaremos a start_simulation, y que se necesitan algunas otras modificaciones).

Dependencias

Para seguir este tutorial, deberás instalar los siguientes paquetes: datasets, evaluate, flwr, torch y transformers. Esto se puede hacer utilizando pip:

pip install datasets evaluate flwr torch transformers

Flujo de trabajo estándar de Hugging Face

Manejo de los datos

Para obtener el conjunto de datos IMDB, utilizaremos la biblioteca datasets de Hugging Face. Luego necesitamos tokenizar los datos y crear dataloaders de PyTorch, todo esto se hace en la función load_data:

import random

import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, DataCollatorWithPadding


DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
CHECKPOINT = "distilbert-base-uncased"

def load_data():
    """Cargar datos de IMDB (entrenamiento y evaluación)"""
    raw_datasets = load_dataset("imdb")
    raw_datasets = raw_datasets.shuffle(seed=42)

    # eliminar división de datos innecesaria
    del raw_datasets["unsupervised"]

    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

    def tokenize_function(examples):
        return tokenizer(examples["text"], truncation=True)

    # Tomaremos una pequeña muestra para reducir el tiempo de cálculo, esto es opcional
    train_population = random.sample(range(len(raw_datasets["train"])), 100)
    test_population = random.sample(range(len(raw_datasets["test"])), 100)

    tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
    tokenized_datasets["train"] = tokenized_datasets["train"].select(train_population)
    tokenized_datasets["test"] = tokenized_datasets["test"].select(test_population)
    tokenized_datasets = tokenized_datasets.remove_columns("text")
    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
    trainloader = DataLoader(
        tokenized_datasets["train"],
        shuffle=True,
        batch_size=32,
        collate_fn=data_collator,
    )

    testloader = DataLoader(
        tokenized_datasets["test"], batch_size=32, collate_fn=data_collator
    )

    return trainloader, testloader
    
trainloader, testloader = load_data()

Entrenamiento y prueba del modelo

Una vez que tenemos una forma de crear nuestro trainloader y testloader, podemos encargarnos del entrenamiento y la prueba. Esto es muy similar a cualquier ciclo de entrenamiento o prueba de PyTorch:

from evaluate import load as load_metric
from transformers import AdamW


def train(net, trainloader, epochs):
    optimizer = AdamW(net.parameters(), lr=5e-5)
    net.train()
    for _ in range(epochs):
        for batch in trainloader:
            batch = {k: v.to(DEVICE) for k, v in batch.items()}
            outputs = net(**batch)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

def test(net, testloader):
    metric = load_metric("accuracy")
    loss = 0
    net.eval()
    for batch in testloader:
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        with torch.no_grad():
            outputs = net(**batch)
        logits = outputs.logits
        loss += outputs.loss.item()
        predictions = torch.argmax(logits, dim=-1)
        metric.add_batch(predictions=predictions, references=batch["labels"])
    loss /= len(testloader.dataset)
    accuracy = metric.compute()["accuracy"]
    return loss, accuracy

Creando el modelo en sí

Para crear el modelo en sí, simplemente cargaremos el modelo pre-entrenado distillBERT utilizando AutoModelForSequenceClassification de Hugging Face:

from transformers import AutoModelForSequenceClassification 


net = AutoModelForSequenceClassification.from_pretrained(
        CHECKPOINT, num_labels=2
    ).to(DEVICE)

Federando el ejemplo

La idea detrás del Aprendizaje Federado es entrenar un modelo entre múltiples clientes y un servidor sin tener que compartir ningún dato. Esto se logra permitiendo que cada cliente entrene el modelo localmente en sus datos y enviar sus parámetros de regreso al servidor, que luego agrega todos los parámetros de los clientes utilizando una estrategia predefinida. Este proceso se simplifica mucho utilizando el framework Flower. Si deseas obtener una visión más completa, asegúrate de consultar esta guía: ¿Qué es el Aprendizaje Federado?

Creando el IMDBClient

Para federar nuestro ejemplo a varios clientes, primero debemos escribir nuestra clase cliente de Flower (heredando de flwr.client.NumPyClient). Esto es muy fácil, ya que nuestro modelo es un modelo estándar de PyTorch:

from collections import OrderedDict

import flwr as fl


class IMDBClient(fl.client.NumPyClient):
        def get_parameters(self, config):
            return [val.cpu().numpy() for _, val in net.state_dict().items()]

        def set_parameters(self, parameters):
            params_dict = zip(net.state_dict().keys(), parameters)
            state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
            net.load_state_dict(state_dict, strict=True)

        def fit(self, parameters, config):
            self.set_parameters(parameters)
            print("Entrenamiento iniciado...")
            train(net, trainloader, epochs=1)
            print("Entrenamiento finalizado.")
            return self.get_parameters(config={}), len(trainloader), {}

        def evaluate(self, parameters, config):
            self.set_parameters(parameters)
            loss, accuracy = test(net, testloader)
            return float(loss), len(testloader), {"accuracy": float(accuracy)}

La función get_parameters permite al servidor obtener los parámetros del cliente. Inversamente, la función set_parameters permite al servidor enviar sus parámetros al cliente. Finalmente, la función fit entrena el modelo localmente para el cliente, y la función evaluate prueba el modelo localmente y devuelve las métricas relevantes.

Ahora podemos iniciar instancias de clientes utilizando:

fl.client.start_numpy_client(server_address="127.0.0.1:8080", 
                                                             client=IMDBClient())

Iniciando el servidor

Ahora que tenemos una forma de instanciar clientes, necesitamos crear nuestro servidor para poder agregar los resultados. Utilizando Flower, esto se puede hacer fácilmente eligiendo primero una estrategia (aquí estamos utilizando FedAvg, que definirá los pesos globales como el promedio de los pesos de todos los clientes en cada ronda) y luego utilizando la función flwr.server.start_server:

def weighted_average(metrics):
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    losses = [num_examples * m["loss"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, _ in metrics]
    return {"accuracy": sum(accuracies) / sum(examples), "loss": sum(losses) / sum(examples)}

# Define la estrategia
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,
    fraction_evaluate=1.0,
    evaluate_metrics_aggregation_fn=weighted_average,
)

# Inicia el servidor
fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=3),
    strategy=strategy,
)

La función weighted_average está ahí para proporcionar una forma de agregar las métricas distribuidas entre los clientes (básicamente, esto nos permite mostrar un promedio agradable de precisión y pérdida para cada ronda).

Poniendo todo junto

Si quieres ver todo junto, deberías revisar el ejemplo de código que escribimos para el repositorio de Flower: https://github.com/adap/flower/tree/main/examples/quickstart_huggingface .

Por supuesto, este es un ejemplo muy básico y se puede agregar o modificar mucho, fue solo para mostrar cómo podríamos federar fácilmente un flujo de trabajo de Hugging Face usando Flower.

Tenga en cuenta que en este ejemplo usamos PyTorch, pero muy bien podríamos haber usado TensorFlow.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

10 millones se registran en la aplicación rival de Twitter de Meta, Threads.

La experiencia de microblogging similar a Twitter sugiere que Meta Platforms se ha estado preparando para desafiar di...

Inteligencia Artificial

Este artículo de IA de la Universidad de Tokio ha aplicado el aprendizaje profundo al problema de la simulación de supernovas.

Investigadores de la Universidad de Tokio han desarrollado un modelo de aprendizaje profundo llamado 3D-Memory In Mem...

Inteligencia Artificial

Los satélites más antiguos de observación de la Tierra de NOAA obtienen 'vida prolongada

La Administración Nacional Oceánica y Atmosférica utilizará un sistema basado en la nube para extender la vida de los...

Inteligencia Artificial

Gafas de realidad virtual para ratones crean escenarios inmersivos para la investigación cerebral

Investigadores de la Universidad Northwestern desarrollaron gafas de realidad virtual para ratones.