Skip to main content

​
SmallCNN

A compact convolutional neural network designed for edge device deployment. Optimized for 28x28 grayscale images with configurable channel dimensions.

​
Constructor

SmallCNN(
    conv1_channels: int = 16,
    conv2_channels: int = 32,
    num_classes: int = 10
) -> None
​
conv1_channels
int
default:"16"
Number of output channels for the first convolutional layer. Controls the model’s feature extraction capacity in the initial layer.
​
conv2_channels
int
default:"32"
Number of output channels for the second convolutional layer. Determines the depth of features in the second layer.
​
num_classes
int
default:"10"
Number of output classes for classification. The final linear layer outputs this many logits.

​
Architecture Details

The SmallCNN architecture consists of:
  • Layer 1: Conv2d(1 → conv1_channels, kernel=3, padding=1) → ReLU → MaxPool2d(2)
  • Layer 2: Conv2d(conv1_channels → conv2_channels, kernel=3, padding=1) → ReLU → MaxPool2d(2)
  • Classifier: Flatten → Linear(conv2_channels × 7 × 7 → num_classes)
The model expects input tensors of shape (batch_size, 1, 28, 28). After two pooling operations, the spatial dimensions reduce from 28×28 to 7×7.

​
forward

Performs forward pass through the network. 
def forward(self, x: torch.Tensor) -> torch.Tensor
​
x
torch.Tensor
Input tensor of shape (batch_size, 1, 28, 28) representing grayscale images.
​
output
torch.Tensor
Classification logits of shape (batch_size, num_classes).

​
Usage Example

import torch
from edge_opt.model import SmallCNN

# Create model with default parameters
model = SmallCNN()

# Create sample input
x = torch.randn(8, 1, 28, 28)  # batch_size=8

# Forward pass
output = model(x)
print(output.shape)  # torch.Size([8, 10])

​
set_deterministic

Configures all random number generators for deterministic behavior across PyTorch, NumPy, and Python’s random module. Essential for reproducible experiments on edge devices. 
def set_deterministic(seed: int) -> None
​
seed
int
required
Random seed value to set across all libraries. Use the same seed to reproduce results.

​
What It Does

This function sets:
  1. Python random: random.seed(seed)
  2. NumPy: np.random.seed(seed)
  3. PyTorch CPU: torch.manual_seed(seed)
  4. PyTorch CUDA: torch.cuda.manual_seed_all(seed)
  5. cuDNN: Sets deterministic=True and benchmark=False
  6. PyTorch algorithms: Enables use_deterministic_algorithms(True)
Enabling deterministic mode may reduce performance. Use for reproducibility during experiments, but consider disabling for production deployments.

​
Usage Example

from edge_opt.model import set_deterministic, SmallCNN
import torch

# Set deterministic behavior at the start
set_deterministic(seed=42)

# All subsequent operations will be deterministic
model = SmallCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Training loop will produce identical results with same seed
for epoch in range(10):
    # ... training code ...
    pass

​
See Also

  • Config - Use ExperimentConfig.seed to manage seeds
  • Data - build_loaders accepts a seed parameter for data loading

Build docs developers (and LLMs) love

Get started for freeTalk to us