I need some insights on high and volatile loss. It is increasing over the course of a few epochs, and then stays in a very volatile and high range.
I have plotted loss, rms and r² score.
Is there an error in my code? Or do I need more training data?
import torch
import numpy as np
from torch.utils.data import DataLoader, TensorDataset, IterableDataset
import os
from torch import optim
from model import Network
from matplotlib import pyplot as plt
from sklearn.metrics import r2_score

path = os.getcwd()

inputs = np.load(path + '/data/in1a.npy')
outputs = np.load(path + '/data/out1a.npy')

batch_size = 400

inputs_0 = torch.Tensor(inputs[0,:])
inputs_1 = torch.Tensor(inputs[1,:])

outputs_0 = torch.Tensor(outputs[0,:])
outputs_1 = torch.Tensor(outputs[1,:])
outputs_2 = torch.Tensor(outputs[2,:])


inputs_tensor = torch.Tensor(inputs)

dataset_in = TensorDataset(inputs_0, inputs_1)
dataset_out = TensorDataset(outputs_0, outputs_1, outputs_2)

dataloader_in = DataLoader(dataset_in, batch_size=batch_size, shuffle=False)
dataloader_out = DataLoader(dataset_out, batch_size=batch_size, shuffle = False)

number_batches_in = int(len(dataset_in)/batch_size)
number_batches_out = int(len(dataset_out)/batch_size)

x = torch.empty(size=(number_batches_in, 800))
y = torch.empty(size=(number_batches_out,1200), dtype=torch.float64)

for index, (x1, x2) in enumerate(dataloader_in):
    batch = torch.cat((x1, x2), 0)
    x[index] = batch

x_mean = torch.mean(x)
x_std = torch.std(x)
x_norm = (x - x_mean) / x_std

for index, (y1, y2, y3) in enumerate(dataloader_out):
    batch = torch.cat((y1, y2, y3), 0)
    y[index] = batch

y_mean = torch.mean(y)
y_std = torch.std(y)
y_norm = (y - y_mean) / y_std


model = Network(800,1200,3,800,200)
SAVE_PATH = "trained/model.dat"
epochs = 200
learning_rate = 0.001
optimizer = optim.Adam(model.parameters(),lr=learning_rate, eps=1e-08)
hist_error = []
hist_loss = []
hist_r2 = []
beta = 0.5

for epoch in range(epochs):
    epoch_error = []
    epoch_loss = []
    epoch_r2 = []
    for x_batch, y_true in zip(x_norm, y_norm):
        optimizer.zero_grad()
        x_batch = torch.unsqueeze(x_batch, 0)
        y_true = torch.unsqueeze(y_true, 0)
        pred = model.forward(x_batch)
        loss = torch.mean(torch.sum((pred - y_true)) ** 2)
        loss.backward()
        optimizer.step()
        error = torch.mean(torch.sqrt((pred - y_true) ** 2)).detach().numpy()
        r2y = torch.squeeze(y_true, 0)
        r2pred = torch.squeeze(pred, 0)
        r2 = r2_score(r2y.detach().numpy(), r2pred.detach().numpy())
        epoch_error.append(error)
        epoch_loss.append(loss.data.detach().numpy())
        epoch_r2.append(r2)
    hist_error.append(np.mean(epoch_error))
    hist_loss.append(np.mean(epoch_loss))
    hist_r2.append(np.mean(epoch_r2))
    print("Epoch %d -- loss %f, RMS error %f, R² score %f " % (epoch+1, hist_loss[-1], hist_error[-1], hist_r2[-1]))
torch.save(model.state_dict(), SAVE_PATH)
print("Model saved to %s" % SAVE_PATH)


f, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True)
ax1.plot(hist_error)
ax1.set_ylabel("Amplitude RMSE")
ax2.plot(hist_loss)
ax2.set_ylabel("Loss")
ax3.plot(hist_r2)
ax3.set_ylabel("R²")
ax3.set_xlabel("Epoch")

plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F


