使用 Hugging Face 和 Flower 進行聯邦學習

釋出於 2023 年 3 月 27 日

訪客

本教程將展示如何利用 Hugging Face，透過 Flower 在多個客戶端上聯邦訓練語言模型。更具體地說，我們將針對 IMDB 評分資料集對預訓練的 Transformer 模型 (distilBERT) 進行微調，以進行序列分類。最終目標是檢測電影評分是正面還是負面。

還有一份筆記本可在此處獲取，但它不是在多個獨立的客戶端上執行，而是利用 Flower 的模擬功能（使用 flwr['simulation']）在 Google Colab 中模擬聯邦設定（這也意味著我們將呼叫 start_simulation 而不是 start_server，並且需要進行一些其他修改）。

依賴項

要按照本教程操作，您需要安裝以下軟體包：datasets、evaluate、flwr、torch 和 transformers。這可以使用 pip 完成

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():
    """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
    
trainloader, testloader = load_data()

訓練和測試模型

一旦我們有了建立訓練器和測試器的方法，我們就可以進行訓練和測試。這與任何 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

為了將我們的示例聯邦化到多個客戶端，我們首先需要編寫我們的 Flower 客戶端類（繼承自 flwr.client.NumPyClient）。這非常容易，因為我們的模型是一個標準的 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("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)}

get_parameters 函式允許伺服器獲取客戶端的引數。相反，set_parameters 函式允許伺服器將其引數傳送給客戶端。最後，fit 函式在本地為客戶端訓練模型，而 evaluate 函式在本地測試模型並返回相關指標。

我們現在可以使用以下方式啟動客戶端例項：

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

啟動伺服器

現在我們有了例項化客戶端的方法，我們需要建立我們的伺服器以聚合結果。使用 Flower，這可以非常容易地完成，首先選擇一個策略（這裡我們使用 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)}

# 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,
)