Model not training and loss is almost constant

I am trying to train a simple model. However, it seems that the model does not learn anything, and I am unsure why I have this problem.
Here is my code:

class Regressor(nn.Module):
    def __init__(self, inp_dim=1, out_dim=1, hidden_units=20):
        super(Regressor, self).__init__()
        self.criterion = torch.nn.MSELoss()
        self.model = torch.nn.Sequential(
            torch.nn.Linear(inp_dim, hidden_units),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_units, hidden_units),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_units, out_dim),
        )
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.1)

    def update(self, x, y_target):
        y_pred = self.model(torch.Tensor(x))
        loss = self.criterion(y_pred, torch.Tensor(y_target))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def predict(self, x):
        with torch.no_grad():
            return self.model(torch.Tensor(x))

And the data and training loop is as follows:

x = torch.linspace(1, 20, 64)
y = 2 * x + x ** 2 + np.random.rand()
model = Regressor()

epochs = 1000
for epoch in range(epochs):
    model.update(x.unsqueeze(-1), y)

I would be grateful for any help.

I don’t think you can do “self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.1)”, inside the class.
By that time, init method hasn’t ended. Regression Model isn’t yet intialized completely. So you can’t pass it’s parameters to optimizer.

Since the optimizer was initialized after the model, it would still work. However, I forgot to squeeze the output of the prediction:

y_pred = self.model(torch.Tensor(x))

Instead, it should be

y_pred = self.model(torch.Tensor(x)).squeeze()
1 Like

yeah.
UserWarning did warn about different shapes. :slight_smile: Didn’t read it.