Skip to main content

​
Overview

The CustomCNNBuilder class enables building custom CNN architectures from declarative layer configurations. It provides a flexible way to construct CNNs by specifying layers, filters, activations, and other parameters in a configuration dictionary.

​
Classes

​
CustomCNNBuilder

Builder class that constructs CNN models from configuration. 
from models.pytorch.cnn_builder import CustomCNNBuilder
Location: app/models/pytorch/cnn_builder.py:13

​
Constructor

__init__(config: dict[str, Any])
​
config
dict[str, Any]
required
Model configuration dictionary containing:
  • num_classes: Number of output classes
  • cnn_config: CNN-specific configuration with layers list
Example: 
config = {
    "num_classes": 10,
    "model_type": "CNN",
    "cnn_config": {
        "layers": [
            {"type": "Conv2D", "params": {"filters": 32, "kernel_size": 3, "activation": "relu"}},
            {"type": "MaxPooling2D", "params": {"pool_size": 2}},
            {"type": "Conv2D", "params": {"filters": 64, "kernel_size": 3, "activation": "relu"}},
            {"type": "Flatten", "params": {}},
            {"type": "Dense", "params": {"units": 128, "activation": "relu"}},
        ]
    }
}

builder = CustomCNNBuilder(config)
model = builder.build()

​
Methods

build() -> nn.Module
Build the CNN model from configuration.
​
model
nn.Module
Built PyTorch CNN model (CustomCNN instance)
Raises:
  • ValueError: If configuration validation fails
Example: 
builder = CustomCNNBuilder(config)
model = builder.build()
print(model)
get_parameters_count() -> tuple[int, int]
Get parameter counts for the model.
​
total_params
int
Total number of parameters
​
trainable_params
int
Number of trainable parameters
validate_config() -> bool
Validate CNN configuration.
​
is_valid
bool
Returns True if configuration is valid
Validation checks:
  • Base validation (num_classes > 0)
  • cnn_config must be present
  • layers list must not be empty
  • Must contain at least one Conv2D layer
  • Must contain at least one transition layer (Flatten or GlobalAvgPool)
  • Must contain at least one Dense layer

​
CustomCNN

PyTorch model class built from layer stack. 
from models.pytorch.cnn_builder import CustomCNN
Location: app/models/pytorch/cnn_builder.py:72

​
Constructor

__init__(
    layers: list[dict[str, Any]],
    num_classes: int,
    input_channels: int = 3,
    input_size: int = 224
)
​
layers
list[dict[str, Any]]
required
List of layer configurations from layer_stack
​
num_classes
int
required
Number of output classes
​
input_channels
int
default:"3"
Number of input channels (3 for RGB images)
​
input_size
int
default:"224"
Input image size (assumes square images)
Example: 
layers = [
    {"type": "Conv2D", "params": {"filters": 32, "kernel_size": 3, "activation": "relu"}},
    {"type": "MaxPooling2D", "params": {"pool_size": 2}},
    {"type": "Flatten", "params": {}},
    {"type": "Dense", "params": {"units": 128, "activation": "relu"}}
]

model = CustomCNN(layers=layers, num_classes=10)

​
Methods

forward(x: torch.Tensor) -> torch.Tensor
Forward pass through the network.
​
x
torch.Tensor
required
Input tensor of shape (batch, channels, height, width)
​
output
torch.Tensor
Output logits of shape (batch, num_classes)
Example: 
import torch

model = CustomCNN(layers, num_classes=10)
x = torch.randn(8, 3, 224, 224)  # Batch of 8 RGB images
logits = model(x)
print(logits.shape)  # torch.Size([8, 10])
get_feature_extractor() -> nn.Sequential
Extract the feature extraction layers (convolutional part).
​
features
nn.Sequential
Sequential module containing all feature extraction layers
Example: 
model = CustomCNN(layers, num_classes=10)
feature_extractor = model.get_feature_extractor()

