训练器

SentenceTransformerTrainer

class sentence_transformers.trainer.SentenceTransformerTrainer(model: SentenceTransformer | None = None, args: SentenceTransformerTrainingArguments | None = None, train_dataset: Dataset | DatasetDict | IterableDataset | dict[str, Dataset] | None = None, eval_dataset: Dataset | DatasetDict | IterableDataset | dict[str, Dataset] | None = None, loss: nn.Module | dict[str, nn.Module] | Callable[[SentenceTransformer], torch.nn.Module] | dict[str, Callable[[SentenceTransformer], torch.nn.Module]] | None = None, evaluator: SentenceEvaluator | list[SentenceEvaluator] | None = None, data_collator: DataCollator | None = None, tokenizer: PreTrainedTokenizerBase | Callable | None = None, model_init: Callable[[], SentenceTransformer] | None = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None)[源代码]

SentenceTransformerTrainer 是一个简单但功能齐全的 PyTorch 训练和评估循环,基于 🤗 Transformers Trainer

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

  • 如果安装了 wandbWandbCallback 可自动将训练指标记录到 W&B

  • 如果可以访问 tensorboardTensorBoardCallback 可将训练指标记录到 TensorBoard。

  • 如果安装了 codecarbonCodeCarbonCallback 可在训练期间跟踪模型的碳排放。

    • 注意:这些碳排放量将被包含在您自动生成的模型卡片中。

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

参数:

重要属性

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

  • model_wrapped – 始终指向最外层的模型,以防一个或多个其他模块包装了原始模型。这是应该用于前向传播的模型。例如,在 DeepSpeed 下,内部模型被包装在 DeepSpeed 中,然后再次被包装在 torch.nn.DistributedDataParallel 中。如果内部模型没有被包装,那么 self.model_wrappedself.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[源代码]

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

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

参数:

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

注意

可以通过子类化训练器来重写此方法,以在自定义用例中删除/自定义此回调

compute_loss(model: SentenceTransformer, inputs: dict[str

, torch.Tensor

| Any

], return_outputs: bool = False, num_items_in_batch=None) torch.Tensor

| tuple

[torch.Tensor

, dict

[str

, Any

]

][源代码]

计算 SentenceTransformer 模型的损失。

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

参数:
  • model (SentenceTransformer) – SentenceTransformer 模型。

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

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

  • num_items_in_batch (int, 可选) – 批次中的项目数。默认为 None。未使用,但 transformers Trainer 需要。

返回:

计算出的损失。如果 `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

[str

]

| 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()

设置优化器。

我们提供了一个效果不错的合理默认值。如果你想使用其他东西,可以在 Trainer 的 init 中通过 optimizers 传递一个元组,或者在子类中继承并重写此方法。

create_optimizer_and_scheduler(num_training_steps: int)

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

我们提供了一个效果不错的合理默认值。如果你想使用其他东西,可以在 Trainer 的 init 中通过 `optimizers` 传递一个元组,或者在子类中子类化并重写此方法(或 `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

][源代码]

运行评估并返回指标。

调用脚本将负责提供一个计算指标的方法,因为它们是任务相关的(将其传递给 init 的 compute_metrics 参数)。

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

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

    如果您希望覆盖 self.eval_dataset,请传递一个数据集。如果它是一个 [~datasets.Dataset],model.forward() 方法不接受的列将自动被移除。如果它是一个字典,它将在每个数据集上进行评估,并将字典键前置到指标名称。数据集必须实现 __len__ 方法。

    <提示>

    如果你传入一个以数据集名称为键、数据集为值的字典,evaluate 将在每个数据集上运行独立的评估。这对于监控训练如何影响其他数据集,或者仅仅为了获得更细粒度的评估非常有用。当与 `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` 上的损失。

    </提示>

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

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

返回:

一个包含评估损失和从预测中计算的潜在指标的字典。该字典还包含来自训练状态的轮次数。

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.DataLoaderbatch_sampler

注意

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

参数:
  • 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 or 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__,则不使用采样器;否则使用随机采样器(必要时适应分布式训练)。

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

使用 optunaRay TuneSigOpt 启动超参数搜索。优化的量由 compute_objective 确定,当没有提供度量时,它默认为一个返回评估损失的函数,否则返回所有度量的总和。

<提示 warning={true}>

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

</提示>

参数:
  • 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 or List[str], 可选, 默认为 “minimize”) – 如果是单目标优化,direction 是 str,可以是 “minimize”“maximize”,当优化验证损失时应选择 “minimize”,当优化一个或多个指标时应选择 “maximize”。如果是多目标优化,direction 是 List[str],可以是 “minimize”“maximize” 的列表,当优化验证损失时应选择 “minimize”,当优化一个或多个指标时应选择 “maximize”

  • backend (str or [~training_utils.HPSearchBackend], 可选) – 用于超参数搜索的后端。将根据已安装的 optuna、Ray Tune 或 SigOpt 默认为其中之一。如果都已安装,则默认为 optuna。

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

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

    每个后端的附加关键字参数

返回:

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

返回类型:

[trainer_utils.BestRunList[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[source]

如果满足以下条件,则可能向数据集中添加一个数据集名称列:

  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] 的实例。在第一种情况下,将弹出回调列表中找到的该类的第一个成员。

返回:

被移除的回调(如果找到)。

返回类型:

[~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[source]

通过惰性地添加数据集名称列来预处理数据集。这在多数据集训练、使用多个损失函数、针对特定数据集的提示词或特定数据集的路由器映射时是必需的。

参数:
  • 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) – (可选)要添加到数据集的路由器映射。可以是一个将列名映射到 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.modelself.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] 类或一个 [~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 (strbool, 可选) – 如果是 str,则为先前 [Trainer] 实例保存的检查点的本地路径。如果是 bool 且等于 True,则加载先前 [Trainer] 实例在 args.output_dir 中保存的最新检查点。如果存在此参数,训练将从此加载的模型/优化器/调度器状态恢复。

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

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

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