The basics of pytorch have been introduced in the previous articles, and now we will complete the complete training routine. Since it is only for practice, the smaller CIFAR10 data set is still selected.
First prepare the dataset:
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, transform=torchvision.transforms.ToTensor(),
download=True)
You can check how many pictures are in the training set and validation set:
train_data_size = len (train_data)
test_data_size = len (test_data)
print ( ' The length of the training data set is: {}' .format(train_data_size))
print ( ' The length of the test data set is: {}' .format(test_data_size))
The output is as follows:
Load the dataset with dataloader:
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
To make the code easier to read and modify, create a new file for writing the neural network model:
#build neural network
import torch
from torch import nn
class Test(nn.Module):
def __init__ ( self ):
super (Test , self ). __init__ ()
self .model = nn.Sequential(
nn.Conv2d( in_channels = 3 , out_channels = 32 , kernel_size = 5 , stride = 1 , padding = 2 ) ,
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2, stride=1),
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2, stride=1),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(in_features=1024, out_features=64),
nn.Linear(in_features=64, out_features=10)
)
def forward(self, x):
x = self.model(x)
return x
if __name__ == '__main__':
test1 = Test()
input = torch.ones((64, 3, 32, 32))
output = test1(input)
print(output.shape)
The method of building a network model has been introduced earlier, so I won’t repeat it here.
Create the network model in the main file:
test1 = Test()
Create a loss function:
loss_fn = nn.CrossEntropyLoss()
Create an optimizer with a learning rate of 0.01:
learning_rate = 0.01
optmizer = torch.optim.SGD(test1.parameters(), lr=learning_rate)
Set some parameters of the training network and record the training times:
total_train_step = 0
Record the number of tests:
total_test_step = 0
Set training rounds:
epoch = 10
Add tensorboard:
writer = SummaryWriter('logs_train')
Start training:
for i in range (epoch):
print ( "------- {} round of training starts-------" .format(i + 1 ))
#Training steps start
test1.train( )
for data in train_dataloader:
imgs , targets = data
outputs = test1(imgs)
#calculate loss value
loss = loss_fn(outputs , targets)
#optimize optmizer.zero_grad() loss.backward
()
optmizer.step() #training times increase 1. Print total_train_step = total_train_step + 1 if
total_train_step % 100 == 0 :
print ( " Training times: {}, loss: {}" .format(total_train_step , loss.item()))
writer.add_scalar( "train_loss" , loss.item() , total_train_step)
# The test step starts
test1.eval()
total_test_loss = 0
#matching correct times
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs , targets = data
outputs = test1(imgs)
loss = loss_fn(outputs , targets)
#loss of the overall test
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax( 1 ) == targets).sum()
total_accuracy = total_accuracy + accuracy
print ( " Overall test set loss on: {}" .format(total_test_loss))
print ( " Positive rate on the overall test set: {}" .format(total_accuracy/test_data_size))
writer.add_scalar( "test_loss" , total_test_loss , total_test_step)
writer.add_scalar ( "test_accuracy" , total_accuracy/test_data_size, total_test_step)
total_test_step = total_test_step + 1
torch.save(test1, "lxw_{}.ptn".format(i))
writer.close()
The output is as follows:
The figure above shows the loss value of the first round of training, and it can be seen that the loss value is gradually decreasing.
The accuracy rate of the last round of training was significantly improved.
View the training process in tensorboard:
The training sample loss value changes as follows:
The test sample loss values are as follows:
The test accuracy is as follows:
