分类头

任何 🤗 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 上更新