Press "Enter" to skip to content

使用Hugging Face和Flower的联合学习

使用Hugging Face和Flower的联合学习 四海 第1张

本教程将展示如何利用Hugging Face和Flower在多个客户端上联合训练语言模型。更具体地说,我们将对IMDB评分数据集上的预训练Transformer模型(distilBERT)进行微调,用于序列分类。最终目标是判断电影评分是积极的还是消极的。

此外,还提供了一个笔记本,但它不是在多个独立的客户端上运行,而是利用Flower的模拟功能(使用flwr['simulation']),以在Google Colab中模拟联合设置(这意味着我们将调用start_simulation而不是start_server,并且需要进行一些其他修改)。

依赖项

为了跟随本教程,您需要安装以下软件包:datasetsevaluateflwrtorchtransformers。可以使用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():
    """加载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

为了将我们的示例联合到多个客户端,我们首先需要编写我们的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("开始训练...")
            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())

启动服务器

现在我们有了实例化客户端的方法,我们需要创建服务器以聚合结果。使用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)}

# 定义策略
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函数提供了一种在客户端分布的指标之间进行聚合的方法(基本上允许我们显示每一轮的平均准确率和损失)。

将所有内容组合在一起

如果您想要查看所有内容的完整示例,请查看我们为Flower存储库编写的代码示例:https://github.com/adap/flower/tree/main/examples/quickstart_huggingface。

当然,这只是一个非常基本的例子,可以添加或修改很多内容。这只是为了展示如何简单地使用Flower来联合Hugging Face的工作流程。

请注意,在这个例子中我们使用了PyTorch,但我们也可以很好地使用TensorFlow

Leave a Reply

Your email address will not be published. Required fields are marked *