训练期间保存检查点

在使用具有许多参数或大数据集的模型时，训练可能需要数天到数周的时间。这引入了大量可能的错误，例如会话超时、服务器重启等，这将导致迄今为止所有进度的完全丢失。为了避免这种情况，PyKEEN 支持内置检查点，允许直接保存当前训练循环状态，并从保存的检查点恢复保存的状态，如常规检查点所示，以及在训练循环失败时保存的失败检查点，如失败时的检查点所示。要更详细地了解检查点的工作原理以及如何以编程方式使用它们，请查看超越管道和技术细节的检查点。为了修复可能的错误和安全回退，请也查看注意事项和可能的错误。

定期检查点

教程 First Steps 展示了如何使用 pykeen.pipeline.pipeline() 函数仅用两行代码设置整个KGEM进行训练和评估。下面展示了一个稍微扩展的示例：

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=1000,
...     ),
... )

要启用检查点，您只需向training_kwargs添加一个checkpoint_name参数。此参数应具有您希望保存在计算机上的检查点文件的名称。

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=1000,
...         checkpoint_name='my_checkpoint.pt',
...     ),
... )

此外，您可以通过设置参数checkpoint_frequency来设置检查点频率，即每隔多少分钟保存一次检查点。默认频率为30分钟，将其设置为0将导致训练循环在每个周期后保存一个检查点。让我们看一个例子。

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=1000,
...         checkpoint_name='my_checkpoint.pt',
...         checkpoint_frequency=5,
...     ),
... )

这里我们定义了一个管道，它将在每次一个周期结束时以及自上次保存以来至少经过5分钟时，在名为my_checkpoint.pt的检查点文件中保存训练循环检查点。假设这个管道在200个周期后崩溃，你可以简单地执行相同的代码，管道将从检查点文件加载最后的状态，并继续训练，就像什么都没发生一样。结果将与你连续运行管道1000个周期而不中断时完全相同。

另一个不错的功能是，使用检查点时，训练循环将在训练循环完成或早停器停止时保存状态。假设你已经成功地将上述KGEM训练了1000个周期，但现在决定要用2000个周期来测试模型，你只需要更改周期数并执行代码即可：

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,  # more epochs than before
...         checkpoint_name='my_checkpoint.pt',
...         checkpoint_frequency=5,
...     ),
... )

上述代码将在完成1000个周期后加载保存的状态，并继续训练到2000个周期，产生与一开始就运行2000个周期完全相同的结果。

默认情况下，您的检查点将保存在PYKEEN_HOME目录中，该目录在pykeen.constants中定义，这是您主目录中的一个子目录，例如~/.data/pykeen/checkpoints（通过pystow配置）。可选地，您可以通过设置checkpoint_directory参数来指定保存检查点的路径，该参数可以是一个字符串或一个pathlib.Path对象，包含您所需的根路径，如下例所示：

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,
...         checkpoint_name='my_checkpoint.pt',
...         checkpoint_directory='doctests/checkpoint_dir',
...     ),
... )

失败时的检查点

在你只想在训练循环可能失败时保存检查点的情况下，你可以使用参数 checkpoint_on_failure=True，例如：

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,
...         checkpoint_on_failure=True,
...     ),
... )

此选项与常规检查点不同，因为常规检查点仅在成功完成一个周期后保存。当由于训练循环失败而保存检查点时，无法保证所有随机状态都能正确恢复，这可能会导致该特定训练循环的可重复性出现问题。因此，这些检查点会以不同的检查点名称保存，即使在您也选择使用上述定义的常规检查点时，例如使用以下代码，这些检查点也会保存在给定的checkpoint_directory中，名称为PyKEEN_just_saved_my_day_{datetime}.pt。

>>> from pykeen.pipeline import pipeline
>>> pipeline_result = pipeline(
...     dataset='Nations',
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,
...         checkpoint_name='my_checkpoint.pt',
...         checkpoint_on_failure=True,
...     ),
... )

注意：请谨慎使用此参数，因为每次失败的训练循环都会创建一个单独的检查点文件。

自带数据时的检查点

在继续训练或恢复训练后使用模型时，确保实体标签到标识符（entity_to_id）和关系标签到标识符（relation_to_id）的映射与保存检查点时使用的映射相同是至关重要的。如果不相同，任何下游使用都将毫无意义。

如果您使用的是PyKEEN提供的数据集，您将自动受到保护。然而，当使用您自己的数据集时（请参阅Bring Your Own Data），您需要确保这一点。以下是两个典型的例子，展示了如何将您自己的数据与检查点结合使用。

恢复训练

