2
$\begingroup$

Background

I am teaching myself Pytorch, as a Mechanical engineering technology (MET) faculty. My end goal is to replace many data-driven heat transfer and Fluid dynamics models with Neural network approximations. This is a wholly academic exercise to expose my MET students to Neural networks via a familiar environment. I have some experience creating Neural Networks using the Wolfram Language.

The problem statement

Approximate $z_i = x_i^2 + y_i^2$ with a multi-layer feedforward percepteron where $x_i, y_i$ are randomly generated floats.

The issue I am faced with

The NN I created using Pytorch does not converge and I cannot tell if this is because of:

  • Improper layer definitions. I have experimented with three layers (linear - ReLU or Tanh - linear) and the current one. I have experimented with different numbers of outputs from the first linear layer.
  • Not sufficient epochs.
  • Improper learning rate.
  • The code itself is foul.

I would greatly appreciate help or advice on this matter.

I have included my code. I would be happy to provide any other information.

Code

Setting up the data, NN layers, and the optimizer

import torch
import torch.nn as nn
import numpy as np
from sklearn.model_selection import train_test_split


x = np.random.random(1000);
y = np.random.random(1000);
z = x**2 + y**2
input_data = torch.Tensor(np.transpose([x ,y]))
output_data = torch.Tensor(z)

input_training, input_validation, output_training, output_validation = train_test_split(input_data, output_data, random_state=42, test_size=0.15, shuffle=True)


class NonLinearRegression(torch.nn.Module):
  def __init__(self):
    super(NonLinearRegression, self).__init__()
    self.linear_1 = nn.Linear(in_features=2, out_features=10)
    self.act_1 = nn.ReLU()
    self.linear_2 = nn.Linear(in_features=10,out_features=5)
    self.act_2 = nn.ReLU()
    self.linear_3 = nn.Linear(in_features=5,out_features=1)

  def forward(self, y):
    y = self.linear_1(y)
    y = self.act_1(y)
    y = self.linear_2(y)
    y = self.act_2(y)
    y = self.linear_3(y)
    y_pred = y
    return y_pred  


model_nonlinear = NonLinearRegression()   
optimizer = torch.optim.SGD(model_nonlinear.parameters(), lr=1e-6)
criterion = nn.MSELoss(reduction='sum')

The NN training loop

epoch_max = 20000
for epoch in range(epoch_max):        
    total_loss = 0;
    model_nonlinear.train()    
    y_pred = model_nonlinear(input_training)
    loss = criterion(y_pred, output_training)
    loss.backward()
    total_loss += float(loss)
    if (total_loss < 0.001):
        print("Num steps: " + str(epoch))
        break
    optimizer.step()

Validation

input_validation, model_nonlinear(input_validation) #the math does not check out.
$\endgroup$

1 Answer 1

4
$\begingroup$

I think the main thing that is wrong is that your training loop is currently not resetting the gradients between epochs (using optimizer.zero_grad()). This causes gradients to accumulate, which stops your network from learning properly. Making this single change already massively improves the learning of your network, achieving a loss of around 6.5 after 20000 epochs. Some additional changes I've made that improve/speed up the learning even more are the following:

  • Use the Adam optimizer instead of the SGD optimizer, the default learning rate of 0.001 seems to work fine.
  • Use mini batches instead of the full training dataset.
  • Increasing the number of parameters in your model.

These changes result in the following code:

# model definition
class NonLinearRegression(torch.nn.Module):
    def __init__(self):
        super(NonLinearRegression, self).__init__()
        self.linear_1 = nn.Linear(in_features=2, out_features=25)
        self.act_1 = nn.ReLU()
        self.linear_2 = nn.Linear(in_features=25,out_features=10)
        self.act_2 = nn.ReLU()
        self.linear_3 = nn.Linear(in_features=10,out_features=1)

    def forward(self, y):
        y = self.linear_1(y)
        y = self.act_1(y)
        y = self.linear_2(y)
        y = self.act_2(y)
        y = self.linear_3(y)
        y_pred = y
        return y_pred  


model_nonlinear = NonLinearRegression()
# changed optimizer
optimizer = torch.optim.Adam(model_nonlinear.parameters())
criterion = nn.MSELoss(reduction='sum')

# training loop
epoch_max = 2500
for epoch in range(epoch_max):
    # mini-batch
    ix = torch.randint(0, input_training.shape[0], size=(64,))
    total_loss = 0
    y_pred = model_nonlinear(input_training[ix])
    loss = criterion(y_pred.squeeze(), output_training[ix])
    # set gradients to zero
    optimizer.zero_grad()
    loss.backward()
    total_loss += float(loss)
    if (total_loss < 0.001):
        print("Num steps: " + str(epoch))
        break
    optimizer.step()

With the following loss curve and predictions: enter image description here enter image description here

$\endgroup$
3
  • 1
    $\begingroup$ > is currently not resetting the gradients between epochs Wow. That's a big one! Thank you for catching it. I will run your modified code and see how it works. I appreciate it. Is there any reading that you can suggest on the choice of optimizer and layers (the univ. approx. thm. notwithstanding)... or is this purely experiential knowledge? $\endgroup$
    – dearN
    Commented Oct 19, 2022 at 13:47
  • 2
    $\begingroup$ Most of it is experimental knowledge combined with trial and error, when it comes to the choice of optimizer I've found Adam is used most often as it generally just works and is able to achieve quite good results. For that extra bit of performance SGD is sometimes also used, but it seems choosing the correct learning rate is a bit more important. $\endgroup$
    – Oxbowerce
    Commented Oct 19, 2022 at 13:50
  • 1
    $\begingroup$ Your code works. I appreciate the insight into using mini-batches. All these parameters and choices are so close and the "mechanics" of modifying them is so close to my original field which is numerical simulations of physical systems. $\endgroup$
    – dearN
    Commented Oct 19, 2022 at 17:38

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.