Create a Simple Neural Network Class

1. Background Knowledge

PyTorch nn.Module is the foundation for all neural networks. Every neural network in PyTorch inherits from nn.Module, which provides essential functionality like parameter tracking, device placement, and the forward pass mechanism.

Key concepts:

• Parameters: nn.Module automatically registers layers like nn.Linear as trainable parameters via self.layer = nn.Linear(...)
• Forward pass: Defined in forward() method; called automatically when using model(input)
• Linear layer: nn.Linear(in_features, out_features) implements $y = xW^T + b$ where $W \in \mathbb{R}^{out \times in}$, $b \in \mathbb{R}^{out}$
• Shape transformation: Input (batch_size, input_size) → Output (batch_size, output_size)

Prerequisites:

import torch
import torch.nn as nn

2. Algorithm Approach

This is a feedforward neural network with a single linear transformation:

Input: (batch_size, 10) → Linear(10, 5) → Output: (batch_size, 5)

Mathematical operation: $$\mathbf{y}$=$\mathbf{x}$$\mathbf{W}$^T +$\mathbf{b}$$

• $$\mathbf{x}$\in \mathbb{R}^{batch \times 10}$
• $$\mathbf{W}$\in \mathbb{R}^{5 \times 10}$ (weights)
• $$\mathbf{b}$\in \mathbb{R}^{5}$ (bias)
• $$\mathbf{y}$\in \mathbb{R}^{batch \times 5}$

3. Step-by-Step Strategy

1. Define the class inheriting from nn.Module:

class SimpleNet(nn.Module):

2. Initialize in init():

• Call super(SimpleNet, self).init() first
• Create self.linear = nn.Linear(10, 5)

3. Define forward():

• Return self.linear(x) (no other operations needed)

4. Create create_simple_network():

• Instantiate SimpleNet()
• Verify it's nn.Module subclass: isinstance(net, nn.Module)
• Test output shape: net(torch.randn(1, 10)).shape
• Return dictionary with all 3 keys

Complete solution structure:

def create_simple_network():
 class SimpleNet(nn.Module): # Defined inside function
 def __init__(self):
 super().__init__()
 self.linear = nn.Linear(10, 5)

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

 net = SimpleNet()
 return {
 "class_name": "SimpleNet",
 "is_module": isinstance(net, nn.Module),
 "output_shape": list(net(torch.randn(1, 10)).shape)
 }

4. Common Pitfalls

Critical: Class must be named exactly "SimpleNet" for the test.

5. Time & Space Complexity

Forward pass: $O(n \cdot d_{in} \cdot d_{out}) = O(1 \cdot 10 \cdot 5) = O(50)$ per sample

• Matrix multiplication dominates
• Constant time for fixed architecture

Space complexity:

• Parameters: $10 \times 5 + 5 = 55$ scalars
• Activations: $O(batch\_size \cdot 5)$
• Total: $O(1)$ (fixed-size model)

Verification test runs in $O(1)$ time since batch_size=1 and layers are tiny.

This pattern scales to complex networks: more layers → $O($\sum$L_i \cdot d_i \cdot d_{i+1})$.