训练循环

class TrainingLoop(model: Model, triples_factory: CoreTriplesFactory, optimizer: str | Optimizer | type[Optimizer] | None = None, optimizer_kwargs: Mapping[str, Any] | None = None, lr_scheduler: str | LRScheduler | type[LRScheduler] | None = None, lr_scheduler_kwargs: Mapping[str, Any] | None = None, automatic_memory_optimization: bool = True, mode: Literal['training', 'validation', 'testing'] | None = None, result_tracker: str | ResultTracker | type[ResultTracker] | None = None, result_tracker_kwargs: Mapping[str, Any] | None = None)[源代码]

基础类: Generic[SampleType, BatchType], ABC

一个训练循环。

初始化训练循环。

Parameters:
  • model (Model) – 要训练的模型

  • triples_factory (CoreTriplesFactory) – 训练三元组工厂

  • optimizer (Optimizer) – 训练模型时使用的优化器

  • optimizer_kwargs (Mapping[str, Any] | None) – 用于实例化优化器的额外基于关键字的参数(如果需要)。params 将根据 model 自动添加。

  • lr_scheduler (LRScheduler | None) – 在训练模型时你想要使用的学习率调度器

  • lr_scheduler_kwargs (Mapping[str, Any] | None) – 用于实例化LR调度器的额外基于关键字的参数(如果需要)。optimizer 将会自动添加。

  • automatic_memory_optimization (bool) – bool 是否在训练期间自动优化子批次大小,并在评估期间根据当前硬件自动优化批次大小。

  • mode (Literal['training', 'validation', 'testing'] | None) – 归纳训练模式。如果是传导式的则为None。

  • result_tracker (str | ResultTracker | type[ResultTracker] | None) – 结果跟踪器

  • result_tracker_kwargs (Mapping[str, Any] | None) – 用于实例化结果跟踪器的额外基于关键字的参数

属性摘要

checksum

训练循环配置的模型和优化器的校验和。

device

模型使用的设备。

hpo_default

loss

模型使用的损失。

supports_slicing

方法总结

batch_size_search(*, triples_factory[, ...])

在当前设置下找到训练的最大批量大小。

get_normalized_name()

获取训练循环的规范化名称。

sub_batch_and_slice(*, batch_size, sampler, ...)

检查是否需要对子批次和/或切片进行处理,以便在现有硬件上训练模型。

train(triples_factory[, num_epochs, ...])

训练KGE模型。

属性文档

checksum

训练循环配置的模型和优化器的校验和。

device

模型使用的设备。

hpo_default = {'batch_size': {'high': 12, 'low': 4, 'scale': 'power_two', 'type': <class 'int'>}, 'num_epochs': {'high': 1000, 'low': 100, 'q': 100, 'type': <class 'int'>}}
loss

模型使用的损失。

supports_slicing: ClassVar[bool] = False

方法文档

在当前设置下找到训练的最大批量大小。

此方法检查当前模型在给定训练数据和现有硬件条件下可以使用的批量大小。如果可能,该方法将输出确定的批量大小和一个布尔值,表示该批量大小已成功评估。否则,输出将是批量大小1,布尔值将为False。

Parameters:
  • triples_factory (CoreTriplesFactory) – 运行搜索的三元组工厂

  • batch_size (int | None) – 开始搜索时的批量大小。如果为None,则设置batch_size=num_triples(即全批量训练)。

Returns:

包含最大可能批量大小的元组,以及一个指示器,指示使用该大小的评估是否成功。

Raises:

RuntimeError – 如果在训练期间引发运行时错误

Return type:

tuple[int, bool]

classmethod get_normalized_name() str[source]

获取训练循环的规范化名称。

Return type:

str

sub_batch_and_slice(*, batch_size: int, sampler: str | None, triples_factory: CoreTriplesFactory) tuple[int, int | None][source]

检查是否需要对子批次和/或切片进行处理,以便在现有硬件上训练模型。

Parameters:
Return type:

tuple[int, int | None]

train(triples_factory: CoreTriplesFactory, num_epochs: int = 1, batch_size: int | None = None, slice_size: int | None = None, label_smoothing: float = 0.0, sampler: str | None = None, continue_training: bool = False, only_size_probing: bool = False, use_tqdm: bool = True, use_tqdm_batch: bool = True, tqdm_kwargs: Mapping[str, Any] | None = None, stopper: Stopper | None = None, sub_batch_size: int | None = None, num_workers: int | None = None, clear_optimizer: bool = False, checkpoint_directory: None | str | Path = None, checkpoint_name: str | None = None, checkpoint_frequency: int | None = None, checkpoint_on_failure: bool = False, drop_last: bool | None = None, callbacks: str | TrainingCallback | type[TrainingCallback] | None | Sequence[str | TrainingCallback | type[TrainingCallback] | None] = None, callbacks_kwargs: Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None] = None, gradient_clipping_max_norm: float | None = None, gradient_clipping_norm_type: float | None = None, gradient_clipping_max_abs_value: float | None = None, pin_memory: bool = True) list[float] | None[来源]

