Salvar e carregar um modelo

Uma empresa de manufatura quer classificar seus projetos com base em imagens e definir a embalagem de envio adequada. Depois de treinar um modelo altamente preciso em PyTorch, você agora planeja salvar o modelo e seus pesos pré-treinados para uso futuro e para compartilhá-lo com sua equipe, garantindo que eles possam carregá-lo sem complicações.

torch e torch.nn como nn já foram importados. O objeto do modelo pré-treinado está disponível no seu ambiente como model, e sua arquitetura como ManufacturingCNN.

Este exercício faz parte do curso

Deep Learning para Imagens com PyTorch

Ver curso

Instruções do exercício

Salve o modelo pré-treinado como ModelCNN.pth, lembrando de salvar os pesos, não apenas a arquitetura.
Crie uma instância de modelo chamada loaded_model a partir da classe ManufacturingCNN().
Carregue os pesos de ModelCNN.pth em loaded_model passando os pesos para .load_state_dict().

Exercício interativo prático

Experimente este exercício completando este código de exemplo.

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

# Create a new model
loaded_model = ____

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

Editar e executar o código