class Network(nn.Module):
    def __init__(self, input_dim, output_dim, latent_dim, layer_dim1, layer_dim2):
        """
        Parameter:
        input_dim (int): number of inputs
        output_dim (int): number of outputs
        latent_dim (int): number of latent neurons
        Layer_dim (int): number of neurons in hidden layers
        """
        super(Network, self).__init__()
        self.latent_dim = latent_dim

        self.enc1 = nn.Linear(input_dim, layer_dim1)
        self.enc2 = nn.Linear(layer_dim1, layer_dim2)

        self.latent = nn.Linear(layer_dim2, latent_dim*2)

        self.dec1 = nn.Linear(latent_dim, layer_dim2)
        self.dec2 = nn.Linear(layer_dim2, layer_dim1)

        self.out = nn.Linear(layer_dim1, output_dim)

    def encoder(self, x):
        z = F.elu(self.enc1(x))
        z = F.elu(self.enc2(z))
        z = self.latent(z)
        self.mu = z[:,0:self.latent_dim]
        self.log_sigma = z[:,self.latent_dim:]
        self.sigma = torch.exp(self.log_sigma)

        eps = torch.randn(x.size(0), self.latent_dim)
        z_sample = self.mu + self.sigma * eps

        self.kl_loss = kl_divergence(self.mu, self.log_sigma, dim=self.latent_dim)

        return z_sample

    def decoder(self, z):
        x = F.elu(self.dec1(z))
        x = F.elu(self.dec2(x))
        return self.out(x)

    def forward(self, batch):
        self.latent_rep = self.encoder(batch)
        dec_input = self.latent_rep
        return self.decoder(dec_input)


def kl_divergence(means, log_sigma, dim, target_sigma=0.1):
    """
    Computes Kullback–Leibler divergence for arrays of mean and log(sigma)
    """
    target_sigma = torch.Tensor([target_sigma])
    inner = 1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma)
    inner = torch.mean(inner, dim = 0)
    out = 1 / 2. * torch.mean(inner - dim)
    out = out
    return out



The first 20 epochs:
Epoch 1 -- loss 2592.617822, RMS error 0.769075, R² score -0.216357 
Epoch 2 -- loss 1020.093633, RMS error 0.767074, R² score -0.159109 
Epoch 3 -- loss 1855.906876, RMS error 0.758663, R² score -0.079950 
Epoch 4 -- loss 2494.895288, RMS error 0.760361, R² score -0.097916 
Epoch 5 -- loss 1447.344017, RMS error 0.759499, R² score -0.094597 
Epoch 6 -- loss 1500.075461, RMS error 0.760647, R² score -0.133853 
Epoch 7 -- loss 2715.783510, RMS error 0.760349, R² score -0.113949 
Epoch 8 -- loss 4040.605946, RMS error 0.760447, R² score -0.097853 
Epoch 9 -- loss 2115.793215, RMS error 0.758570, R² score -0.090718 
Epoch 10 -- loss 2927.331250, RMS error 0.758347, R² score -0.090624 
Epoch 11 -- loss 5020.661278, RMS error 0.761674, R² score -0.122575 
Epoch 12 -- loss 5662.833613, RMS error 0.759100, R² score -0.088969 
Epoch 13 -- loss 2570.623434, RMS error 0.757806, R² score -0.048262 
Epoch 14 -- loss 1901.247553, RMS error 0.758893, R² score -0.064967 
Epoch 15 -- loss 13875.499270, RMS error 0.760261, R² score -0.087031 
Epoch 16 -- loss 16458.334455, RMS error 0.758352, R² score -0.057363 
Epoch 17 -- loss 17253.359936, RMS error 0.760583, R² score -0.080694 
Epoch 18 -- loss 17851.229393, RMS error 0.760331, R² score -0.077037 
Epoch 19 -- loss 24284.866304, RMS error 0.758882, R² score -0.066919 
Epoch 20 -- loss 20997.437961, RMS error 0.760005, R² score -0.084857