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

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

torch и torch.nn импортированы как nn. Предобученный объект модели доступен в рабочей среде под именем model, а её архитектура — под именем ManufacturingCNN.

Это упражнение является частью курса

Глубокое обучение для работы с изображениями на PyTorch

Посмотреть курс

Инструкции к упражнению

Сохраните предобученную модель под именем ModelCNN.pth, не забыв сохранить веса, а не только архитектуру.
Создайте экземпляр модели с именем loaded_model на основе класса ManufacturingCNN().
Загрузите веса из ModelCNN.pth в loaded_model, передав их в метод .load_state_dict().

Интерактивное практическое упражнение

Попробуйте выполнить это упражнение, дополнив этот пример кода.

# Save the model
torch.____(model.____, ____)

# Create a new model
loaded_model = ____

# Load the saved model
loaded_model.____(torch.____('ModelCNN.pth'))
print(loaded_model)

Редактировать и запускать код