Linear regression using Pytorch
I have classification problem. I am using Pytorch, My input is sequence of length 341 and output one of three classes {0,1,2}, I want to train linear regression model using pytorch, I created the following class but during the training, the loss values start to have numbers then inf then NAN. I do not know how to fix that . Also I tried to initialize the weights for linear model but it is the same thing. Any suggestions.
class regression(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.input_dim = input_dim
        # One layer
        self.linear = nn.Linear(input_dim, 1)

    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred

criterion = torch.nn.MSELoss()

def fit(model, data_loader, optim, epochs):
    for epoch in range(epochs):

        for i, (X, y) in enumerate(data_loader):
            X = X.float()
            y = y.unsqueeze(1).float()
            X = Variable(X, requires_grad=True)
            y = Variable(y, requires_grad=True)
            # Make a prediction for the input X
            pred = model(X)
            #loss = (y-pred).pow(2).mean()
            loss = criterion(y, pred)
            optim.zero_grad()
            loss.backward()
            optim.step()
            print(loss)
            print(type(loss))
        # Give some feedback after each 5th pass through the data
        if epoch % 5 == 0:
            print(&quot;Epoch&quot;, epoch, f&quot;loss: {loss}&quot;)
    return None
regnet = regression(input_dim=341)
optim = SGD(regnet.parameters(), lr=0.01)
fit(regnet, data_loader, optim=optim, epochs=5)
pred = regnet(torch.Tensor(test_set.data_info).float())
pred = pred.detach().numpy()


I would additionally suggest to replace MSE with CrossEntropy Loss as it is better suited for multi-class classificiation problems.
import random
import torch
from torch import nn, optim
from matplotlib import pyplot as plt

# Generate random dataset with your shape to test
# Replace this with your own dataset
data = []
for label in [0, 1, 2]:
    for i in range(1000):
        data.append((torch.rand(341), label))

# train test split
random.shuffle(data)
train, val = data[:1500], data[1500:]    


 def run_gradient_descent(model, data_train, data_val, batch_size=64, learning_rate=0.01, weight_decay=0, num_epochs=10):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    iters, losses = [], []
    iters_sub, train_acc, val_acc = [], [] ,[]
    train_loader = torch.utils.data.DataLoader(data_train, batch_size=batch_size, shuffle=True)

    # training
    n = 0 # the number of iterations
    for epoch in range(num_epochs):
        for xs, ts in iter(train_loader):
            if len(ts) != batch_size:
                continue
            zs = model(xs)
            loss = criterion(zs, ts) # compute the total loss
            loss.backward() # compute updates for each parameter
            optimizer.step() # make the updates for each parameter
            optimizer.zero_grad() # a clean up step for PyTorch
            # save the current training information
            iters.append(n)
            losses.append(float(loss)/batch_size) # compute *average* loss
            if n % 10 == 0:
                iters_sub.append(n)
                train_acc.append(get_accuracy(model, data_train))
                val_acc.append(get_accuracy(model, data_val))
             # increment the iteration number
            n += 1

    # plotting
    plt.title(&quot;Training Curve (batch_size={}, lr={})&quot;.format(batch_size, learning_rate))
    plt.plot(iters, losses, label=&quot;Train&quot;)
    plt.xlabel(&quot;Iterations&quot;)
    plt.ylabel(&quot;Loss&quot;)
    plt.show()
    plt.title(&quot;Training Curve (batch_size={}, lr={})&quot;.format(batch_size, learning_rate))
    plt.plot(iters_sub, train_acc, label=&quot;Train&quot;)
    plt.plot(iters_sub, val_acc, label=&quot;Validation&quot;)
    plt.xlabel(&quot;Iterations&quot;)
    plt.ylabel(&quot;Accuracy&quot;)
    plt.legend(loc='best')
    plt.show()
    return model


def get_accuracy(model, data):
    loader = torch.utils.data.DataLoader(data, batch_size=500)
    correct, total = 0, 0
    for xs, ts in loader:
        zs = model(xs)
        pred = zs.max(1, keepdim=True)[1] # get the index of the max logit
        correct += pred.eq(ts.view_as(pred)).sum().item()
        total += int(ts.shape[0])
    return correct / total


class MyRegression(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(MyRegression, self).__init__()
        # One layer
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.linear(x)


model = MyRegression(341, 3)
run_gradient_descent(model, train, val, batch_size=64, learning_rate=0.01, num_epochs=10)

