分類頭

任何 🤗 SetFit 模型都由兩部分組成：一個 SentenceTransformer 嵌入體和一個分類頭。

本指南將向您展示

內建的邏輯迴歸分類頭
內建的可微分分類頭
自定義分類頭的要求

邏輯迴歸分類頭

當初始化一個新的 SetFit 模型時，預設選擇 scikit-learn 邏輯迴歸頭。事實證明，將其應用於微調後的句子轉換器主體上時非常有效，並且它仍然是推薦的分類頭。使用邏輯迴歸頭初始化一個新的 SetFit 模型很簡單

>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5")
>>> model.model_head
LogisticRegression()

要使用額外的引數初始化邏輯迴歸頭（或任何其他頭），您可以使用 SetFitModel.from_pretrained() 上的 head_params 引數

>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", head_params={"solver": "liblinear", "max_iter": 300})
>>> model.model_head
LogisticRegression(max_iter=300, solver='liblinear')

可微分分類頭

SetFit 還提供了 SetFitHead 作為專門的 torch 分類頭。它使用線性層將嵌入對映到類。可以透過將 SetFitModel.from_pretrained() 上的 use_differentiable_head 引數設定為 True 來使用它

>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True)
>>> model.model_head
SetFitHead({'in_features': 384, 'out_features': 2, 'temperature': 1.0, 'bias': True, 'device': 'cuda'})

預設情況下，這將假定為二元分類。要更改這一點，還可以透過 head_params 將 out_features 設定為您正在使用的類別數量。

>>> from setfit import SetFitModel

>>> model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True, head_params={"out_features": 5})
>>> model.model_head
SetFitHead({'in_features': 384, 'out_features': 5, 'temperature': 1.0, 'bias': True, 'device': 'cuda'})

與預設的邏輯迴歸頭不同，可微分分類頭只支援 [0, num_classes) 範圍內的整數標籤。

使用可微分分類頭進行訓練

使用 SetFitHead 可以解鎖一些新的 TrainingArguments，這些引數不會用於基於 sklearn 的頭。請注意，使用 SetFit 進行訓練在幕後包含兩個階段：**微調嵌入**和**訓練分類頭**。因此，一些訓練引數可以是元組，其中兩個值分別用於這兩個階段。對於許多這些情況，第二個值僅在分類頭可微分時才使用。例如

batch_size：（Union[int, Tuple[int, int]]，預設為 (16, 2)）- 元組中的第二個值決定了訓練可微分 SetFitHead 時的批次大小。
num_epochs：（Union[int, Tuple[int, int]]，預設為 (1, 16)）- 元組中的第二個值決定了訓練可微分 SetFitHead 時的 epoch 數量。實際上，訓練分類頭的 num_epochs 通常更大。這有兩個原因
1. 此訓練階段不使用對比對進行訓練，因此與微調嵌入模型不同，每個標註的訓練文字只獲得一個訓練樣本。
2. 此訓練階段涉及從頭開始訓練分類器，而不是微調一個已經有能力的模型。我們需要更多的訓練步驟。
end_to_end：（bool，預設為 False）- 如果為 True，則在分類器訓練階段端到端訓練整個模型。否則，凍結 Sentence Transformer 主體並僅訓練頭部。
body_learning_rate：（Union[float, Tuple[float, float]]，預設為 (2e-5, 1e-5)）- 元組中的第二個值決定了在分類器訓練階段 Sentence Transformer 主體的學習率。這僅在 end_to_end 為 True 時相關，因為否則在訓練分類器時 Sentence Transformer 主體會被凍結。
head_learning_rate（float，預設為 1e-2）- 此值決定了在分類器訓練階段可微分頭的學習率。僅在使用可微分頭時使用。
l2_weight（float，*可選*）- 可選的 L2 權重，僅在使用可微分頭時傳遞給分類器訓練階段的 AdamW 最佳化器，用於模型主體和頭部。

例如，使用可微分分類頭的完整訓練指令碼可能如下所示

from setfit import SetFitModel, Trainer, TrainingArguments, sample_dataset
from datasets import load_dataset

# Initializing a new SetFit model
model = SetFitModel.from_pretrained("BAAI/bge-small-en-v1.5", use_differentiable_head=True, head_params={"out_features": 2})

# Preparing the dataset
dataset = load_dataset("SetFit/sst2")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=32)
test_dataset = dataset["test"]

