使用 Optimum 高效執行 SetFit 模型

SetFit 是一種用於少樣本文字分類的技術，它使用對比學習來微調 Sentence Transformers，適用於幾乎沒有標註資料的領域。它的效能與現有基於大型語言模型的最新方法相當，但無需提示，並且訓練效率高（通常在 GPU 上只需幾秒鐘，在 CPU 上只需幾分鐘）。

在本筆記本中，您將學習如何使用 Optimum Onnx 進一步壓縮 SetFit 模型，以實現更快的 GPU 推理和部署。

1. 設定開發環境

我們的第一步是安裝 SetFit。執行以下單元格將為我們安裝所有必需的包。

!pip install setfit accelerate -qqq

2. 建立效能基準

在訓練和最佳化任何模型之前，我們先定義一個性能基準，用於比較我們的模型。通常，在生產環境中部署機器學習模型涉及權衡以下幾個約束：

模型效能：模型在精心設計的測試集上的表現如何？
延遲：我們的模型能以多快的速度提供預測？
記憶體：我們可以在哪種雲實例或裝置上儲存和載入模型？

下面的類定義了一個簡單的基準，用於衡量給定 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 資料集，這是一個情感文字集合，分為兩類：積極和消極。

我們首先從 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"]

使用以下程式碼行下載已微調的模型並進行評估。或者，取消註釋其下方的程式碼以從頭開始微調基礎模型。

請注意，我們在 Google Colab 上使用免費的 T4 GPU 進行評估。

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

我們來建立一個使用分詞器、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.43ms 縮短到 2.19ms，速度提高了 6.13 倍！

為了進一步提升效能，我們建議增加推理批次大小，因為這也會大大提高吞吐量。例如，將批次大小設定為 128 會將延遲進一步降低到 0.3ms，在批次大小為 2048 時則降至 0.2ms。

< > 在 GitHub 上更新