# Pytorch Neural Network that tries to approximate $z_i = x_i^2 + y_i^2$ not converging to solution

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.


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
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])