Démarrage rapide 🤗 Transformateurs

Construisons un système d’apprentissage fédéré à l’aide des transformateurs Hugging Face et de Flower !

Nous nous appuierons sur Hugging Face pour fédérer l’entraînement de modèles de langage sur plusieurs clients à l’aide de Flower. Plus précisément, nous mettrons au point un modèle Transformer pré-entraîné (distilBERT) pour la classification de séquences sur un ensemble de données d’évaluations IMDB. L’objectif final est de détecter si l’évaluation d’un film est positive ou négative.

Dépendances

Pour suivre ce tutoriel, tu devras installer les paquets suivants : datasets, evaluate, flwr, torch, et transformers. Cela peut être fait en utilisant pip :

$ pip install datasets evaluate flwr torch transformers

Flux de travail standard pour le visage

Traitement des données

Pour récupérer le jeu de données IMDB, nous utiliserons la bibliothèque datasets de Hugging Face. Nous devons ensuite tokeniser les données et créer des PyTorch dataloaders, ce qui est fait dans la fonction 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():
    """Load IMDB data (training and eval)"""
    raw_datasets = load_dataset("imdb")
    raw_datasets = raw_datasets.shuffle(seed=42)
    # remove unnecessary data split
    del raw_datasets["unsupervised"]
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
    def tokenize_function(examples):
        return tokenizer(examples["text"], truncation=True)
    # We will take a small sample in order to reduce the compute time, this is optional
    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

Former et tester le modèle

Une fois que nous avons trouvé un moyen de créer notre trainloader et notre testloader, nous pouvons nous occuper de l’entraînement et du test. C’est très similaire à n’importe quelle boucle d’entraînement ou de test 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

Créer le modèle lui-même

Pour créer le modèle lui-même, nous allons simplement charger le modèle distillBERT pré-entraîné en utilisant le AutoModelForSequenceClassification de Hugging Face :

from transformers import AutoModelForSequenceClassification

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

Fédérer l’exemple

Création du client IMDBC

Pour fédérer notre exemple à plusieurs clients, nous devons d’abord écrire notre classe de client Flower (héritant de flwr.client.NumPyClient). C’est très facile, car notre modèle est un modèle PyTorch standard :

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("Training Started...")
            train(net, trainloader, epochs=1)
            print("Training Finished.")
            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 fonction get_parameters permet au serveur d’obtenir les paramètres du client. Inversement, la fonction set_parameters permet au serveur d’envoyer ses paramètres au client. Enfin, la fonction fit forme le modèle localement pour le client, et la fonction evaluate teste le modèle localement et renvoie les mesures correspondantes.

Démarrer le serveur

Maintenant que nous avons un moyen d’instancier les clients, nous devons créer notre serveur afin d’agréger les résultats. Avec Flower, cela peut être fait très facilement en choisissant d’abord une stratégie (ici, nous utilisons FedAvg, qui définira les poids globaux comme la moyenne des poids de tous les clients à chaque tour) et en utilisant ensuite la fonction 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 strategy
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,
    fraction_evaluate=1.0,
    evaluate_metrics_aggregation_fn=weighted_average,
)

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

La fonction weighted_average est là pour fournir un moyen d’agréger les mesures réparties entre les clients (en gros, cela nous permet d’afficher une belle moyenne de précision et de perte pour chaque tour).

Tout assembler

Nous pouvons maintenant démarrer des instances de clients en utilisant :

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

Et ils pourront se connecter au serveur et démarrer la formation fédérée.

If you want to check out everything put together, you should check out the full code example .

Bien sûr, c’est un exemple très basique, et beaucoup de choses peuvent être ajoutées ou modifiées, il s’agissait juste de montrer avec quelle simplicité on pouvait fédérer un flux de travail Hugging Face à l’aide de Flower.

Notez que dans cet exemple, nous avons utilisé PyTorch, mais nous aurions très bien pu utiliser TensorFlow.