# Use for transfer learning
x = torch.randn(8, 3, 224, 224)
features = feature_extractor(x)
print(features.shape)

​
Supported Layer Types

​
Conv2D

Convolutional layer with activation. Parameters:
​
filters
int
required
Number of output filters/channels
​
kernel_size
int
default:"3"
Size of the convolution kernel
​
activation
str
default:"relu"
Activation function: "relu", "leaky_relu", "gelu", "swish", "none"
​
padding
str
default:"same"
Padding mode: "same" or "valid"
Example: 
{"type": "Conv2D", "params": {"filters": 64, "kernel_size": 3, "activation": "relu"}}

​
MaxPooling2D

Max pooling layer.
​
pool_size
int
default:"2"
Size of the pooling window
Example: 
{"type": "MaxPooling2D", "params": {"pool_size": 2}}

​
AveragePooling2D

Average pooling layer.
​
pool_size
int
default:"2"
Size of the pooling window

​
BatchNorm

Batch normalization layer. Example: 
{"type": "BatchNorm", "params": {}}

​
Dropout

Dropout layer (2D for features, 1D for classifier).
​
rate
float
default:"0.25"
Dropout rate (0-1)
Example: 
{"type": "Dropout", "params": {"rate": 0.5}}

​
Flatten

Flatten spatial dimensions. Required before Dense layers. Example: 
{"type": "Flatten", "params": {}}

​
GlobalAvgPool

Global average pooling. Alternative to Flatten, reduces parameters. Example: 
{"type": "GlobalAvgPool", "params": {}}

​
Dense

Fully connected layer with activation.
​
units
int
required
Number of output units
​
activation
str
default:"relu"
Activation function
Example: 
{"type": "Dense", "params": {"units": 256, "activation": "relu"}}

​
Activation Functions

Supported activation functions (case-sensitive):
  • "relu" - ReLU activation
  • "leaky_relu" - Leaky ReLU with alpha=0.1
  • "gelu" - Gaussian Error Linear Unit
  • "swish" - Swish/SiLU activation
  • "none" - Identity (no activation)

​
Complete Example

from models.pytorch.cnn_builder import CustomCNNBuilder
import torch

# Define configuration
config = {
    "num_classes": 10,
    "model_type": "CNN",
    "architecture": "Custom",
    "cnn_config": {
        "layers": [
            # First conv block
            {"type": "Conv2D", "params": {"filters": 32, "kernel_size": 3, "activation": "relu"}},
            {"type": "BatchNorm", "params": {}},
            {"type": "Conv2D", "params": {"filters": 32, "kernel_size": 3, "activation": "relu"}},
            {"type": "MaxPooling2D", "params": {"pool_size": 2}},
            {"type": "Dropout", "params": {"rate": 0.25}},
            
            # Second conv block
            {"type": "Conv2D", "params": {"filters": 64, "kernel_size": 3, "activation": "relu"}},
            {"type": "BatchNorm", "params": {}},
            {"type": "Conv2D", "params": {"filters": 64, "kernel_size": 3, "activation": "relu"}},
            {"type": "MaxPooling2D", "params": {"pool_size": 2}},
            {"type": "Dropout", "params": {"rate": 0.25}},
            
            # Classifier
            {"type": "GlobalAvgPool", "params": {}},
            {"type": "Dense", "params": {"units": 128, "activation": "relu"}},
            {"type": "Dropout", "params": {"rate": 0.5}},
        ]
    }
}

# Build model
builder = CustomCNNBuilder(config)
model = builder.build()

# Get model info
total, trainable = builder.get_parameters_count()
print(f"Total parameters: {total:,}")
print(f"Trainable parameters: {trainable:,}")

# Run inference
x = torch.randn(4, 3, 224, 224)
output = model(x)
print(f"Output shape: {output.shape}")  # torch.Size([4, 10])

​
See Also

Build docs developers (and LLMs) love

Get started for freeTalk to us