Python 如何保存和加载google NLP改革者模型

我正在使用最近的NLP模型

我读过几篇文章，但大部分我都在浏览colab，它有所有的建模步骤和测试功能。我现在面临的问题是，即使使用google TPU，该模型也需要很长时间来训练，因此我需要保存训练过的模型，我猜想它的工作原理与GPT-2模型类似，因为该模型可以在多个会话中训练，因为它允许随时停止训练：

This will take at least 30 minutes to run to completion, but can safely
# be interrupted by selecting "Runtime > Interrupt Execution" 

但是我还没有找到一个关于如何在训练后保存和加载模型的示例。在GPT-2的情况下，会自动为每个新模型创建一个新目录，使用它只需指向该新目录，但对于这个目录，我不知道如何加载以前训练过的模型

编辑：

在笔记本中，我看到了以下代码：

# Set up a Trainer.
output_dir = os.path.expanduser('~/train_dir/')
!rm -f ~/train_dir/model.pkl  # Remove old model
trainer = trax.supervised.Trainer(
    model=trax.models.ReformerLM,
    loss_fn=trax.layers.CrossEntropyLoss,
    optimizer=trax.optimizers.Adam,
    lr_schedule=trax.lr.MultifactorSchedule,
    inputs=trax.supervised.inputs.Inputs(my_inputs),
    output_dir=output_dir,
    has_weights=True)

这将删除以前的模型，我查看了该目录，发现如下：

我使用pickle加载这个model.pkl文件，并将其复制到我的Gdrive文件夹：

with open('model.pkl', 'rb') as handle:
    reformer_model = pickle.load(handle)

reformer_model

但这只是一本关于Weights的字典，不是一个可以直接使用的模型：

如果删除行“！rm-f~/train_dir/model.pkl#remove old model”，并将output_dir更改为指向保存的模型所在的文件夹，则将加载该模型并从停止的位置继续培训。如果该目录中没有模型，它将创建一个新模型

from google.colab import drive
drive.mount('/content/gdrive',  force_remount=True)
%cd /content/gdrive/My\ Drive/


# Train tiny model with Trainer.
output_dir = "CornellMovieDialog/Model/"
trainer = trax.supervised.Trainer(
    model=tiny_transformer_lm,
    loss_fn=trax.layers.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adafactor,  # Change optimizer params here.
    lr_schedule=trax.lr.MultifactorSchedule,  # Change lr schedule here.
    inputs=copy_inputs,
    output_dir=output_dir)