Pytorch save and load model

Model saving and loading

1 Only save and load model parameters

torch.save(model.state_dict(), PATH)   ### Save the parameters of the model to this address , Suffix named pt

model = model(*args, **kwargs)    ### Defining models 
model.load_state_dict(torch.load(PATH))  ## Import model parameters 

2 Save and load the entire model

torch.save(model,path)

model=torch.load(path)

This method can directly save the whole model , There is no need to redefine the model when applying .

Define network structure

The simplest network structure is defined here . Full connection layer of two layers

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1=nn.Linear(1,3)   ### Linear layer 
        self.layer2=nn.Linear(3,1)
    def forward(self,x):
        x=self.layer1(x)
        x=torch.relu(x)   ###relu Activation function 
        x=self.layer2(x)
        return x

Training neural network

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt


epoches=2000


#  Learning rate is defined as 0.01
learning_rate=0.01

#  Create a model 
model=Net()

optimizer = optim.Adam(params=model.parameters(), lr=learning_rate, weight_decay=1e-5)
criterion = nn.MSELoss() # Define the loss function 
#  Use the optimizer to update the network weights ,lr For learning rate ,
for i in range(epoch): # Set training epoch Time 
	model.train() # Set the state of the model to train
	for j in Sample: # Traverse each sample 
		optimizer.zero_grad() # Gradient cleanup , Prepare for this gradient calculation 
		output = model(j)
		loss = criterion(output, target)
		loss.backward()    ### The average recorded here loss
		optimizer.step() # Update network weights 	

        if (epoch+1) % 10==0:   ## Print the current status every ten times 
        print("Epoch {} / {},loss {:.4f}".format(epoch+1,num_epoches,loss.item()))	

Pytorch Model preservation

##torch.save()  You can save dictionary type data 
save_checkpoint({'loss': i, 'state_dict': model.state_dict()},dir)

def save_checkpoint(state, dic): #state Is the weight and state of the model   dic Is the directory where the model is saved 
	if not os.path.exists(dir):
		os.makedirs(directory)
	fileName = directory + 'last.pth'
	torch.save(state,fileName)# Use torch.save Function directly saves the trained model