saving_restoring_trained_models.py

# -*- coding: utf-8 -*-
"""saving/restoring-trained-models.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1-3bs2d73DEVa_v-1vAUasTuuwFdYsB74
"""

import torch
import numpy as np
from torchvision import datasets,transforms
import matplotlib.pyplot as plt

from torch import nn
from torch import optim
import torch.nn.functional as F

class Network(nn.Module):
  def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(784,256)
    #self.drop1 = nn.Dropout(p=0.3)
    self.fc2 = nn.Linear(256,128)
    self.fc3 = nn.Linear(128,64)
    self.fc4 = nn.Linear(64,10)

  def forward(self,x):
    x = self.fc1(x)
    x = F.relu(x)
    #x = self.drop1(x)
    x = self.fc2(x)
    x = F.relu(x)
    x = self.fc3(x)
    x = F.relu(x)
    x = self.fc4(x)

    return x

model = Network()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(),lr=0.003)

"""# Saving The Model"""

checkpoint = {"input_size":784,
              "output_size":10,
              "hidden_layers":[layer.out_features for layer in model.modules() if type(layer)!=Network],
              "state_dict":model.state_dict()}

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

def load_checkpoint(filepath,model):
  checkpoint = torch.load(filepath)
  for id,layer in enumerate(model.modules()):
    if type(layer)!=Network:
      assert layer.out_features == checkpoint['hidden_layers'][id-1]

  model.load_state_dict(checkpoint["state_dict"])
  return model

load_checkpoint('checkpoint.pth',model)