分类头

任何 🤗 SetFit 模型都由两部分组成：SentenceTransformer 嵌入主体和一个分类头。

本指南将向您展示

内置的逻辑回归分类头
内置的可微分分类头
自定义分类头的要求

逻辑回归分类头

当初始化新的 SetFit 模型时，默认选择 scikit-learn 逻辑回归头。这已被证明在应用于微调的 sentence transformer 主体之上时非常有效，并且仍然是推荐的分类头。使用逻辑回归头初始化新的 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，这些 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，形状为 [num_inputs, embedding_size]）-> torch.Tensor，形状为 [num_inputs] - 此方法对嵌入进行分类。输出必须是 [0, num_classes) 范围内的整数。
predict_proba 方法：（self，torch.Tensor，形状为 [num_inputs, embedding_size]）-> torch.Tensor，形状为 [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' 键和 torch.Tensor 值的字典，形状为 [batch_size, num_classes]。

自定义不可微分头

自定义不可微分头必须遵循以下要求

predict 方法：（self，np.array，形状为 [num_inputs, embedding_size]）-> np.array，形状为 [num_inputs] - 此方法对嵌入进行分类。输出必须是 [0, num_classes) 范围内的整数。
predict_proba 方法：（self，np.array，形状为 [num_inputs, embedding_size]）-> np.array，形状为 [num_inputs, num_classes] - 此方法将嵌入分类为每个类别的概率。对于每个输入，大小为 num_classes 的数组之和必须为 1。应用 np.argmax(output, dim=-1) 应产生 predict 的输出。
fit 方法：（self，np.array，形状为 [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 上更新