Сохранение обученной модели в PyTorch: лучшие практики

Быстрый ответ

Для сохранения модели в PyTorch используйте функцию torch.save(). Эта функция предназначена для сохранения параметров модели:

torch.save(model.state_dict(), 'my_model.pth')

Чтобы загрузить модель, примените функцию model.load_state_dict(). Это тоже достаточно просто:

model = MyModel() # заранее определите класс модели
model.load_state_dict(torch.load('my_model.pth'))

Этот способ подразумевает сохранение исключительно обучаемых параметров, благодаря чему получается компактный и переносимый файл модели.

Сохранение и загрузка: Подробное руководство

Паузы в долгих обучающих сессиях

При длительных сессиях обучения удобно использовать контрольные точки, в которых сохраняется больше данных, чем при обычном сохранении состояния модели:

checkpoint = {
    'epoch': epoch,  # количество выполненных эпох
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,  # значение функции потерь
    # можно добавлять и другие метрики, нужные для ваших экспериментов
}
torch.save(checkpoint, 'pitstop.pth')

При загрузке контрольной точки сначала следует инициализировать модель и оптимизатор:

checkpoint = torch.load('pitstop.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

# Если загрузка происходит в середине эпохи…
epoch = checkpoint['epoch']
loss = checkpoint['loss']

При загрузке модели для использования в инференсе не забывайте переключить её в режим model.eval().

Сохранение полной модели: когда важен контекст

Если есть необходимость сохранить всю модель, особенно в случае потери исходного кода модели, можно использовать:

torch.save(model, 'the_whole_shebang.pth')

Для восстановления модели используется:

model = torch.load('the_whole_shebang.pth')

Совместимость между разными версиями PyTorch

Рекомендуется указывать используемую версию PyTorch при сохранении моделей для избежания проблем совместимости в будущем.

Особенности сохранения при распределённом обучении

Если вы используете модель в контексте DataParallel или DistributedDataParallel, сохраняйте параметры следующим образом:

torch.save(model.module.state_dict(), 'all_for_one.pth')

Загрузка параметров в модель другой архитектуры

Для загрузки параметров в модель другой архитектуры надо сначала определить класс модели, а потом загрузить состояние.

Визуализация

Сохранение и восстановление модели в PyTorch состоит из следующих этапов:

1. Обучение: модель обучается выполнять некоторые задачи.
2. Сериализация: применяем `torch.save` для упаковки модели.
3. Создание файла: присваиваем модели имя для удобства управления.
4. Хранение: сохраняем модель на любой носитель информации.
5. Развертывание: модель готова к использованию.

Примечание: Сериализация здесь выступает как упаковка чемодана перед путешествием!

Полезные материалы

Сохранять модель напрямую через model.parameters() не стоит, так как это может вызвать ошибки. Используйте state_dict(). Дайте вашим файлам расширения .pt или .pth для следования конвенциям PyTorch. Когда работаете с DataParallel, не забывайте правильно обработать модель перед загрузкой состояния. Полезно сохранять в контрольных точках не только состояние модели и оптимизатора, но и прошедшие эпохи и различные метрики. Периодически проверяйте свои модели на сериализуемость и ищите баланс между объемом хранения и удобством развертывания.