Plotting training and validation loss curve in my PyTorch model

Question

I developed my model using the following code.

class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data.drop('cnt', axis=1).values
        self.targets = data['cnt'].values

    def __len__(self):
        return len(self.targets)

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.targets[idx]
        return x, y

class MyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_hidden_layers):
        super(MyModel, self).__init__()
        self.input_layer = nn.Linear(input_size, hidden_size)
        self.hidden_layers = nn.ModuleList()
        for i in range(num_hidden_layers):
            self.hidden_layers.append(nn.Linear(hidden_size, hidden_size))
        self.output_layer = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = torch.relu(self.input_layer(x))
        for layer in self.hidden_layers:
            x = torch.relu(layer(x))
        x = self.output_layer(x)
        return x

input_size = 57
hidden_size = 128
output_size = 1
num_hidden_layers = 2
learning_rate = 0.001
num_epochs = 100

model = MyModel(input_size, hidden_size, output_size, num_hidden_layers)

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

train_dataset = MyDataset(train_data)
val_dataset = MyDataset(val_data)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs.float())
        loss = criterion(outputs.squeeze(), targets.float())
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * inputs.size(0)
    train_loss /= len(train_loader.dataset)

    model.eval()
    val_loss = 0.0
    val_preds = []
    val_targets = []
    with torch.no_grad():
        for inputs, targets in val_loader:
            outputs = model(inputs.float())
            loss = criterion(outputs.squeeze(), targets.float())
            val_loss += loss.item() * inputs.size(0)
            val_preds += outputs.squeeze().tolist()
            val_targets += targets.tolist()
        val_loss /= len(val_loader.dataset)
        val_r2 = r2_score(val_targets, val_preds)
        val_mse = mean_squared_error(val_targets, val_preds, squared=True)

    print(f'Epoch {epoch+1}/{num_epochs}, '
          f'Training Loss: {train_loss:.4f}, '
          f'Validation Loss: {val_loss:.4f}, '
          f'Validation R^2: {val_r2:.4f}, '
          f'Validation MSE: {val_mse:.4f}')

And it gives me the results showing the training loss and validation loss in each epoch. result of each epoch

Then I try to plot training and validation loss curve using the following code.

import matplotlib.pyplot as plt

train_losses = []
val_losses = []

for epoch in range(num_epochs):    
    train_losses.append(train_loss)
    val_losses.append(val_loss)

plt.plot(train_losses, label='Training Loss')
plt.plot(val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

However, the training loss and validation loss curves in the plot do not align with the losses of each epoch above. Why is my plot not showing the "curves"? plot

Answer 1

Note that you print train_loss and val_loss within the fitting loop and from what you posted it seems that train_losses and val_losses for plotting is filled afterwards with a constant value (probably the last value assigned in the fitting loop).

So just do

...
    print(f'Epoch {epoch+1}/{num_epochs}, '
          f'Training Loss: {train_loss:.4f}, '
          f'Validation Loss: {val_loss:.4f}, '
          f'Validation R^2: {val_r2:.4f}, '
          f'Validation MSE: {val_mse:.4f}')
    train_losses.append(train_loss)
    val_losses.append(val_loss)

And then plot train_losses and val_losses as before.