[New into ML] Why am I getting "nan" while training, how to fix it?

1

Short context: I was learning PyTorch and ML basics, here I was just writing some code and was trying to understand how the stuffs are working

Here is the sample data I’ve created

import torch

x = torch.tensor([[1, 10], [2, 20], [3, 30], [4, 40], [5, 50], [6, 60], [7, 70], [8, 80], [9, 90], [10, 100]], dtype=torch.float)
y = (5 * x[:, 0] + 6 * x[:, 1] + 1000).unsqueeze(dim=1)

x.shape, y.shape

(torch.Size([10, 2]), torch.Size([10, 1]))

and here is my training area

class LinearRegressionVersion3(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.weights = torch.nn.Parameter(torch.tensor([[0], [0]], requires_grad=True, dtype=torch.float))
    self.bias = torch.nn.Parameter(torch.tensor(0, requires_grad=True, dtype=torch.float))

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    # Corrected matrix multiplication order
    return x @ self.weights + self.bias

modelv3 = LinearRegressionVersion3()
modelv3.to(device="cuda")

MSEloss = torch.nn.MSELoss()
optimizer = torch.optim.SGD(params=modelv3.parameters(), lr=0.01)

for _ in range(50_000):
  modelv3.train()
  y_pred = modelv3(x)
  loss = MSEloss(y_pred, y)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  modelv3.eval()

print(modelv3.state_dict())

OrderedDict({'weights': tensor([[nan],
        [nan]], device='cuda:0'), 'bias': tensor(nan, device='cuda:0')})

The problem: I am getting the either nan or the weights and biases which are far away from the read one!

Stuff, I have tried: I have tried to change the lr with 0.1, 0.5, 0.01, 0.05, 0.005 and 0.001, except for lr as 0.001, everytime I am getting is nan, in training loop I have tried epocs with 10_000, 50_000, 100_000 and 500_000, but still getting the same issues!

Tools I have tried: I have tried some AI tools to getting help, but it’s just changing either lror epochs , I am totally confused, what’s the issue, is it with the formula, the sample data I made or something else!?

2

Could you try using the torch.nn.Linear module directly and see if your model converges? If it does, your initialization might be the issue. Proper initialization of model weights can have a huge impact on the training.

Just replace the modelv3 = LinearRegressionVersion3() with modelv3 = torch.nn.Linear(2, 1)

2 Likes
3

I agree with @Tanmay_Pani: initializing weights and biases to zero is quite unusual and could be the reason why your model doesn’t train properly. Using torch.nn.Linear does the same thing as the layer you’ve implemented (up to transposition of the weights matrix), but with weights and biases initialized as explained here.

I’d like to add that you should also try an even lower learning rate. Your sample data seems to contain quite large values (in the thousands), so it wouldn’t be surprising that the optimal learning rate is much lower than the lowest you’ve tried (0.001). Also, you’re saying that 0.001 is the only one not giving you nan, so it’s another reason to try an even lower learning rate.

1 Like