Trainer

SentenceTransformerTrainer

class sentence_transformers.trainer.SentenceTransformerTrainer(model: SentenceTransformer | None = None, args: SentenceTransformerTrainingArguments = 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)[source]

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

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

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

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

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

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

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

参数:

重要属性

  • 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 或 [~transformers.TrainerCallback]) – [~transformers.TrainerCallback] 类或 [~transformers.TrainerCallback] 的实例。在第一种情况下，将实例化该类的成员。

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

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

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

参数:

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

注意

此方法可以被训练器的子类覆盖，以在自定义用例中删除/自定义此回调

static add_prompts_or_dataset_name_transform(batch: dict[str, list[Any]], prompts: dict[str, str] | str | None = None, prompt_lengths: dict[str, int] | int | None = None, dataset_name: str | None = None, transform: Callable[[dict[str, list[Any]]], dict[str, list[Any]]] | None = None, **kwargs) → dict[str, list[Any]][source]

一个转换/映射函数，用于将 prompts 或数据集名称添加到批次。

参数:

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

  • prompts (dict[str, str] | str | None, optional) – 字符串提示的可选列名映射，或用于所有列的字符串提示。默认为 None。

  • prompt_lengths (dict[str, int] | int | None, optional) – 提示名称到提示 token 长度的可选映射，或者当提示是字符串时的提示 token 长度。默认为 None。

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

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

返回值:

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

返回类型:

dict[str, list[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]][source]

计算 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, optional) – 是否返回损失以及输出。默认为 False。

  • num_items_in_batch (int, optional) – 批次中的项目数。默认为 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[source]

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

参数:

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

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

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

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

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

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

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

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

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

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)

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

参数:

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

运行评估并返回指标。

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

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

参数:

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

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

    <Tip>

    如果您传递一个字典，其中数据集名称作为键，数据集作为值，则 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 上的损失。

    </Tip>

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

  • metric_key_prefix (str, optional, defaults to “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) → BatchSampler | None[source]

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

注意

覆盖此方法以提供自定义批次采样器。

参数:

  • dataset (Dataset) – 要从中采样的数据集。

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

  • drop_last (bool) – 如果为 True，并且数据集大小不能被批次大小整除，则删除最后一个不完整的批次。

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

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

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

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

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

参数:

eval_dataset (torch.utils.data.Dataset, optional) – 如果提供，将覆盖 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[source]

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

注意

覆盖此方法以提供自定义多数据集批次采样器。

参数:

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

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

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

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

get_num_trainable_parameters()

获取可训练参数的数量。

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

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

参数:

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

get_test_dataloader(test_dataset: Dataset | DatasetDict | IterableDataset) → DataLoader[source]

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

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

参数:

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

get_train_dataloader() → DataLoader[source]

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

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

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

使用 optuna 或 Ray Tune 或 SigOpt 启动超参数搜索。优化的量由 compute_objective 确定，当没有提供指标时，默认函数返回评估损失；否则，返回所有指标的总和。

<Tip warning={true}>

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

</Tip>

参数:

  • hp_space (Callable[[“optuna.Trial”], Dict[str, float]], optional) – 定义超参数搜索空间的函数。将默认为 [~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], optional) – 从 evaluate 方法返回的指标计算要最小化或最大化的目标的函数。将默认为 [~trainer_utils.default_compute_objective]。

  • n_trials (int, optional, defaults to 100) – 要测试的试验运行次数。

  • direction (str 或 List[str], optional, defaults to “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]], 可选) – 定义 trial/run 名称的函数。默认为 None。

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

    传递给 optuna.create_study 或 ray.tune.run 的额外关键字参数。更多信息请参阅

返回值:

所有关于多目标优化的最佳运行的信息。实验总结可以在 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]) → None

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

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

参数:

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

maybe_add_prompts_or_dataset_name_column(dataset_dict: DatasetDict | Dataset | None, prompts: dict[str, dict[str, str]] | dict[str, str] | str | None = None, dataset_name: str | None = None) → DatasetDict | Dataset | None[source]

可能向数据集添加 prompts 或数据集名称。如果满足以下条件，我们会向数据集添加 dataset_name 列：

  1. 损失是一个字典，并且数据集是一个 DatasetDict，或者

  2. prompts 包含到数据集名称的映射。

prompts 有 4 种情况

  1. str：所有数据集和列的单个 prompt。

  2. dict[str, str]：列到 prompt 的映射。

  3. dict[str, str]：数据集到 prompt 的映射。

  4. dict[str, dict[str, str]]：数据集到列到 prompt 的映射。

数据集有 2 种情况

  1. Dataset：单个数据集。

  2. DatasetDict：数据集字典。

不允许 3A，并且 2A 没有意义。

参数:

dataset_dict (DatasetDict | Dataset | None) – 要向其添加 prompts 或数据集名称的数据集。

返回值:

添加了 prompts 或数据集名称的数据集。

返回类型:

DatasetDict | Dataset | None

pop_callback(callback)

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

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

参数:

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

返回值:

移除的 callback，如果找到的话。

返回类型:

[~transformers.TrainerCallback]

propagate_args_to_deepspeed(auto_find_batch_size=False)

根据 Trainer args 设置 deepspeed 插件中的值

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

将 self.model 和 self.tokenizer 上传到 🤗 模型 hub 中 self.args.hub_model_id 指定的仓库。

参数:

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

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

  • token (str，可选，默认为 None) – 具有写入权限以覆盖 Trainer 原始 args 的 Token。

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

返回值:

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

remove_callback(callback)

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

参数:

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

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

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

仅从主进程保存。

train(resume_from_checkpoint: bool | str | None = None, trial: optuna.Trial | Dict[str, Any] = 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]，可选) – trial 运行或用于超参数搜索的超参数字典。

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

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