CODEWITHSUNDEEP
pythonpytorch
Code_Pi
Code_Pi

Reputation: 31

Writing a pytorch neural net class that has functions for both model fitting and prediction

I want a PyTorch neural net to predict y using x where, for example,

x = torch.tensor([[6,2],[5,2],[1,3],[7,6]]).float()
y = torch.tensor([1,5,2,5]).float()

For this, I have written the following PyTorch class which can fit y using x. But, I am not able to extend the code to predict using new values of x, x_test. Can somebody please guide me in this regard?

import torch
import torch.nn as nn
from torch.optim import SGD


class MyNeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(2,4,bias=True)
        self.layer2 = nn.Linear(4,1,bias=True)
        self.loss = nn.MSELoss()

    def fit(self,x,y):
        def forward(self,x):
            x = self.layer1(x)
            x = self.layer2(x)

            return x.squeeze()

        
        def compile(self,x):
            forward(self,x)
            opt = SGD(self.parameters(),lr=0.01) # parameters from first forward pass
        
        
        compile(self,x)
        
        losses = []
        for _ in range(100):
            opt.zero_grad() # flush previous epoch's gradient
            res = forward(self,x)
            loss_value = self.loss(res,y)
            loss_value.backward() # compute gradient
            opt.step() # Perform iteration using gradient above
            losses.append(loss_value.item())
        y_train_hat = forward(self,x)
        return y_train_hat.detach().numpy()


    def predict(self,x_test):

         """
         
         need help in this part
         
         """

        # y_test_hat = forward(self,x_test)
        # return y_test_hat.detach().numpy()

I changed the code provided by Ivan. The following is the working version of the code:

class MyNeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(2, 4, bias=True)
        self.layer2 = nn.Linear(4, 1, bias=True)
        self.loss = nn.MSELoss()
        self.compile_()

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x.squeeze()

    def fit(self, x, y):
     
        losses = []
        for epoch in range(100):
            
            ## Inference
            res = self(x)  #instead of self(self,x)
            loss_value = self.loss(res,y)

            ## Backpropagation
            self.opt.zero_grad()  # flush previous epoch's gradient; this should be the first step, I think
            loss_value.backward() # compute gradient
           
            self.opt.step()       # Perform iteration using gradient above

            ## Logging
            losses.append(loss_value.item())
        
    def compile_(self):
        self.opt = SGD(self.parameters(), lr=0.01)
       
    def predict(self,x_test):
        self.eval()
        y_test_hat = self(x_test) #instead of forward(self,x_test)
        return y_test_hat.detach().numpy()
        self.train()

x = torch.tensor([[6,2],[5,2],[1,3],[7,6]]).float()
y = torch.tensor([1,5,2,5]).float()

model = MyNeuralNet()
model.fit(x,y)
model.predict(x)

Is this correct?

Upvotes: 0

Views: 474

Answers (1)

Ivan
Ivan

Reputation: 40618

Here is a minimal example with some modifications carried out to your code:

class MyNeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(2, 4, bias=True)
        self.layer2 = nn.Linear(4, 1, bias=True)
        self.loss = nn.MSELoss()
        self.compile_()

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x.squeeze()

    def fit(self, x, y):
        losses = []
        for epoch in range(100):
            ## Inference
            res = self(self,x)
            loss_value = self.loss(res,y)

            ## Backpropagation
            loss_value.backward() # compute gradient
            self.opt.zero_grad()  # flush previous epoch's gradient
            self.opt.step()       # Perform iteration using gradient above

            ## Logging
            losses.append(loss_value.item())
        
    def compile_(self, x):
        self.opt = SGD(self.parameters(), lr=0.01)
       
    def predict(self, x_test):
        self.eval()
        y_test_hat = self(x_test)
        return y_test_hat.detach().numpy()
        self.train()

Which can be used like this for training:

>>> model = MyNeuralNet()
>>> model.fit(x, y) # fit on some data (x, y)

Then perform inference with some other data x_test:

>>> model.predict(x_test)

Upvotes: 2

Related Questions