__init__ for layers in PyTorch

Introduction

The __init__ method sets up a layer's parts when you create it. It tells the layer what it needs to work.

When you want to create a new neural network layer with custom settings.

When you need to define the size or shape of the layer's weights.

When you want to add layers like linear, convolution, or dropout inside your model.

When you want to prepare your layer to process input data correctly.

Syntax

PyTorch

class MyLayer(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.linear = torch.nn.Linear(input_size, output_size)

The __init__ method runs once when you create the layer.

Always call super().__init__() to set up the base class properly.

Examples

This layer has a linear part that changes 10 inputs into 5 outputs.

PyTorch

class SimpleLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

This layer randomly ignores some inputs during training to help the model learn better.

PyTorch

class DropoutLayer(torch.nn.Module):
    def __init__(self, drop_prob):
        super().__init__()
        self.dropout = torch.nn.Dropout(drop_prob)

Sample Model

This code defines a layer that changes 3 numbers into 2 numbers. It then runs one example input through it and prints the result.

PyTorch

import torch

class MyLayer(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.linear = torch.nn.Linear(input_size, output_size)

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

# Create layer with 3 inputs and 2 outputs
layer = MyLayer(3, 2)

# Create a sample input tensor
input_tensor = torch.tensor([[1.0, 2.0, 3.0]])

# Get output from the layer
output = layer(input_tensor)
print(output)

Important Notes

Inside __init__, you only set up parts; actual data flows in the forward method.

Use super().__init__() to make sure PyTorch knows this is a special layer.

Keep __init__ simple: just create layers and parameters here.

Summary

__init__ prepares your layer when you make it.

Always call super().__init__() first inside __init__.

Define parts like linear or dropout layers inside __init__.