How to save and load a pytorch model

[Tutorial is under development]

Saving Model

import torch
import torch.nn as nn
from torch.optim import SGD
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import random

# generate synthetic data
X = []
Y = []
for _ in range(10):
    x1 = random.randint(1,20)
    x2 = random.randint(1,20)
    X.append([x1, x2])
    Y.append([x1 + x2])

# convert lists into tensor objects (it is a good practice to convert tensor elements in float)
# objects since it will be multiplied by float weights
X = torch.tensor(X).float()
Y = torch.tensor(Y).float()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
X = X.to(device)
Y = Y.to(device)

class ThisDataset(Dataset):
    def __init__(self,x,y):
        self.x = torch.tensor(x).float()
        self.y = torch.tensor(y).float()
    def __len__(self):
        return len(self.x)
    def __getitem__(self, ix):
        return self.x[ix], self.y[ix]
ds = ThisDataset(X, Y)

# create a neural network which can learn multiplication operation
class MulNeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_to_deep_layer = nn.Linear(2,8) # 2 input features and 8 output features
        self.deep_layer_activation = nn.ReLU() # relu activation, other activation functions are sigmod, softmax, tanh
        self.deep_layer_to_output_layer = nn.Linear(8,1)

    def forward(self, x):
        '''
        It is mandatory to use forward function name in pytorch since it is a reserved name in PyTorch
        for forward propogation
        '''
        x = self.input_to_deep_layer(x)
        x = self.deep_layer_activation(x)
        x = self.deep_layer_to_output_layer(x)
        return x

mynet = MulNeuralNet().to(device)

def custom_mean_squared_error(y_pred, y):
    loss = (y_pred-y)**2
    loss = loss.mean()
    return loss

loss_func = custom_mean_squared_error

opt = SGD(mynet.parameters(), lr = 0.001)
loss_track = []
dl = DataLoader(ds, batch_size=5, shuffle=True)
for _ in range(100):
    for data in dl:
        x, y = data
        opt.zero_grad()
        loss_value = loss_func(mynet(x),y)
        loss_value.backward()
        opt.step()
        loss_track.append(loss_value)

torch.save(mynet.to('cpu').state_dict(), 'mynetmodel.pth')

Loading Model:

import torch
import torch.nn as nn
from torch.optim import SGD
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import random

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# create a neural network which can learn multiplication operation
class MulNeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_to_deep_layer = nn.Linear(2,8) # 2 input features and 8 output features
        self.deep_layer_activation = nn.ReLU() # relu activation, other activation functions are sigmod, softmax, tanh
        self.deep_layer_to_output_layer = nn.Linear(8,1)

    def forward(self, x):
        '''
        It is mandatory to use forward function name in pytorch since it is a reserved name in PyTorch
        for forward propogation
        '''
        x = self.input_to_deep_layer(x)
        x = self.deep_layer_activation(x)
        x = self.deep_layer_to_output_layer(x)
        return x

mynet = MulNeuralNet().to(device)

state_dict = torch.load('mynetmodel.pth')
mynet.load_state_dict(state_dict)
mynet.to(device)

val_x = [[7,22]]
val_x = torch.tensor(val_x).float().to(device)
print(mynet(val_x))