Custom Module - ( Linear Regression with PyTorch )

Heading h2

Syntax

import torch
import torch.nn as nn
import torch.optim as optim

class LinearRegression(nn.Module):
    def __init__(self, input_dim):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(input_dim, 1)

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

model = LinearRegression(input_dim=1)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

Example

import torch
import torch.nn as nn
import torch.optim as optim

class LinearRegression(nn.Module):
    def __init__(self, input_dim):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(input_dim, 1)

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

model = LinearRegression(input_dim=1)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Training loop
epochs = 100
for epoch in range(epochs):
    # Forward pass
    y_pred = model(x)

    # Compute loss
    loss = criterion(y_pred, y)

    # Backward pass and optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print loss
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, epochs, loss.item()))

Output

Epoch [1/100], Loss: 1453.5778
Epoch [2/100], Loss: 829.6146
Epoch [3/100], Loss: 478.0710
Epoch [4/100], Loss: 279.0025
...

Explanation

The custom module in PyTorch allows you to define a neural network architecture as a subclass of nn.Module. In this example, we define a linear regression model with a single input and output.

We define the __init__ function to initialize the nn.Linear module with the input and output dimensions. The forward function defines the forward pass of the neural network.

To train the model, we define a criterion, optimizer, and a loop over the epochs. In each epoch, we compute the forward pass of the model, compute the loss, perform the backward pass and optimization.

Use

Custom modules can be used to define more complex neural network architectures with multiple layers, inputs, and outputs. They can be used for various tasks such as regression, classification, and more.

Important Points

• Custom modules are defined as subclasses of nn.Module
• __init__ function initializes the neural network layers
• forward function defines the forward pass of the neural network
• Custom modules can be used for various tasks such as regression, classification, etc.

Summary

In summary, custom modules in PyTorch allow you to define neural network architectures as subclasses of nn.Module. They can be used to define more complex neural networks for various tasks such as regression, classification, and more. By defining a custom module, you gain more control over the architecture of the neural network and can customize it to fit your specific requirements.