Creating a neural network from scratch using PyTorch’s torch.nn.Module is a simpler process that allows for flexibility and customizability. The first step involves subclassing torch.nn.Module to define the architecture of your network. You’ll need to implement two main functions: __init__ and forward.
Within the __init__ method, you should define all the layers your network will use. For instance, if you’re building a simple feedforward neural network with one hidden layer, you might include a linear layer and an activation function.
import torch
import torch.nn as nn
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(10, 5) # Input layer to hidden layer
self.fc2 = nn.Linear(5, 2) # Hidden layer to output layer
self.relu = nn.ReLU() # Activation function
def forward(self, x):
x = self.relu(self.fc1(x)) # Pass input through first layer and apply activation
x = self.fc2(x) # Pass to output layer
return x
Now that you’ve defined the structure, you can instantiate your model. Make sure to pass the input tensor with the correct dimensions when you call the model. PyTorch will automatically handle the gradient computations during backpropagation, thanks to its dynamic computation graph.
model = SimpleNN() input_tensor = torch.randn(1, 10) # Batch size of 1, 10 features output = model(input_tensor) # Forward pass print(output)
Next, when working with custom architectures, it’s crucial to think about how you handle different types of layers, especially when they may require specific initialization or behavior. For instance, if you want to add dropout or batch normalization, you can simply include those layers in your __init__ method.
self.dropout = nn.Dropout(p=0.5) # 50% dropout rate self.batch_norm = nn.BatchNorm1d(5) # Batch normalization for hidden layer
In the forward method, you can then apply these layers as part of the forward pass:
def forward(self, x):
x = self.relu(self.fc1(x))
x = self.dropout(x) # Apply dropout
x = self.batch_norm(x) # Apply batch normalization
x = self.fc2(x)
return x
It is also important to consider how you plan to handle the training and evaluation of your model. You may want to implement additional methods for saving and loading your model, or for adjusting the learning rate during training. Using the built-in optimizers in PyTorch can significantly streamline this process.
import torch.optim as optim optimizer = optim.Adam(model.parameters(), lr=0.001) # Using Adam optimizer criterion = nn.CrossEntropyLoss() # Loss function for multi-class classification
With the model, loss function, and optimizer set up, you’re ready to loop through your data, performing forward passes, calculating loss, and updating the model’s weights. Remember to zero the gradients before each backward pass to prevent accumulation from previous iterations.
for data, labels in dataloader:
optimizer.zero_grad() # Zero gradients
outputs = model(data) # Forward pass
loss = criterion(outputs, labels) # Calculate loss
loss.backward() # Backpropagation
optimizer.step() # Update weights
As you extend your model, consider how each component interacts and the potential need for custom forward passes to accommodate unique architectures. Deep learning often requires experimentation, and having a solid understanding of the underlying mechanics will help you design more effective models.
Now retrieving an image set.
Google Play gift code - give the gift of games, apps and more (Email or Text Message Delivery - US Only)
$25.00 (as of July 15, 2026 15:40 GMT +00:00 - More infoProduct prices and availability are accurate as of the date/time indicated and are subject to change. Any price and availability information displayed on [relevant Amazon Site(s), as applicable] at the time of purchase will apply to the purchase of this product.)Best practices for extending torch.nn.Module in custom architectures
When extending torch.nn.Module, always call super() in your __init__ method to ensure the base class is properly initialized. That is vital for PyTorch to track your submodules and parameters correctly.
Another best practice is to register all trainable parameters as nn.Module attributes. Avoid creating parameters or layers dynamically inside the forward method; doing so breaks PyTorch’s ability to track these parameters for optimization.
For more complex architectures, consider modularizing your code by breaking down large networks into smaller, reusable submodules. This makes debugging easier and improves readability.
class ResidualBlock(nn.Module):
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(channels)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(channels)
def forward(self, x):
residual = x
out = self.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += residual
out = self.relu(out)
return out
class CustomResNet(nn.Module):
def __init__(self):
super(CustomResNet, self).__init__()
self.initial_conv = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.bn = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.layer1 = ResidualBlock(64)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(64, 10)
def forward(self, x):
x = self.relu(self.bn(self.initial_conv(x)))
x = self.layer1(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
In this example, ResidualBlock encapsulates a common pattern and is reused inside the CustomResNet. This design pattern not only improves clarity but also enables easy expansion—for instance, adding more residual blocks or modifying their internals without touching the main network class.
When defining your forward method, avoid side effects such as in-place modifications of inputs or global state changes. This ensures that your model remains compatible with PyTorch’s autograd system and tools like torch.jit.
For initialization, PyTorch layers come with sensible defaults, but sometimes you want to customize the weights explicitly. Do this inside the __init__ or in a separate initialization method called from __init__. Use torch.nn.init functions to set weights and biases.
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
class CustomResNet(nn.Module):
def __init__(self):
super(CustomResNet, self).__init__()
# ... layer definitions ...
self._initialize_weights()
Keep your model’s API clean and intuitive. If you need to expose additional functionality—like returning intermediate activations or applying different forward modes—consider adding optional arguments to forward or separate methods. Avoid overloading forward with too many responsibilities.
Finally, when building large models, leverage PyTorch’s device-agnostic coding by moving modules and tensors explicitly to the desired device (cpu or cuda). This prevents subtle bugs that arise from mixing devices.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CustomResNet().to(device)
for data, labels in dataloader:
data = data.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(data)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
