In this tutorial, you will learn
- How to train a model in PyTorch
- How to evaluate model performance
- How to save and load models
Assumption
You know some basic programming (array, loop etc.)
What are you going to do?
#1: We will create a dataset with 1000 rows like the one below where:
y = x1 * x2
#2: We will build a neural network to predict y from x1 and x2. Note that the choice of the equation ( y = x1 * x2) is arbitrary. You can try different equations and see that the model can predict almost anything!
Step#1
Go to Google Colab: https://colab.research.google.com/. Log in if you are not logged in already. You should see a screen like the one below:
Step#2
Click New Notebook:
Step#3
You will get a place to write code. Copy the following code and paste it there:
import torch
from torch import nn
from sklearn.metrics import r2_score
class MyMachine(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Sequential(
nn.Linear(2,5),
nn.ReLU(),
nn.Linear(5,1)
)
def forward(self, x):
x = self.fc(x)
return x
def get_dataset():
X = torch.rand((1000,2))
x1 = X[:,0]
x2 = X[:,1]
y = x1 * x2
return X, y
def train():
model = MyMachine()
model.train()
X, y = get_dataset()
NUM_EPOCHS = 1000
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=1e-5)
criterion = torch.nn.MSELoss(reduction='mean')
for epoch in range(NUM_EPOCHS):
optimizer.zero_grad()
y_pred = model(X)
y_pred = y_pred.reshape(1000)
loss = criterion(y_pred, y)
loss.backward()
optimizer.step()
print(f'Epoch:{epoch}, Loss:{loss.item()}')
torch.save(model.state_dict(), 'model.h5')
def test():
model = MyMachine()
model.load_state_dict(torch.load("model.h5"))
model.eval()
X, y = get_dataset()
with torch.no_grad():
y_pred = model(X)
print(r2_score(y, y_pred))
train()
test()Step#4
Hit the run button:
Step#5
It will take a few seconds to run. Then you will see an output like the below:
Explanation
As you see in the last two lines, we are first training the model and then testing:
Both train and test functions use the MyMachine class and get_dataset function. So, let us see the structure of those first.
MyMachine
As you can see from the definition of MyMachine:
it is a simple neural network of the following structure:
Notice that we only define the forward pass mechanism. PyTorch will handle backward passes (backpropagation) automatically (with the help of computational graphs; if you are curious how).
get_dataset()
On the first line, we define a 2-dimensional array X (1000 rows and 2 columns) with random numbers. The columns represent x1 and x2, respectively. Then we derive an array y (1000 rows and 1 column) by multiplying x1 by x2.
train()
On the first line, we create an instance of MyMachine (model). On the second line, we are informing PyTorch that we are using the model for training at this moment.
An optimizer is used to update the weights and biases of the neural network. Here we are using Adam optimizer.
After that, we define a Mean-squared error (MSE) loss function.
Epoch simply means the number of times all data is passed for training.
At every iteration in the for loop, with:
we reset gradient values to ensure gradients are not adjusted multiple times.
After we pass the data through the model, we get the prediction (y_pred) in the form of an array of arrays of single values, for example:
[
[0.8425],
[0.7415],
[0.5790],
...
]Before comparing the prediction (y_pred) with the ground truth (y), we have to make its structure the same as y:
Then the structure will be like this:
[
0.8425,
0.7415,
0.5790,
...
]After calculating the loss, we derive the gradients for all the weights and biases by:
Then, by running
we update all weights and biases based on the derived gradients.
Finally, we save the trained model in the file model.h5:
test()
Again, on the first line, we create an instance of MyMachine (model). Then, we load the parameters (weights and biases) of the machine with the trained parameters we just saved. On the third line, we inform PyTorch that we are using the model in test (evaluation) mode.
After that, we get a new dataset (as X, y will be formed with newly generated random numbers) and pass X through the model. Finally, we print the R-squared value of the ground truth (y) and the prediction (y_pred). Simply speaking, the R-squared value indicates the prediction performance. R-squared value 1 means perfect prediction without any error (which is not practical). A low R-squared value indicates poor prediction.
