使用 Optimum 高效运行 SetFit 模型

SetFit 是一种用于少样本文本分类的技术，它使用对比学习在几乎没有或没有标记数据的领域中微调 Sentence Transformers。它实现了与基于大型语言模型的现有最先进方法相当的性能，但不需要提示，并且训练效率很高（在 GPU 上通常只需几秒钟，在 CPU 上只需几分钟）。

在本笔记本中，您将学习如何进一步压缩 SetFit 模型，以便在使用 Optimum Onnx 在 GPU 上进行更快的推理和部署。

1. 设置开发环境

我们的第一步是安装 SetFit。运行以下单元格将为我们安装所有必需的软件包。

!pip install setfit accelerate -qqq

2. 创建性能基准

在我们训练和优化任何模型之前，让我们先定义一个性能基准，我们可以使用它来比较我们的模型。通常，在生产环境中部署 ML 模型涉及到几个约束之间的权衡

模型性能：模型在精心制作的测试集上的表现如何？
延迟：我们的模型可以多快地交付预测？
内存：我们可以在哪个云实例或设备上存储和加载我们的模型？

下面的类定义了一个简单的基准，用于衡量给定 SetFit 模型和测试数据集的每个指标

from pathlib import Path
from time import perf_counter

import evaluate
import numpy as np
import torch
from tqdm.auto import tqdm

metric = evaluate.load("accuracy")


class PerformanceBenchmark:
    def __init__(self, model, dataset, optim_type):
        self.model = model
        self.dataset = dataset
        self.optim_type = optim_type

    def compute_accuracy(self):
        preds = self.model.predict(self.dataset["text"])
        labels = self.dataset["label"]
        accuracy = metric.compute(predictions=preds, references=labels)
        print(f"Accuracy on test set - {accuracy['accuracy']:.3f}")
        return accuracy

    def compute_size(self):
        state_dict = self.model.model_body.state_dict()
        tmp_path = Path("model.pt")
        torch.save(state_dict, tmp_path)
        # Calculate size in megabytes
        size_mb = Path(tmp_path).stat().st_size / (1024 * 1024)
        # Delete temporary file
        tmp_path.unlink()
        print(f"Model size (MB) - {size_mb:.2f}")
        return {"size_mb": size_mb}

    def time_model(self, query="that loves its characters and communicates something rather beautiful about human nature"):
        latencies = []
        # Warmup
        for _ in range(10):
            _ = self.model([query])
        # Timed run
        for _ in range(100):
            start_time = perf_counter()
            _ = self.model([query])
            latency = perf_counter() - start_time
            latencies.append(latency)
        # Compute run statistics
        time_avg_ms = 1000 * np.mean(latencies)
        time_std_ms = 1000 * np.std(latencies)
        print(rf"Average latency (ms) - {time_avg_ms:.2f} +\- {time_std_ms:.2f}")
        return {"time_avg_ms": time_avg_ms, "time_std_ms": time_std_ms}

    def run_benchmark(self):
        metrics = {}
        metrics[self.optim_type] = self.compute_size()
        metrics[self.optim_type].update(self.compute_accuracy())
        metrics[self.optim_type].update(self.time_model())
        return metrics

除此之外，我们将创建一个简单的函数来绘制此基准报告的性能。

import matplotlib.pyplot as plt
import pandas as pd


def plot_metrics(perf_metrics):
    df = pd.DataFrame.from_dict(perf_metrics, orient="index")

    for idx in df.index:
        df_opt = df.loc[idx]
        plt.errorbar(
            df_opt["time_avg_ms"],
            df_opt["accuracy"] * 100,
            xerr=df_opt["time_std_ms"],
            fmt="o",
            alpha=0.5,
            ms=df_opt["size_mb"] / 15,
            label=idx,
            capsize=5,
            capthick=1,
        )

    legend = plt.legend(loc="lower right")

    plt.ylim(63, 95)
    # Use the slowest model to define the x-axis range
    xlim = max([metrics["time_avg_ms"] for metrics in perf_metrics.values()]) * 1.2
    plt.xlim(0, xlim)
    plt.ylabel("Accuracy (%)")
    plt.xlabel("Average latency with batch_size=1 (ms)")
    plt.show()

3. 训练/评估 bge-small SetFit 模型

在优化任何模型之前，让我们先训练一些基线模型作为参考点。我们将使用 sst-2 数据集，这是一个情感文本的集合，分为 2 类：正面，负面

让我们首先从 Hub 加载数据集

from datasets import load_dataset

dataset = load_dataset("SetFit/sst2")
dataset

DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 6920
    })
    validation: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 872
    })
    test: Dataset({
        features: ['text', 'label', 'label_text'],
        num_rows: 1821
    })
})

我们使用完整数据集训练 SetFit 模型。回想一下，SetFit 在少样本场景中表现出色，但这次我们有兴趣实现最大精度。

train_dataset = dataset["train"]
test_dataset = dataset["validation"]

