I’m trying to replicate a very simple linear regression NN but it appears to just default to predicting the middle value of whatever data it trains on. For example if the target data has a range 0-1 then the NN simply outputs 0.5 constantly (I guess as a loss function compromise).
It’s been a while since I’ve done a beginners network like this so maybe there’s a mistake somewhere in my model that a second pair of eyes can see?
I’ve copied a reproducible example below which runs on a GPU.
In the example I do the following:

Create fake data of the form y = 2*x + 1
Create a simple 1 in 1 out linear regression model
Train the model
Test the model and output the predicted values for the training data

import torch
import numpy as np
from torchvision import transforms
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
import random

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

torch.cuda.benchmark=True
torch.cuda.fastest =True
device = torch.device("cuda")

plt.ioff()
plt.close('all')

path = '/home/corona/c1307135/'

""" CREATING THE DATASET """

x = np.random.uniform(0,1,64)
print('made x data')
random.shuffle(x)
print('shuffled x data')
err = np.random.uniform(0,0.005,len(x))
m = 2
c = 1
y = (m*x) + c 
print('made y data')

plt.figure()
plt.plot(x,y,'k.')
plt.savefig(path + 'test.png')
plt.close('all')

data = np.vstack([x,y]).T

transform = transforms.Compose([transforms.ToTensor()])

x_tensor = torch.tensor(x,dtype=torch.float).unsqueeze(-1)
y_tensor = torch.tensor(y,dtype=torch.float)

print('made tensors')

""" INITIALISING THE WEIGHTS """

def weight_init(m):
    if isinstance(m, torch.nn.Linear):
        torch.nn.init.xavier_uniform_(m.weight.data)
        m.bias.data.zero_()
        
""" CREATING THE MODEL """

class NN(torch.nn.Module):
    
    def __init__(self,p):
        super().__init__()
             
        self.linear = torch.nn.Sequential(
            torch.nn.Linear(p,1) )
     
    def forward(self,x):
        output = self.linear(x)

        return output

""" INSTANTIATING THE MODEL """
        
model = NN(1).train()
model = model.to(device).to(torch.float)
model.apply(weight_init)
print('initialised weights')
num_iterations = 500
initial_lr = 1e-3
optimiser = torch.optim.Adam(model.parameters(), lr = initial_lr)
loss_fn = torch.nn.MSELoss()

""" TRAINING THE MODEL """

for i in range(num_iterations):
    target = y_tensor.to(device).to(torch.float)
    x_data = x_tensor.to(device).to(torch.float)
    for idx in x_data:
        pred = model(idx).squeeze(-1)
        optimiser.zero_grad()
        loss = loss_fn(pred,target)
        loss.backward()
        optimiser.step()
    if (i+1) % 1 == 0:
        print("[iteration %03d] loss: %.4f" % (i+1, loss.item()))

print("Learned parameters:")
for name, param in model.named_parameters():
    print(name,param.data.cpu().numpy())

""" TESTING THE MODEL """
    
model.eval()

predicted = model(x_tensor.to(device).to(torch.float)).detach().cpu().numpy()
print(predicted)

Perhaps this is as a result of the model ‘seeing’ the training data in the same order each epoch but I would still  expect it to overfit before 100 epochs.
Many thanks for any advice in advance!