Hugging FaceとFlowerを使用したフェデレーテッドラーニング

'Hugging FaceとFlowerでのフェデレーテッドラーニング'

このチュートリアルでは、Hugging Faceを使用して、Flowerを介して複数のクライアント上で言語モデルのトレーニングをフェデレートする方法を紹介します。具体的には、IMDBの評価データセットを使用して、事前トレーニングされたTransformerモデル（distilBERT）をシーケンス分類のために微調整します。最終的な目標は、映画の評価がポジティブかネガティブかを検出することです。

ノートブックはこちらでご利用いただけますが、複数のクライアントで実行する代わりに、Google Colab内でフェデレーテッド環境をエミュレートするためにFlowerのシミュレーション機能（flwr['simulation']）を使用します（これはまた、start_serverを呼び出す代わりにstart_simulationを呼び出す必要があり、その他の変更が必要です）。

依存関係

このチュートリアルに従うためには、以下のパッケージをインストールする必要があります：datasets、evaluate、flwr、torch、およびtransformers。これはpipを使用して行うことができます：

• 大規模言語モデルの高速推論：Habana Gaudi2アクセラレータ上のBLOOMZ
• 🤗 Transformersを使用してTensorFlowとTPUで言語モデルをトレーニングする
• ビジョン言語モデルの高速化：Habana Gaudi2上のBridgeTower

pip install datasets evaluate flwr torch transformers

標準的なHugging Faceのワークフロー

データの処理

IMDBデータセットを取得するために、Hugging Faceのdatasetsライブラリを使用します。その後、データをトークン化し、PyTorchのデータローダーを作成する必要があります。これはすべて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():
    """IMDBデータをロードする（トレーニングと評価）"""
    raw_datasets = load_dataset("imdb")
    raw_datasets = raw_datasets.shuffle(seed=42)

    # 不要なデータ分割を削除
    del raw_datasets["unsupervised"]

    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

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

    # 計算時間を短縮するために、小さなサンプルを取得します（オプション）
    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()

モデルのトレーニングとテスト

trainloaderとtestloaderを作成する方法がわかれば、トレーニングとテストを行うことができます。これは、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

モデルの作成

モデル自体を作成するには、Hugging FaceのAutoModelForSequenceClassificationを使用して事前学習済みのdistillBERTモデルをロードするだけです：

from transformers import AutoModelForSequenceClassification 


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

例のフェデレーテッド学習

フェデレーテッド学習のアイデアは、データを共有せずに複数のクライアントとサーバー間でモデルをトレーニングすることです。これは、各クライアントがデータを使用してモデルをローカルでトレーニングし、そのパラメーターをサーバーに送信し、サーバーがあらかじめ定義された戦略を使用してすべてのクライアントのパラメーターを集約することによって行われます。このプロセスは、Flowerフレームワークを使用することで非常に簡単に行うことができます。より完全な概要が必要な場合は、次のガイドを参照してください：フェデレーテッド学習とは何ですか？

IMDBClientの作成

複数のクライアントに例をフェデレートするためには、まずflwr.client.NumPyClientを継承したFlowerクライアントクラスを作成する必要があります。これは非常に簡単で、モデルは標準の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("トレーニング開始...")
            train(net, trainloader, epochs=1)
            print("トレーニング終了.")
            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)}

get_parameters関数は、サーバーがクライアントのパラメーターを取得するのに使用されます。逆に、set_parameters関数は、サーバーがクライアントにパラメーターを送信するのに使用されます。最後に、fit関数はクライアントのためにモデルをローカルでトレーニングし、evaluate関数はモデルをローカルでテストし、関連するメトリックを返します。

次のコマンドを使用してクライアントインスタンスを起動できます：

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

サーバーの起動

クライアントをインスタンス化する方法があるので、結果を集計するためにサーバーを作成する必要があります。これは、まず戦略を選択する（ここではFedAvgを使用しています。これにより、各ラウンドでクライアントの重みの平均がグローバルな重みとして定義されます）し、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)}

# 戦略を定義する
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,
    fraction_evaluate=1.0,
    evaluate_metrics_aggregation_fn=weighted_average,
)

# サーバーを起動する
fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=3),
    strategy=strategy,
)

weighted_average関数は、クライアント間で分散されたメトリックを集約する方法を提供するためのものです（基本的には、各ラウンドごとに平均的な精度と損失を表示するためのものです）。

すべてをまとめる

すべてをまとめたものを確認したい場合は、次のコード例をチェックしてください：https://github.com/adap/flower/tree/main/examples/quickstart_huggingface 。

もちろん、これは非常に基本的な例であり、追加や変更ができる余地がたくさんあります。これは、Hugging Faceワークフローを簡単にFederateする方法を示すためのものでした。

この例では、PyTorchを使用しましたが、TensorFlowを使用することもできました。

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