supervised learning with mse loss with Gaussian parameterized by both mu and sigma

parameterized_mu_sigma.py

"""
supervised learning with mse loss with Gaussian parameterized by both mu and sigma

"""
import numpy as np
import torch


N = 200 # number of training points
d = 5 # dimension of training inputs


def custom_loss(outputs1, outputs2, target):
    return ((outputs1 - target)**2 / (2 * outputs2**2) + 0.5 * np.log(2 * np.pi) + torch.log(outputs2)).sum() / outputs1.data.nelement()


def mse_loss(output, target):
    return ((output - target) ** 2).sum() / output.data.nelement()



def data():

    # two separated gaussian inputs
    x1 = np.random.normal(loc=-5.,
                         size=(int(N/2), d),
                         scale=1.0)
    y1 = -1. * np.ones((int(N/2), 1))
    x2 = np.random.normal(loc=5.,
                          size=(int(N/2), d),
                          scale=1.0)
    y2 = np.ones((int(N/2), 1))
    x = np.vstack([x1, x2])
    y = np.vstack([y1, y2])
    return torch.Tensor(x), torch.Tensor(y)


class Network1(torch.nn.Module):
    def __init__(self):
        super(Network1, self).__init__()
        self.linear1 = torch.nn.Linear(d, 1)
    def forward(self, x):
        out1 = torch.tanh(self.linear1(x))
        return out1

class Network2(torch.nn.Module):
    def __init__(self):
        super(Network2, self).__init__()
        self.linear1 = torch.nn.Linear(d, 1)
    def forward(self, x):
        out1 = torch.relu(self.linear1(x)) + 1e-3
        return out1


def test_tdata():
    x, y = data()
    assert x.shape == (N, d)
    assert y.shape == (N, 1)


def main():
    test_tdata()

    model1 = Network1()
    model2 = Network2()
    learning_rate = 1e-3
    num_epochs = 1000

    x, y = data()
    xt, yt = x[:int(3*N/4)], y[:int(3*N/4)]

    criterion = custom_loss#mse_loss
    optimizer = torch.optim.SGD(list(model1.parameters())+ list(model2.parameters()), lr=learning_rate)
    for epoch in range(num_epochs):
        optimizer.zero_grad()
        outputs1= model1(xt)
        outputs2 = model2(xt)
        loss = criterion(outputs1, outputs2, yt)
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f"epoch: {epoch} loss: {loss}")

    xte, yte = x[int(3*N/4):], y[int(3*N/4):]
    outputs1 = model1(xte)
    correct = (torch.round(outputs1) == yte).sum()
    accuracy = correct / yte.shape[0]
    print(f"test accuracy: {100 *accuracy}%")


if __name__ == "__main__":
    main()