训练KGE模型。

注意

梯度裁剪是一种避免梯度爆炸问题的技术。按范数裁剪和按值裁剪是两种替代实现。

Parameters:
  • triples_factory (CoreTriplesFactory) – 训练三元组。

  • num_epochs (int) – 训练模型的轮数。

  • batch_size (int | None) – 如果设置了,则用于小批量训练的批量大小。否则自动找到最大可能的batch_size。

  • slice_size (int | None) – >0 在使用切片时,评分函数的除数。这通常仅适用于LCWA训练循环,并且仅适用于已实现切片功能的模型。

  • label_smoothing (float) – (0 <= label_smoothing < 1) 如果大于零,则使用标签平滑。

  • sampler (str | None) – (None 或 ‘schlichtkrull’) 使用的采样器类型。目前,R-GCN中的sLCWA是唯一使用schlichtkrull采样的用户。

  • continue_training (bool) – 如果设置为False,(重新)初始化模型的权重。否则继续训练。

  • only_size_probing (bool) – 评估仅针对两个批次进行,以测试内存占用,特别是在GPU上。

  • use_tqdm (bool) – 是否应该为每个epoch显示进度条?

  • use_tqdm_batch (bool) – 是否应为批处理显示进度条(在epoch进度条内部)?

  • tqdm_kwargs (Mapping[str, Any] | None) – 传递给 tqdm 的关键字参数,用于管理进度条。

  • stopper (Stopper | None) – 一个pykeen.stopper.EarlyStopper的实例,用于检查是否应提前停止训练

  • sub_batch_size (int | None) – 如果提供,将每个批次分割成子批次,以避免大型模型/小型GPU的内存问题。

  • num_workers (int | None) – 用于加载数据的子CPU工作进程的数量。如果为None,数据将在主进程中加载。

  • clear_optimizer (bool) – 是否在训练后删除优化器实例(因为优化器可能会由于例如Adam中的动量而增加额外的内存消耗)。

  • checkpoint_directory (None | str | Path) – 一个可选的目录,用于存储检查点文件。如果为None,则使用由pykeen.constants.PYKEEN_HOME定义的目录中名为checkpoints的子目录。除非环境变量PYKEEN_HOME被覆盖,否则这将是~/.pykeen/checkpoints

  • checkpoint_name (str | None) – 用于保存检查点的文件名。如果给定的文件名已经存在,将加载并使用该文件以继续训练。

  • checkpoint_frequency (int | None) – 保存检查点的频率,单位为分钟。设置为0将在每个epoch后保存一个检查点。

  • checkpoint_on_failure (bool) – 是否在发生 RuntimeError 或 MemoryError 时保存检查点。此选项与普通检查点不同,因为普通检查点仅在成功完成一个周期后保存。当由于训练循环失败而保存检查点时,无法保证所有随机状态都能正确恢复,这可能会导致该特定训练循环的可重复性出现问题。因此,这些检查点将以不同的检查点名称保存,该名称将在给定的 checkpoint_root 中为 PyKEEN_just_saved_my_day_{datetime}.pt

  • drop_last (bool | None) – 是否在每个epoch中丢弃最后一个批次以防止较小的批次。默认为False,除非模型包含批量归一化层。可以显式提供以覆盖。

  • callbacks (str | TrainingCallback | type[TrainingCallback] | None | Sequence[str | TrainingCallback | type[TrainingCallback] | None]) – 一个可选的 pykeen.training.TrainingCallback 或回调实例的集合,定义了 几种功能之一。它们的接口受到了Keras的启发。

  • callbacks_kwargs (Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None]) – 用于实例化训练回调的额外基于关键字的参数。

  • gradient_clipping_max_norm (float | None) – 用于梯度裁剪的最大梯度范数。如果为None,则不使用梯度范数裁剪。

  • gradient_clipping_norm_type (float | None) – 用于最大梯度范数的梯度范数类型,参见 torch.nn.utils.clip_grad_norm_()

  • gradient_clipping_max_abs_value (float | None) – 梯度中的最大绝对值,参见 torch.nn.utils.clip_grad_value_()。如果为 None,则不使用梯度裁剪。

  • pin_memory (bool) – 是否在数据加载器中使用内存固定,参见 https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning

Returns:

每个时期的损失。

Return type:

list[float] | 无