Сохранение и загрузка моделей PyTorch

В проектах глубокого обучения сохранение и загрузка моделей — это важнейший этап. Основные причины включают:

  1. Восстановление после прерывания обучения: Можно возобновить обучение с контрольной точки, если процесс обучения неожиданно прерван
  2. Развёртывание модели: Разверните обученную модель в производственной среде
  3. Совместное использование модели: Удобно делиться результатами модели между членами команды
  4. Трансферное обучение: Сохраняйте предобученные модели для других задач
  5. Оценка эффективности: Сохраняйте модели с разных этапов обучения для сравнения

Основные методы сохранения и загрузки

Сохранение всей модели

Это самый простой метод: сохранение архитектуры и параметров модели:

Пример

import torch
import torchvision.models as models

# Создание и обучение модели
model = models.resnet18(pretrained=True)
# ... код обучения ...

# Сохранить всю модель
torch.save(model, 'model.pth')

# Загрузить всю модель
loaded_model = torch.load('model.pth')

Преимущества：

  • Простой и интуитивно понятный код
  • Сохраняет полную структуру модели

Недостатки：

  • Больший размер файла
  • Зависит от определения класса модели

Сохранение только параметров модели (Рекомендуется)

Более рекомендуемый способ — сохранять только state_dict модели:

Пример

# Сохранить параметры модели
torch.save(model.state_dict(), 'model_weights.pth')

# Загрузить параметры модели
model = models.resnet18()  # Сначала нужно создать модель с той же архитектурой
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # Установить режим оценки

Преимущества：

  • Меньший размер файла
  • Более гибкий, можно загружать в разные архитектуры
  • Лучшая совместимость

Сохранение и загрузка состояния обучения

В реальных проектах обычно также необходимо сохранять состояние оптимизатора, эпоху и другую информацию:

Пример

# Сохранить контрольную точку
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    # Можно добавить другую информацию, которую нужно сохранить
}

torch.save(checkpoint, 'checkpoint.pth')

# Загрузить чекпоинт
checkpoint = torch.load('checkpoint.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()  # Или model.train() в зависимости от ваших потребностей

Загрузка моделей на разных устройствах

Обработка совместимости CPU/GPU

Пример

# Указывайте map_location при сохранении
torch.save(model.state_dict(), 'model_weights.pth')

# Загрузка на CPU (когда модель обучалась на GPU)
device = torch.device('cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))

# Загрузка на GPU
device = torch.device('cuda')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
model.to(device)

Загрузка моделей, обученных на нескольких GPU

Пример

# Сохранение модели с несколькими GPU
torch.save(model.module.state_dict(), 'multigpu_model.pth')

# Загрузка на один GPU
model = ModelClass()
model.load_state_dict(torch.load('multigpu_model.pth'))

Преобразование моделей и совместимость

Совместимость версий PyTorch

Пример

# Укажите _use_new_zipfile_serialization=True при сохранении для лучшей совместимости
torch.save(model.state_dict(), 'model.pth', _use_new_zipfile_serialization=True)

Преобразование в TorchScript

Пример

# Преобразуйте модель в формат TorchScript
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'model_scripted.pt')

# Загрузите модель TorchScript
loaded_script = torch.jit.load('model_scripted.pt')

Лучшие практики и типичные проблемы

Лучшие практики

  1. Соглашения об именовании: используйте осмысленные имена файлов, напримерresnet18_epoch50.pth
  2. Периодическое сохранение: Сохраняйте чекпоинт каждые несколько эпох
  3. Проверка загрузки: Проверьте функциональность загрузки сразу после сохранения
  4. Документация: Записывайте архитектуру модели и параметры обучения
  5. Контроль версий: Включайте файлы модели в систему контроля версий

Решения распространённых проблем

Проблема 1：Отсутствуют ключ(и) в state_dict
Решение: Убедитесь, что архитектура модели точно совпадает, или используйтеstrict=Falseпараметр:

model.load_state_dict(torch.load('model.pth'), strict=False)

Проблема 2：Недостаточно памяти CUDA
Решение: Сначала поместите его на CPU при загрузке:

Пример

model.load_state_dict(torch.load('model.pth', map_location='cpu'))

Проблема 3: Не удаётся загрузить модели старых версий
Решение: Попробуйте загрузить в разных версиях PyTorch или преобразовать формат модели

Примеры практического применения

Процесс сохранения и загрузки модели классификации изображений

Полный пример кода

Пример

import torch
import torch.nn as nn
import torch.optim as optim

# Определим простую модель
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)
   
    def forward(self, x):
        return self.fc(x)

# Инициализация
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# Имитация процесса обучения
for epoch in range(5):
    # Имитация шагов обучения
    inputs = torch.randn(32, 10)
    labels = torch.randint(0, 2, (32,))
   
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
   
    # Сохраняем чекпоинт каждые 2 эпохи
    if epoch % 2 == 0:
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item(),
        }
        torch.save(checkpoint, f'checkpoint_epoch{epoch}.pth')
        print(f'Чекпоинт сохранён на эпохе {epoch}')

# Финальное сохранение
torch.save(model.state_dict(), 'final_model.pth')

# Пример загрузки
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('final_model.pth'))
loaded_model.eval()

# Тестирование загруженной модели
test_input = torch.randn(1, 10)
with torch.no_grad():
    output = loaded_model(test_input)
print(f'Тестовый вывод: {output}')
Другие расширения