以下示例展示了如何使用自定义三元组工厂来处理从包含标记三元组的文件派生的训练、验证和测试数据集。注意在创建validation和testing三元组工厂时如何使用entity_to_id和relation_to_id参数，以确保这些数据集与训练数据集使用相同的映射创建。由于checkpoint_name设置为'my_checkpoint.pt'，PyKEEN将检查点保存在~/.data/pykeen/checkpoints/my_checkpoint.pt。

>>> from pykeen.pipeline import pipeline
>>> from pykeen.triples import TriplesFactory
>>> from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, NATIONS_VALIDATE_PATH
>>> training = TriplesFactory.from_path(
...     path=NATIONS_TRAIN_PATH,
... )
>>> validation = TriplesFactory.from_path(
...     path=NATIONS_VALIDATE_PATH,
...     entity_to_id=train.entity_to_id,
...     relation_to_id=train.relation_to_id,
... )
>>> testing = TriplesFactory.from_path(
...     path=NATIONS_TEST_PATH,
...     entity_to_id=train.entity_to_id,
...     relation_to_id=train.relation_to_id,
... )
>>> pipeline_result = pipeline(
...     training=training,
...     validation=validation,
...     testing=testing,
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,
...         checkpoint_name='my_checkpoint.pt',
...     ),
... )

当你确定上面显示的数据集是相同的，你可以简单地重新运行该代码，PyKEEN 将自动从上次停止的地方恢复训练。然而，如果你只更改了数据集或对其进行了采样，你需要确保在从检查点恢复训练时映射是正确的。这可以通过以下方式从检查点加载映射来完成：

>>> import torch
>>> from pykeen.constants import PYKEEN_CHECKPOINTS
>>> checkpoint = torch.load(PYKEEN_CHECKPOINTS.joinpath('my_checkpoint.pt'))

你现在已经加载了包含映射的检查点，现在可以使用这些映射以以下方式创建与检查点中保存的模型匹配的映射：

>>> from pykeen.triples import TriplesFactory
>>> from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, NATIONS_VALIDATE_PATH
>>> training = TriplesFactory.from_path(
...     path=NATIONS_TRAIN_PATH,
...     entity_to_id=checkpoint['entity_to_id_dict'],
...     relation_to_id=checkpoint['relation_to_id_dict'],
... )
>>> validation = TriplesFactory.from_path(
...     path=NATIONS_VALIDATE_PATH,
...     entity_to_id=checkpoint['entity_to_id_dict'],
...     relation_to_id=checkpoint['relation_to_id_dict'],
... )
>>> testing = TriplesFactory.from_path(
...     path=NATIONS_TEST_PATH,
...     entity_to_id=checkpoint['entity_to_id_dict'],
...     relation_to_id=checkpoint['relation_to_id_dict'],
... )

现在你可以简单地使用与上面相同的代码来恢复管道：

>>> pipeline_result = pipeline(
...     training=training,
...     validation=validation,
...     testing=testing,
...     model='TransE',
...     optimizer='Adam',
...     training_kwargs=dict(
...         num_epochs=2000,
...         checkpoint_name='my_checkpoint.pt',
...     ),
... )

如果你觉得这工作量太大，我们仍然为你提供了保障，因为PyKEEN会在后台检查提供的三元组工厂映射是否与检查点中提供的匹配，如果不匹配，它会警告你。

手动加载模型

除了像上面展示的那样使用检查点恢复训练外，您还可以手动从检查点加载模型以进行调查或执行预测任务。这可以通过以下方式完成：

>>> import torch
>>> from pykeen.constants import PYKEEN_CHECKPOINTS
>>> from pykeen.pipeline import pipeline
>>> from pykeen.triples import TriplesFactory
>>> checkpoint = torch.load(PYKEEN_CHECKPOINTS.joinpath('my_checkpoint.pt'))

你现在已经加载了包含模型以及entity_to_id和relation_to_id映射的检查点。要将这些加载到PyKEEN中，你只需要执行以下操作：

>>> from pykeen.datasets.nations import NATIONS_TRAIN_PATH
>>> train = TriplesFactory.from_path(
...     path=NATIONS_TRAIN_PATH,
...     entity_to_id=checkpoint['entity_to_id_dict'],
...     relation_to_id=checkpoint['relation_to_id_dict'],
... )

… 现在加载模型并将训练三元组工厂传递给模型

>>> from pykeen.models import TransE
>>> my_model = TransE(triples_factory=train)
>>> my_model.load_state_dict(checkpoint['model_state_dict'])

现在你已经加载了模型，并确保三元组工厂中的映射与模型权重对齐。尽情享受吧！

待办事项

关于从hpo_pipeline恢复的教程。

注意事项和可能的错误

在使用检查点并尝试多种配置时，这会导致生成多个不同的检查点，从而存在覆盖检查点的固有风险。这种情况通常会在你更改KGEM的配置但未更改checkpoint_name参数时发生。为了防止这种情况发生，PyKEEN会对检查点的配置和当前配置进行哈希和比较。如果这些配置不匹配，PyKEEN将不接受该检查点并会引发错误。