使用以下代码行下载已经微调的模型并进行评估。或者，取消注释下面的代码以从头开始微调基础模型。

请注意，我们使用免费的 T4 GPU 在 Google Colab 上执行评估。

# Evaluate the uploaded model!
from setfit import SetFitModel

small_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
pb = PerformanceBenchmark(model=small_model, dataset=test_dataset, optim_type="bge-small (PyTorch)")
perf_metrics = pb.run_benchmark()

Model size (MB) - 127.33
Accuracy on test set - 0.906
Average latency (ms) - 17.42 +\- 4.47

# # Fine-tune the base model and Evaluate!
# from setfit import SetFitModel, Trainer, TrainingArguments

# # Load pretrained model from the Hub
# small_model = SetFitModel.from_pretrained(
#    "BAAI/bge-small-en-v1.5"
# )
# args = TrainingArguments(num_iterations=20)

# # Create trainer
# small_trainer = Trainer(
#    model=small_model, args=args, train_dataset=train_dataset
# )
# # Train!
# small_trainer.train()

# # Evaluate!
# pb = PerformanceBenchmark(
#    model=small_trainer.model, dataset=test_dataset, optim_type="bge-small (base)"
# )
# perf_metrics = pb.run_benchmark()

让我们绘制结果以可视化性能

plot_metrics(perf_metrics)

setfit_torch

4. 使用 Optimum ONNX 和 CUDAExecutionProvider 进行压缩

我们将使用 Optimum 的 ONNX Runtime 支持和 CUDAExecutionProvider，因为它速度快，同时还支持动态形状。

!pip install optimum[onnxruntime-gpu] -qqq

optimum-cli 使将模型导出到 ONNX 并应用 SOTA 图优化/内核融合变得非常容易。

!optimum-cli export onnx \
  --model moshew/bge-small-en-v1.5_setfit-sst2-english \
  --task feature-extraction \
  --optimize O4 \
  --device cuda \
  bge_auto_opt_O4

我们可能会看到一些警告，但这些警告无需担心。稍后我们将看到它不影响模型性能。

首先，我们将创建性能基准的子类，以允许基准测试 ONNX 模型。

class OnnxPerformanceBenchmark(PerformanceBenchmark):
    def __init__(self, *args, model_path, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_path = model_path

    def compute_size(self):
        size_mb = Path(self.model_path).stat().st_size / (1024 * 1024)
        print(f"Model size (MB) - {size_mb:.2f}")
        return {"size_mb": size_mb}

然后，我们可以使用 "CUDAExecutionProvider" 提供程序加载转换后的 SentenceTransformer 模型。也可以随意尝试其他提供程序，例如 "TensorrtExecutionProvider" 和 "CPUExecutionProvider"。前者可能比 "CUDAExecutionProvider" 更快，但需要更多安装步骤。

import torch
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForFeatureExtraction

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('bge_auto_opt_O4', model_max_length=512)
ort_model = ORTModelForFeatureExtraction.from_pretrained('bge_auto_opt_O4', provider="CUDAExecutionProvider")

让我们创建一个使用 tokenizer、ONNX Runtime (ORT) 模型和 SetFit 模型头的类。

from setfit.exporters.utils import mean_pooling


class OnnxSetFitModel:
    def __init__(self, ort_model, tokenizer, model_head):
        self.ort_model = ort_model
        self.tokenizer = tokenizer
        self.model_head = model_head

    def predict(self, inputs):
        encoded_inputs = self.tokenizer(
            inputs, padding=True, truncation=True, return_tensors="pt"
        ).to(self.ort_model.device)

        outputs = self.ort_model(**encoded_inputs)
        embeddings = mean_pooling(
            outputs["last_hidden_state"], encoded_inputs["attention_mask"]
        )
        return self.model_head.predict(embeddings.cpu())

    def __call__(self, inputs):
        return self.predict(inputs)

我们可以像这样初始化此模型

model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
onnx_setfit_model = OnnxSetFitModel(ort_model, tokenizer, model.model_head)

# Perform inference
onnx_setfit_model(test_dataset["text"][:2])

array([0, 0])

是时候对这个 ONNX 模型进行基准测试了。

pb = OnnxPerformanceBenchmark(
    onnx_setfit_model,
    test_dataset,
    "bge-small (optimum ONNX)",
    model_path="bge_auto_opt_O4/model.onnx",
)
perf_metrics.update(pb.run_benchmark())

plot_metrics(perf_metrics)

setfit_onnx

通过应用 ONNX，我们将每个样本的延迟从 13.43 毫秒提高到 2.19 毫秒，速度提高了 6.13 倍！

为了进一步改进，我们建议增加推理批次大小，因为这也可以大大提高吞吐量。例如，将批次大小设置为 128 会将延迟进一步降低到 0.3 毫秒，而在批次大小为 2048 时则降低到 0.2 毫秒。

< > 在 GitHub 上更新