Как сохранить и загрузить модель реформатора Google NLP

Question

Я работаю с последней моделью НЛП от Google

Я прочитал несколько постов, но в основном я курсирую над 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

Но это просто словарь с весами, а не модель для непосредственного использования:

Answer 1

если вы удалите строку «! Rm -f ~ / train_dir / model.pkl # Удалить старую модель» и измените 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)