如果你想用新的配置覆盖之前的检查点文件，你必须明确地删除它。这种行为的原因有三点：

这允许通过简单地重新运行完全相同的代码，以一种非常容易和用户友好的方式恢复中断的训练循环。
通过明确要求命名检查点文件，用户可以控制文件的命名，从而更容易保持概览。
为每次运行隐式创建新的检查点文件将导致大多数用户无意中在文件系统中生成大量未使用的检查点，这些检查点在运行许多实验时很容易累积到数百GB。

超越管道和技术细节的检查点

目前，PyKEEN 仅支持训练循环的检查点，这些检查点在类 pykeen.training.TrainingLoop 中实现。当使用如上定义的 pykeen.pipeline.pipeline() 函数时，管道实际上使用了训练循环功能。因此，这些检查点保存的是训练循环的状态，而不是管道本身的状态。因此，检查点不会包含管道中的评估结果。然而，PyKEEN 确保使用训练循环检查点的管道的最终结果与没有检查点的连续运行完全相同，包括评估结果！

为了展示如何在不使用管道的情况下使用检查点功能，我们首先定义一个KGEM：

>>> from pykeen.models import TransE
>>> from pykeen.training import SLCWATrainingLoop
>>> from pykeen.triples import TriplesFactory
>>> from torch.optim import Adam
>>> triples_factory = Nations().training
>>> model = TransE(
...     triples_factory=triples_factory,
...     random_seed=123,
... )
>>> optimizer = Adam(params=model.get_grad_params())
>>> training_loop = SLCWATrainingLoop(model=model, optimizer=optimizer)

此时，我们已经在训练循环中设置好了模型、数据集和优化器，并准备使用training_loop的方法pykeen.training.TrainingLoop.train()来训练模型。要启用检查点，您只需将函数参数checkpoint_name设置为您希望的名称。此外，您可以通过设置参数checkpoint_frequency来设置检查点的频率，即每隔多少分钟保存一次检查点，该参数接受一个整数值。默认频率为30分钟，将其设置为0将使训练循环在每个epoch后保存一个检查点。可选地，您可以通过设置checkpoint_directory参数来指定保存检查点的路径，该参数可以是一个字符串或一个pathlib.Path对象，包含您希望的根路径。如果您没有设置checkpoint_directory参数，您的检查点将保存在pykeen.constants中定义的PYKEEN_HOME目录中，该目录是您主目录下的一个子目录，例如~/.data/pykeen/checkpoints。

这是一个示例：

>>> losses = training_loop.train(
...     num_epochs=1000,
...     checkpoint_name='my_checkpoint.pt',
...     checkpoint_frequency=5,
... )

通过这段代码，我们已经使用上述定义的KGEM开始了训练循环。训练循环将在my_checkpoint.pt文件中保存一个检查点，该文件将保存在~/.data/pykeen/checkpoints/目录中，因为我们没有设置checkpoint_directory参数。检查点文件将在训练循环开始后5分钟或上次保存检查点并且当前轮次结束时保存，即当一个轮次需要10分钟时，检查点将在10分钟后保存。此外，当早停器停止训练循环或最后一个轮次完成时，检查点总是会被保存。

假设你是有预见性的，保存了检查点，而你的训练循环在200个周期后崩溃了。现在你想从最后一个检查点恢复。你所要做的就是重新运行完全相同的代码如上所述， PyKEEN 将会从给定的检查点顺利开始。由于 PyKEEN 存储了所有的随机状态以及模型、优化器和早停器的状态，结果将与不间断运行训练循环完全相同。当然，PyKEEN 也会继续保存新的检查点，即使是从之前的检查点恢复。

除了可以恢复中断的训练循环外，您还可以恢复成功完成的训练循环。例如，上述训练循环在1000个周期后成功完成，但您希望从该状态开始训练相同的模型，直到2000个周期。您只需将上述代码中的参数num_epochs更改为：

>>> losses = training_loop.train(
...     num_epochs=2000,
...     checkpoint_name='my_checkpoint.pt',
...     checkpoint_frequency=5,
... )

现在训练循环将从1000个周期的状态恢复，并继续训练直到2000个周期。

如失败时的检查点所示，你也可以仅在训练循环失败时保存检查点。为此，你只需设置参数checkpoint_on_failure=True，如下所示：

>>> losses = training_loop.train(
...     num_epochs=2000,
...     checkpoint_directory='/my/secret/dir',
...     checkpoint_on_failure=True,
... )

这段代码将在训练循环失败时保存一个检查点。注意我们如何通过将checkpoint_directory参数设置为/my/secret/dir来选择一个新的检查点目录。