API Reference

Complete API documentation for Neural Networks Classification Suite

CNNClassifier

torch.nn.Module

Basic CNN architecture for multi-class image classification with configurable layers and dropout. 

class CNNClassifier(nn.Module):
    def __init__(
        self,
        num_classes: int = 8,
        in_channels: int = 1,
        dropout_rate: float = 0.5,
        use_batch_norm: bool = True
    )

Parameters

Parameter Type Default Description
num_classes int 8 Number of output classes
in_channels int 1 Number of input channels (1 for grayscale, 3 for RGB)
dropout_rate float 0.5 Dropout probability for regularization
use_batch_norm bool True Whether to use batch normalization

Example

import torch
from nn_classification.models import CNNClassifier

# Initialize model
model = CNNClassifier(
    num_classes=8,
    in_channels=1,
    dropout_rate=0.5
)

# Forward pass
x = torch.randn(32, 1, 224, 224)  # batch_size=32, grayscale, 224x224
output = model(x)
print(output.shape)  # torch.Size([32, 8])

Returns

Tensor of shape (batch_size, num_classes) containing logits

CNNWithAttention

torch.nn.Module

CNN architecture enhanced with CBAM (Convolutional Block Attention Module) for improved feature extraction. 

class CNNWithAttention(nn.Module):
    def __init__(
        self,
        num_classes: int = 8,
        in_channels: int = 1,
        attention_type: str = 'cbam',
        reduction_ratio: int = 16
    )

Parameters

Parameter Type Default Description
num_classes int 8 Number of output classes
in_channels int 1 Number of input channels
attention_type str 'cbam' Type of attention: 'cbam', 'se', or 'spatial'
reduction_ratio int 16 Channel reduction ratio for attention modules

Methods

get_attention_maps(x: Tensor) -> Dict[str, Tensor]

Returns attention weights for visualization

Example

from nn_classification.models import CNNWithAttention

# Initialize model with CBAM attention
model = CNNWithAttention(
    num_classes=8,
    attention_type='cbam',
    reduction_ratio=16
)

# Get predictions and attention maps
x = torch.randn(1, 1, 224, 224)
output = model(x)
attention_maps = model.get_attention_maps(x)

# Visualize attention
import matplotlib.pyplot as plt
plt.imshow(attention_maps['layer3'].squeeze().cpu())
plt.colorbar()
plt.title('Attention Heatmap')
plt.show()

VisionTransformer

torch.nn.Module

Vision Transformer (ViT) implementation for image classification with configurable architecture. 

class VisionTransformer(nn.Module):
    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        num_classes: int = 8,
        dim: int = 768,
        depth: int = 12,
        heads: int = 12,
        mlp_dim: int = 3072,
        dropout: float = 0.1,
        emb_dropout: float = 0.1
    )

Note

Vision Transformers require significant computational resources and larger datasets for optimal performance.

Trainer

Training Pipeline

High-level training API with automatic mixed precision, gradient accumulation, and monitoring. 

class Trainer:
    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        config: Dict[str, Any]
    )

Configuration

config = {
    'learning_rate': 1e-3,
    'epochs': 50,
    'batch_size': 32,
    'optimizer': 'adamw',
    'scheduler': 'cosine',
    'mixed_precision': True,
    'gradient_clip': 1.0,
    'early_stopping_patience': 10,
    'save_best': True,
    'wandb_project': 'nn-classification'
}

trainer = Trainer(model, train_loader, val_loader, config)
trainer.train()

Data Augmentation

torchvision.transforms

Advanced augmentation techniques including MixUp, CutMix, and CutOut.

MixUp

from nn_classification.data.augmentation import mixup_data

# In training loop
inputs, targets = next(iter(train_loader))
inputs, targets_a, targets_b, lam = mixup_data(inputs, targets, alpha=1.0)

outputs = model(inputs)
loss = lam * criterion(outputs, targets_a) + (1 - lam) * criterion(outputs, targets_b)

CutMix

from nn_classification.data.augmentation import cutmix

# Apply CutMix augmentation
inputs, targets_a, targets_b, lam = cutmix(inputs, targets, beta=1.0)
outputs = model(inputs)
loss = lam * criterion(outputs, targets_a) + (1 - lam) * criterion(outputs, targets_b)