训练器

SparseEncoderTrainer

class sentence_transformers.sparse_encoder.SparseEncoderTrainer(model: SparseEncoder | None = None, args: SparseEncoderTrainingArguments | None = None, train_dataset: Dataset | DatasetDict | dict[str, Dataset] | None = None, eval_dataset: Dataset | DatasetDict | dict[str, Dataset] | None = None, loss: Module | dict[str, Module] | Callable[[SparseEncoder], Module] | dict[str, Callable[[SparseEncoder], Module]] | None = None, evaluator: SentenceEvaluator | list[SentenceEvaluator] | None = None, data_collator: SparseEncoderDataCollator | None = None, tokenizer: PreTrainedTokenizerBase | Callable | None = None, model_init: Callable[[], SparseEncoder] | None = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, optimizers: tuple[Optimizer, LambdaLR] = (None, None), preprocess_logits_for_metrics: Callable[[Tensor, Tensor], Tensor] | None = None)

SparseEncoderTrainer 是一个简单但功能完整的 PyTorch 训练和评估循环，基于 SentenceTransformerTrainer，而 SentenceTransformerTrainer 又基于 🤗 Transformers Trainer。

此训练器集成了对各种 transformers.TrainerCallback 子类的支持，例如

• 如果安装了 wandb，则 WandbCallback 会自动将训练指标记录到 W&B

• 如果可以访问 tensorboard，则 TensorBoardCallback 会将训练指标记录到 TensorBoard。

• 如果安装了 codecarbon，则 CodeCarbonCallback 会在训练期间跟踪模型的碳排放。

  • 注意：这些碳排放将包含在自动生成的模型卡中。

有关集成回调以及如何编写自己的回调的更多信息，请参阅 Transformers Callbacks 文档。

参数:

• model (SparseEncoder, 可选) – 要训练、评估或用于预测的模型。如果未提供，则必须传递 model_init。

• args (SparseEncoderTrainingArguments, 可选) – 用于调整训练的参数。如果未提供，则默认为 SparseEncoderTrainingArguments 的基本实例，其中 output_dir 设置为当前目录中名为 tmp_trainer 的目录。

• train_dataset (Union[datasets.Dataset, datasets.DatasetDict, datasets.IterableDataset, Dict[str, datasets.Dataset]], 可选) – 用于训练的数据集。必须具有损失函数接受的格式，请参阅训练概述 > 数据集格式。

• eval_dataset (Union[datasets.Dataset, datasets.DatasetDict, datasets.IterableDataset, Dict[str, datasets.Dataset]], 可选) –

  用于评估的数据集。必须具有损失函数接受的格式，请参阅训练概述 > 数据集格式。

