🔐 Как сохранить обученную нейронную сеть Python PyTorch

Как сохранить обученную нейронную сеть в PyTorch?

Чтобы сохранить обученную нейронную сеть в PyTorch, вы можете воспользоваться функцией torch.save().

# Пример
import torch

# Создание и обучение модели
model = your_model()
train(model)

# Сохранение модели
torch.save(model.state_dict(), 'path_to_save_model.pt')

Функция torch.save() сохраняет состояние модели в формате .pt или .pth в указанном пути. В состоянии модели хранятся значения всех параметров и весов модели, необходимых для воспроизведения и использования обученной модели на новых данных.

Чтобы загрузить сохраненную модель, используйте функцию torch.load().

# Пример
import torch

# Загрузка сохраненной модели
model = your_model()
model.load_state_dict(torch.load('path_to_saved_model.pt'))

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

Как сохранить обученную нейронную сеть с использованием Python и PyTorch

В этой статье мы рассмотрим, как сохранить обученную нейронную сеть с использованием Python и библиотеки глубокого обучения PyTorch. PyTorch предоставляет средства для сохранения моделей, которые могут быть восстановлены и использованы позже для предсказания или дообучения. Давайте посмотрим на несколько шагов, чтобы достичь этой цели.

Шаг 1: Импорт необходимых библиотек

Первым шагом будет импорт необходимых библиотек - PyTorch и torch.nn. Убедитесь, что у вас установлена последняя версия PyTorch.

import torch
import torch.nn as nn

Шаг 2: Определение модели нейронной сети

Затем определим структуру и параметры нашей нейронной сети. Для этого создадим класс модели, который наследуется от базового класса nn.Module.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(in_features, hidden_size)
        self.fc2 = nn.Linear(hidden_size, out_features)
        
    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

Обратите внимание, что мы определили слои нейронной сети в __init__ методе, а затем определили ее прямой проход в методе forward.

Шаг 3: Обучение и сохранение модели

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

# Создание экземпляра модели
model = NeuralNetwork()

# Определение функции потерь и оптимизатора
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# Цикл обучения
for epoch in range(num_epochs):
    # Обработка пакетов данных
    for inputs, labels in dataloader:
        # Обнуление градиентов параметров
        optimizer.zero_grad()
        
        # Прямой проход через модель
        outputs = model(inputs)
        
        # Вычисление функции потерь
        loss = criterion(outputs, labels)
        
        # Обратное распространение ошибки и оптимизация
        loss.backward()
        optimizer.step()

# Сохранение модели
torch.save(model.state_dict(), 'trained_model.pth')

В этом коде мы создаем экземпляр модели, определяем функцию потерь и оптимизатор, а затем проходим через итерации обучения, вычисляя предсказания модели и обратное распространение ошибки для оптимизации параметров. После завершения обучения мы сохраняем состояние модели в файл 'trained_model.pth'.

Шаг 4: Восстановление обученной модели

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

# Создание экземпляра модели
model = NeuralNetwork()

# Загрузка весов модели
model.load_state_dict(torch.load('trained_model.pth'))

# Установка режима оценки (не требуется градиент)
model.eval()

# Использование модели для предсказаний
outputs = model(inputs)

В этом коде мы создаем экземпляр модели, загружаем сохраненные веса из файла 'trained_model.pth', устанавливаем режим оценки (выключаем градиенты) и используем модель для предсказания на новых данных.

Заключение

В этой статье мы рассмотрели, как сохранить обученную нейронную сеть с использованием Python и PyTorch. Мы прошли через шаги по созданию и обучению модели, а затем показали, как сохранить ее в файл. Далее мы рассмотрели, как восстановить обученную модель для использования в предсказаниях. Код приведен для справки и поможет вам начать работу с сохранением и восстановлением нейронных сетей в PyTorch.