Newbie, non linear regression not converging

First im new to pytorh and DL, I want to create a simple non linear regression model, but apparently is not converging, i tried to change some hyperparams without sucess. This is the code, i guess im making wrong something obvius.

    import torch
    
    x1 = torch.arange(1,600,1,  dtype=torch.float32).view(-1,1)
    y1 = x1*x1
    
    model = torch.nn.Sequential(
        torch.nn.Linear(1,500),
        torch.nn.ReLU(),
        torch.nn.Linear(500,1)
    )
    
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
    
    for i in range (10000):
        y_pred = model(x1)
        loss = criterion(y_pred, y1)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
        if i%100 == 0:
            print(loss)

print(model(torch.tensor([6], dtype=torch.float32)))

Your code works fine if I use Adam and increase the number of epochs.
The loss starts with a huge value since you are not normalizing the inputs (and outputs), which could help training the model.
E.g. if I use this native normalization:

x1 = torch.arange(1,600,1,  dtype=torch.float32).view(-1,1)
y1 = x1*x1

mean = x1.mean()
std = x1.std()
x1 = x1 - mean
x1 = x1 / std

mean = y1.mean()
std = y1.std()
y1 = y1 - mean
y1 = y1 / std

model = torch.nn.Sequential(
    torch.nn.Linear(1,500),
    torch.nn.ReLU(),
    torch.nn.Linear(500,1)
)

criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for i in range (10000):
    y_pred = model(x1)
    loss = criterion(y_pred, y1)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if i%100 == 0:
        print(loss)
        
        
import matplotlib.pyplot as plt

plt.plot(y1 * std + mean)
plt.plot(y_pred.detach() * std + mean)

I can overfit the model pretty quickly:
image

1 Like

@ptrblck many thanks.