• loss (Optional[Union[torch.nn.Module, Dict[str, torch.nn.Module], Callable[[SparseEncoder], torch.nn.Module], Dict[str, Callable[[SparseEncoder]]]], 可选) – 用于训练的损失函数。可以是损失类实例，也可以是将数据集名称映射到损失类实例的字典，或者是给定模型返回损失类实例的函数，或者是将数据集名称映射到给定模型返回损失类实例的函数的字典。实际上，后两者主要用于超参数优化。如果未提供 loss，则默认为 SparseMultipleNegativesRankingLoss。

• evaluator (Union[SentenceEvaluator, List[SentenceEvaluator]], 可选) – 用于训练期间有用的评估指标的评估器实例。您可以带或不带 eval_dataset 使用 evaluator，反之亦然。通常，evaluator 返回的指标比 eval_dataset 返回的损失值更有用。评估器列表将被封装在 SequentialEvaluator 中以按顺序运行。

• callbacks (List of [transformers.TrainerCallback], 可选) –

  用于自定义训练循环的回调列表。这些回调将被添加到[此处](callback)详细说明的默认回调列表中。

  如果要移除使用的默认回调之一，请使用 [Trainer.remove_callback] 方法。

• optimizers (Tuple[:class:`torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, 可选, 默认为 (None, None)) – 包含要使用的优化器和调度器的元组。默认为模型上的 torch.optim.AdamW 实例和由 args 控制的 transformers.get_linear_schedule_with_warmup() 给出的调度器。

重要属性

• model – 始终指向核心模型。如果使用 Transformers 模型，它将是 [PreTrainedModel] 子类。

• model_wrapped – 在一个或多个其他模块封装原始模型的情况下，始终指向最外部的模型。这是应该用于前向传播的模型。例如，在 DeepSpeed 下，内部模型首先被 DeepSpeed 封装，然后再次被 torch.nn.DistributedDataParallel 封装。如果内部模型尚未被封装，则 self.model_wrapped 与 self.model 相同。

• is_model_parallel – 模型是否已切换到模型并行模式（不同于数据并行，这意味着某些模型层分布在不同的 GPU 上）。

• place_model_on_device – 是否自动将模型放置在设备上 - 如果使用模型并行或 DeepSpeed，或者如果默认的 TrainingArguments.place_model_on_device 被覆盖为返回 False，则它将设置为 False。

• is_in_train – 模型当前是否正在运行 train（例如，当在 train 中调用 evaluate 时）

add_callback(callback)

将回调添加到当前 [~transformers.TrainerCallback] 列表中。

参数:

callback (type or [~transformers.TrainerCallback]) – 一个 [~transformers.TrainerCallback] 类或 [~transformers.TrainerCallback] 的实例。在第一种情况下，将实例化该类的一个成员。

static add_dataset_name_transform(batch: dict[str, list[Any]], dataset_name: str | None = None, transform: Callable[[dict[str, list[Any]]], dict[str, list[Any]]] | None = None, **kwargs) → dict[str, list[Any]]

添加提示或数据集名称到批次的转换/映射函数。

参数:

• batch (dict[str, list[Any]]) – 数据批次，其中每个键是列名，每个值是值的列表。

• dataset_name (str | None, 可选) – 此数据集的名称，仅当存在多个使用不同损失的数据集时。默认为 None。

• transform (Callable[[dict[str, list[Any]]], dict[str, list[Any]]], 可选) – 在添加提示等之前应用于批次的可选转换函数。默认为 None。

返回:

添加了提示和/或数据集名称的“即时”转换批次。

返回类型:

dict[str, list[Any]]

add_model_card_callback(default_args_dict: dict[str, Any]) → None

添加负责自动跟踪模型卡自动生成所需数据的回调

此方法在 SparseEncoderTrainer 类的 __init__ 方法中调用。

参数:

default_args_dict (Dict[str, Any]) – 默认训练参数的字典，以便我们可以确定哪些参数已为模型卡更改。

compute_loss(model: SparseEncoder, inputs: dict[str, Tensor | Any], return_outputs: bool = False, num_items_in_batch=None) → Tensor | tuple[Tensor, dict[str, Any]]

计算 SparseEncoder 模型的损失。

它使用 self.loss 来计算损失，可以是单个损失函数，也可以是不同数据集的损失函数字典。如果损失是字典，则预期数据集名称在输入中以“dataset_name”键传递。这在 add_dataset_name_column 方法中自动完成。请注意，即使 return_outputs = True，输出也将为空，因为 SparseEncoder 损失不返回输出。

参数:

• model (SparseEncoder) – SparseEncoder 模型。

• inputs (Dict[str, Union[torch.Tensor, Any]]) – 模型的输入数据。

• return_outputs (bool, 可选) – 是否随损失一起返回输出。默认为 False。

• num_items_in_batch (int, 可选) – 批次中的项目数量。默认为 None。未使用，但 Transformers 训练器需要。

返回:

计算出的损失。如果 return_outputs 为 True，则返回损失和输出的元组。否则，仅返回损失。

返回类型:

Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, Any]]]

create_model_card(language: str | None = None, license: str | None = None, tags: str | list[str] | None = None, model_name: str | None = None, finetuned_from: str | None = None, tasks: str | list] | None = None, dataset_tags: str | list[str] | None = None, dataset: str | list[str] | None = None, dataset_args: str | list[str] | None = None, **kwargs) → None

使用 Trainer 可用的信息创建模型卡草稿。

参数:

• language (str, 可选) – 模型的语言（如果适用）

• license (str, 可选) – 模型的许可证。如果提供给 Trainer 的原始模型来自 Hub 上的仓库，则默认为所使用的预训练模型的许可证。

• tags (str or List[str], 可选) – 要包含在模型卡元数据中的一些标签。

• model_name (str, 可选) – 模型的名称。

• finetuned_from (str, 可选) – 用于微调此模型的模型名称（如果适用）。默认为提供给 Trainer 的原始模型的仓库名称（如果它来自 Hub）。

• tasks (str or List[str], 可选) – 一个或多个任务标识符，包含在模型卡元数据中。

• dataset_tags (str or List[str], 可选) – 一个或多个数据集标签，包含在模型卡元数据中。

• dataset (str or List[str], 可选) – 一个或多个数据集标识符，包含在模型卡元数据中。

• dataset_args (str or List[str], 可选) – 一个或多个数据集参数，包含在模型卡元数据中。

create_optimizer()

设置优化器。

我们提供了一个效果良好的合理默认值。如果您想使用其他优化器，可以通过 optimizers 在 Trainer 的 init 中传递一个元组，或者在子类中重写此方法。

create_optimizer_and_scheduler(num_training_steps: int)

设置优化器和学习率调度器。

我们提供了一个效果良好的合理默认值。如果您想使用其他优化器，可以通过 optimizers 在 Trainer 的 init 中传递一个元组，或者在子类中重写此方法（或 create_optimizer 和/或 create_scheduler）。

create_scheduler(num_training_steps: int, optimizer: Optimizer | None = None)

设置调度器。训练器的优化器必须在此方法调用之前或作为参数传递时设置好。

参数:

num_training_steps (int) – 要执行的训练步骤数。

evaluate(eval_dataset: Dataset | dict[str, Dataset] | None = None, ignore_keys: list[str] | None = None, metric_key_prefix: str = 'eval') → dict[str, float]

运行评估并返回指标。

调用脚本将负责提供一个计算指标的方法，因为它们是任务相关的（将其传递给初始化 compute_metrics 参数）。

您也可以通过子类化并重写此方法来注入自定义行为。

参数:

• eval_dataset (Union[Dataset, Dict[str, Dataset]), 可选) –

  如果您希望覆盖 self.eval_dataset，请传递一个数据集。如果它是 [~datasets.Dataset]，则模型 model.forward() 方法不接受的列将自动删除。它必须实现 __len__ 方法。

  <Tip>

  如果您传入一个字典，其中键是数据集名称，值是数据集，则评估将在每个数据集上单独运行。这对于监控训练如何影响其他数据集或简单地获得更细粒度的评估很有用。当与 load_best_model_at_end 一起使用时，请确保 metric_for_best_model 明确引用其中一个数据集。例如，如果您为两个数据集 data1 和 data2 传入 {“data1”: data1, “data2”: data2}，您可以使用 metric_for_best_model=”eval_data1_loss” 来使用 data1 上的损失，并使用 metric_for_best_model=”eval_data2_loss” 来使用 data2 上的损失。

  </Tip>

• ignore_keys (List[str], 可选) – 模型输出中（如果它是字典）在收集预测时应忽略的键列表。

• metric_key_prefix (str, 可选, 默认为 “eval”) – 用作指标键前缀的可选前缀。例如，如果前缀是“eval”（默认），则指标“bleu”将命名为“eval_bleu”

返回:

包含评估损失和从预测中计算出的潜在指标的字典。该字典还包含来自训练状态的 epoch 号。

get_batch_sampler(dataset: Dataset, batch_size: int, drop_last: bool, valid_label_columns: list[str] | None = None, generator: Generator | None = None, seed: int = 0) → BatchSampler | None

根据 self.args 中的 batch_sampler 参数返回相应的批采样器。此批采样器类支持 __len__ 和 __iter__ 方法，并用作 torch.utils.data.DataLoader 的 batch_sampler。

注意

重写此方法以提供自定义批采样器。

参数:

• dataset (Dataset) – 要采样的Dataset。

• batch_size (int) – 每个批次的样本数。

• drop_last (bool) – 如果为 True，当数据集大小不能被批次大小整除时，丢弃最后一个不完整的批次。

• valid_label_columns (List[str]) – 要检查标签的列名列表。在数据集中找到的第一个 valid_label_columns 中的列名将用作标签列。

• generator (torch.Generator, 可选) – 用于打乱索引的可选随机数生成器。

• seed (int) – 随机数生成器的种子，以确保可重现性。默认为 0。

get_eval_dataloader(eval_dataset: Dataset | DatasetDict | IterableDataset | None = None) → DataLoader

返回评估 [~torch.utils.data.DataLoader]。

如果您想注入一些自定义行为，请子类化并重写此方法。

参数:

eval_dataset (torch.utils.data.Dataset, 可选) – 如果提供，将覆盖 self.eval_dataset。如果它是 [~datasets.Dataset]，则模型 model.forward() 方法不接受的列将自动删除。它必须实现 __len__ 方法。

get_learning_rates()

返回 self.optimizer 中每个参数的学习率。

get_multi_dataset_batch_sampler(dataset: ConcatDataset, batch_samplers: list[BatchSampler], generator: Generator | None = None, seed: int | None = 0) → BatchSampler

根据 self.args 中的 multi_dataset_batch_sampler 参数返回相应的多数据集批采样器。此批采样器类支持 __len__ 和 __iter__ 方法，并用作 torch.utils.data.DataLoader 的 batch_sampler。

注意

重写此方法以提供自定义的多数据集批采样器。

参数:

• dataset (ConcatDataset) – 所有数据集的串联。

• batch_samplers (List[BatchSampler]) – 串联数据集中每个数据集的批采样器列表。

• generator (torch.Generator, 可选) – 用于打乱索引的可选随机数生成器。

• seed (int, 可选) – 随机数生成器的可选种子

get_num_trainable_parameters()

获取可训练参数的数量。

get_optimizer_group(param: str | Parameter | None = None)

如果给定参数，则返回该参数的优化器组；否则，返回所有参数的优化器组。

参数:

param (str 或 torch.nn.parameter.Parameter，可选) – 需要返回优化器组的参数。

get_test_dataloader(test_dataset: Dataset | DatasetDict | IterableDataset) → DataLoader

返回训练 [~torch.utils.data.DataLoader]。

如果您想注入一些自定义行为，请子类化并重写此方法。

参数:

test_dataset (torch.utils.data.Dataset，可选) – 要使用的测试数据集。如果它是 [~datasets.Dataset]，则 `model.forward()` 方法不接受的列将自动移除。它必须实现 `__len__`。

get_train_dataloader() → DataLoader

返回训练 [~torch.utils.data.DataLoader]。

如果 `train_dataset` 未实现 `__len__`，则不使用采样器；否则，使用随机采样器（如有必要，会适应分布式训练）。

如果您想注入一些自定义行为，请子类化并重写此方法。

hyperparameter_search(hp_space: Callable[[optuna.Trial], dict[str, float]] | None = None, compute_objective: Callable[[dict[str, float]], float] | None = None, n_trials: int = 20, direction: str | list[str] = 'minimize', backend: str | HPSearchBackend | None = None, hp_name: Callable[[optuna.Trial], str] | None = None, **kwargs) → BestRun | list[BestRun]

使用 optuna、Ray Tune 或 SigOpt 启动超参数搜索。优化的数量由 compute_objective 决定，如果未提供度量，它默认返回评估损失，否则返回所有度量的总和。

<提示 warning={true}>

要使用此方法，您需要在初始化您的 [Trainer] 时提供 model_init：每次新运行时都需要重新初始化模型。这与 optimizers 参数不兼容，因此您需要子类化 [Trainer] 并重写方法 [~Trainer.create_optimizer_and_scheduler] 以自定义优化器/调度器。

</Tip>

参数:

• hp_space (Callable[[“optuna.Trial”], Dict[str, float]]，可选) – 定义超参数搜索空间的函数。将根据您的后端默认设置为 [~trainer_utils.default_hp_space_optuna]、[~trainer_utils.default_hp_space_ray] 或 [~trainer_utils.default_hp_space_sigopt]。

• compute_objective (Callable[[Dict[str, float]], float]，可选) – 一个函数，用于根据 evaluate 方法返回的指标计算要最小化或最大化的目标。将默认设置为 [~trainer_utils.default_compute_objective]。

• n_trials (int，可选，默认为 100) – 要测试的试验运行次数。

• direction (str 或 List[str]，可选，默认为 “minimize”) – 如果是单目标优化，`direction` 是 `str` 类型，可以是 “minimize” 或 “maximize”，在优化验证损失时应选择 “minimize”，在优化一个或多个指标时应选择 “maximize”。如果是多目标优化，`direction` 是 `List[str]` 类型，可以是 “minimize” 和 “maximize” 的列表，在优化验证损失时应选择 “minimize”，在优化一个或多个指标时应选择 “maximize”。

• backend (str 或 [~training_utils.HPSearchBackend]，可选) – 用于超参数搜索的后端。将根据安装情况默认使用 optuna、Ray Tune 或 SigOpt。如果都已安装，则默认使用 optuna。

• hp_name (Callable[[“optuna.Trial”], str]]，可选) – 定义试验/运行名称的函数。将默认为 None。

• kwargs (Dict[str, Any]，可选) –

  每个后端的额外关键字参数

  • optuna：来自 [optuna.study.create_study](https://docs.optuna.cn/en/stable/reference/generated/optuna.study.create_study.html) 的参数，以及来自 [optuna.study.Study.optimize](https://docs.optuna.cn/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize) 的 `timeout`、`n_jobs` 和 `gc_after_trial` 参数。

  • ray：来自 [tune.run](https://docs.rayai.org.cn/en/latest/tune/api_docs/execution.html#tune-run) 的参数。如果 `kwargs` 中未设置 `resources_per_trial`，则默认使用 1 个 CPU 核和 1 个 GPU（如果可用）。如果 `kwargs` 中未设置 `progress_reporter`，则使用 [ray.tune.CLIReporter](https://docs.rayai.org.cn/en/latest/tune/api/doc/ray.tune.CLIReporter.html)。

  • sigopt：来自 [sigopt.Connection.set_proxies](https://docs.sigopt.com/support/faq#how-do-i-use-sigopt-with-a-proxy) 的 `proxies` 参数。

返回:

关于最佳运行或多目标优化最佳运行的所有信息。实验摘要可在 Ray 后端的 run_summary 属性中找到。

返回类型:

[trainer_utils.BestRun 或 List[trainer_utils.BestRun]]

is_local_process_zero() → bool

此进程是否为本地（例如，如果在多台机器上以分布式方式训练，则指单台机器上的）主进程。

is_world_process_zero() → bool

此进程是否为全局主进程（当在多台机器上以分布式方式训练时，此值仅对一个进程为 True）。

log(logs: dict[str, float], start_time: float | None = None) → None

在监控训练的各种对象上记录 logs。

子类化并重写此方法以注入自定义行为。

参数:

• logs (Dict[str, float]) – 要记录的值。

• start_time (Optional[float]) – 训练开始时间。

maybe_add_dataset_name_column(dataset: DatasetDict | Dataset | None, prompts: dict[str, dict[str, str]] | dict[str, str] | str | None = None, router_mapping: dict[str, dict[str, str]] | dict[str, str] | None = None, dataset_name: str | None = None) → DatasetDict | Dataset | None

也许会向数据集添加数据集名称列，如果：

1. 数据集是 DatasetDict，并且满足以下条件之一：

   1. 损失是一个字典，或

   2. 提示包含数据集名称的映射，或

   3. router_mapping 包含数据集名称的映射。

参数:

dataset (DatasetDict | Dataset | None) – 要添加提示或数据集名称的数据集。

返回:

已添加提示或数据集名称的数据集。

返回类型:

DatasetDict | Dataset | None

pop_callback(callback)

从当前 [~transformers.TrainerCallback] 列表中移除一个回调并返回它。

如果未找到回调，则返回 None（且不引发错误）。

参数:

callback (type 或 [~transformers.TrainerCallback]) – 一个 [~transformers.TrainerCallback] 类或其实例。在第一种情况下，将弹出在回调列表中找到的该类的第一个成员。

返回:

如果找到，则为已移除的回调。

返回类型:

[~transformers.TrainerCallback]

preprocess_dataset(dataset: DatasetDict | Dataset | None = None, prompts: dict[str, dict[str, str]] | dict[str, str] | str | None = None, router_mapping: dict[str, dict[str, str]] | dict[str, str] | None = None, dataset_name: str | None = None) → DatasetDict | Dataset | None

预处理数据集，可选地延迟添加数据集名称列，这对于具有多个损失的多数据集训练、数据集特定提示或数据集特定路由映射是必需的。

参数:

• dataset (DatasetDict | Dataset | None) – 要预处理的数据集。如果为 None，则不进行预处理。

• prompts (dict[*str*, *dict*[*str*, *str*]] | *dict*[*str*, *str*] | *str* | *None*) – 要添加到数据集的可选提示。如果是一个字符串，则将其用作所有数据集的单一提示，但它也可以是一个将数据集名称映射到提示的字典，一个将列名映射到提示的字典，或者一个将数据集名称映射到列名再映射到提示的嵌套字典。

• router_mapping (dict[*str*, *dict*[*str*, *str*]] | *dict*[*str*, *str*] | *None*) – 要添加到数据集的可选路由映射。可以是列名到 [~sentence_transformers.models.Router] 路由的字典映射，也可以是数据集名称到列名再到路由的嵌套字典映射。

• dataset_name (str | None) – 数据集名称，用于具有多个损失的多数据集训练。

返回:

预处理后的数据集，可能已将数据集名称添加为惰性列。

返回类型:

DatasetDict | Dataset | None

propagate_args_to_deepspeed(auto_find_batch_size=False)

根据 Trainer 参数设置 deepspeed 插件中的值。

push_to_hub(commit_message: str | None = 'End of training', blocking: bool = True, token: str | None = None, revision: str | None = None, **kwargs) → str

将 self.model 和 self.processing_class 上传到 self.args.hub_model_id 仓库的 🤗 模型中心。

参数:

• commit_message (str，可选，默认为 “End of training”) – 推送时提交的消息。

• blocking (bool，可选，默认为 True) – 函数是否应仅在 `git push` 完成后返回。

• token (str，可选，默认为 None) – 具有写入权限的令牌，用于覆盖 Trainer 的原始参数。

• revision (str，可选) – 要提交的 git 修订版本。默认为“main”分支的头部。

• kwargs (Dict[str, Any]，可选) – 传递给 [~Trainer.create_model_card] 的额外关键字参数。

返回:

如果 blocking=False，则返回模型被推送到的仓库 URL；如果 blocking=True，则返回一个跟踪提交进度的 Future 对象。

remove_callback(callback)

从当前 [~transformers.TrainerCallback] 列表中移除一个回调。

参数:

callback (type 或 [~transformers.TrainerCallback]) – 一个 [~transformers.TrainerCallback] 类或其实例。在第一种情况下，将移除在回调列表中找到的该类的第一个成员。

save_model(output_dir: str | None = None, _internal_call: bool =False)

将保存模型，以便您可以使用 from_pretrained() 重新加载它。

仅从主进程保存。

set_initial_training_values(args: TrainingArguments, dataloader: DataLoader, total_train_batch_size: int)

计算并返回以下值：- num_train_epochs - num_update_steps_per_epoch - num_examples - num_train_samples - epoch_based - len_dataloader - max_steps

train(resume_from_checkpoint: bool | str | None = None, trial: optuna.Trial | dict[str, Any] | None = None, ignore_keys_for_eval: list[str] | None = None, **kwargs)

主要训练入口点。

参数:

• resume_from_checkpoint (str 或 bool，可选) – 如果是 `str`，则为由 [Trainer] 的先前实例保存的检查点的本地路径。如果是 `bool` 且等于 True，则加载由 [Trainer] 的先前实例保存的 *args.output_dir* 中的最后一个检查点。如果存在，训练将从此处加载的模型/优化器/调度器状态继续。

• trial (optuna.Trial 或 Dict[str, Any]，可选) – 用于超参数搜索的试验运行或超参数字典。

• ignore_keys_for_eval (List[str]，可选) – 模型输出（如果是一个字典）中应在训练期间收集评估预测时忽略的键列表。

• kwargs (Dict[str, Any]，可选) – 用于隐藏已弃用参数的额外关键字参数。