# Preparing the training arguments
args = TrainingArguments(
    batch_size=(32, 16),
    num_epochs=(3, 8),
    end_to_end=True,
    body_learning_rate=(2e-5, 5e-6),
    head_learning_rate=2e-3,
    l2_weight=0.01,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()
# ***** Running training *****
#   Num examples = 66
#   Num epochs = 3
#   Total optimization steps = 198
#   Total train batch size = 3
# {'embedding_loss': 0.2204, 'learning_rate': 1.0000000000000002e-06, 'epoch': 0.02}                                                                                 
# {'embedding_loss': 0.0058, 'learning_rate': 1.662921348314607e-05, 'epoch': 0.76}                                                                                  
# {'embedding_loss': 0.0026, 'learning_rate': 1.101123595505618e-05, 'epoch': 1.52}                                                                                  
# {'embedding_loss': 0.0022, 'learning_rate': 5.393258426966292e-06, 'epoch': 2.27}                                                                                  
# {'train_runtime': 36.6756, 'train_samples_per_second': 172.758, 'train_steps_per_second': 5.399, 'epoch': 3.0}                                                     
# 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 198/198 [00:30<00:00,  6.45it/s] 
# The `max_length` is `None`. Using the maximum acceptable length according to the current model body: 512.
# Epoch: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 8/8 [00:07<00:00,  1.03it/s]

# Evaluating
metrics = trainer.evaluate(test_dataset)
print(metrics)
# => {'accuracy': 0.8632619439868204}

# Performing inference
preds = model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
print(preds)
# => tensor([1, 0, 0], device='cuda:0')

自定義分類頭

除了兩個內建選項外，SetFit 還允許您指定自定義分類頭。支援兩種形式的頭：自定義**可微分**頭或自定義**不可微分**頭。這兩種頭都必須實現以下兩個方法

自定義可微分頭

自定義可微分頭必須滿足以下要求

必須是 nn.Module 的子類。
一個 predict 方法：(self, torch.Tensor with shape [num_inputs, embedding_size]) -> torch.Tensor with shape [num_inputs] - 此方法對嵌入進行分類。輸出必須是 [0, num_classes) 範圍內的整數。
一個 predict_proba 方法：(self, torch.Tensor with shape [num_inputs, embedding_size]) -> torch.Tensor with shape [num_inputs, num_classes] - 此方法將嵌入分類為每個類別的機率。對於每個輸入，大小為 num_classes 的張量必須總和為 1。應用 torch.argmax(output, dim=-1) 應該得到 predict 的輸出。
一個 get_loss_fn 方法：(self) -> nn.Module - 返回一個已初始化的損失函式，例如 torch.nn.CrossEntropyLoss()。
一個 forward 方法：(self, Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor] - 給定 Sentence Transformer 主體的輸出，即包含 'input_ids'、'token_type_ids'、'attention_mask'、'token_embeddings' 和 'sentence_embedding' 鍵的字典，返回一個包含 'logits' 鍵和形狀為 [batch_size, num_classes] 的 torch.Tensor 值的字典。

自定義不可微分頭

自定義不可微分頭必須滿足以下要求

一個 predict 方法：(self, np.array with shape [num_inputs, embedding_size]) -> np.array with shape [num_inputs] - 此方法對嵌入進行分類。輸出必須是 [0, num_classes) 範圍內的整數。
一個 predict_proba 方法：(self, np.array with shape [num_inputs, embedding_size]) -> np.array with shape [num_inputs, num_classes] - 此方法將嵌入分類為每個類別的機率。對於每個輸入，大小為 num_classes 的陣列必須總和為 1。應用 np.argmax(output, dim=-1) 應該得到 predict 的輸出。
一個 fit 方法：(self, np.array with shape [num_inputs, embedding_size], List[Any]) -> None - 此方法必須接受一個嵌入的 numpy 陣列和相應的標籤列表。標籤本身不一定是整數。

許多來自 sklearn 的分類器已經滿足這些要求，例如 RandomForestClassifier、MLPClassifier、KNeighborsClassifier 等。

使用您的自定義（非）可微分分類頭初始化 SetFit 模型時，建議使用常規的 __init__ 方法

from setfit import SetFitModel
from sklearn.svm import LinearSVC
from sentence_transformers import SentenceTransformer

# Initializing a new SetFit model
model_body = SentenceTransformer("BAAI/bge-small-en-v1.5")
model_head = LinearSVC()
model = SetFitModel(model_body, model_head)

然後，訓練和推理可以正常進行，例如

from setfit import Trainer, TrainingArguments, sample_dataset
from datasets import load_dataset

# Preparing the dataset
dataset = load_dataset("SetFit/sst2")
train_dataset = sample_dataset(dataset["train"], label_column="label", num_samples=32)
test_dataset = dataset["test"]

# Preparing the training arguments
args = TrainingArguments(
    batch_size=32,
    num_epochs=3,
)

# Preparing the trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()

# Evaluating
metrics = trainer.evaluate(test_dataset)
print(metrics)
# => {'accuracy': 0.8638110928061504}

# Performing inference
preds = model.predict([
    "It's a charming and often affecting journey.",
    "It's slow -- very, very slow.",
    "A sometimes tedious film.",
])
print(preds)
# => tensor([1, 0, 0], dtype=torch.int32)

< > 在 GitHub 上更新