PyTorch CNN - Flatten Mismatch After Pooling Change

import torch
import torch.nn as nn
import torch.nn.functional as F

# io.thecodeforge: Standard CNN architecture for image classification
# Input assumption: 3-channel RGB images at 224x224 resolution, 10-class output
class ForgeCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super(ForgeCNN, self).__init__()

        # Block 1: 3 input channels -> 16 feature maps
        # padding=1 with kernel_size=3 preserves spatial dimensions (same padding)
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(16)  # stabilizes training, especially early on

        # Block 2: 16 feature maps -> 32 feature maps
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)

        # MaxPool2d with kernel=2, stride=2 halves both spatial dimensions
        # After pool1: 224 -> 112. After pool2: 112 -> 56.
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # AdaptiveAvgPool2d fixes output to 7x7 regardless of input resolution
        # This eliminates the dimension recalculation problem when you change upstream layers
        self.adaptive_pool = nn.AdaptiveAvgPool2d(output_size=(7, 7))

        # Flattened size: 32 channels * 7 * 7 = 1568 — fixed regardless of input resolution
        self.fc1 = nn.Linear(32 * 7 * 7, 256)
        self.dropout = nn.Dropout(p=0.4)
        self.fc2 = nn.Linear(256, num_classes)

        # Verify dimensions at init time — fails loudly in development, not silently in production
        self._verify_dimensions()

    def _verify_dimensions(self):
        """Pass a dummy tensor at init to catch shape errors immediately."""
        with torch.no_grad():
            dummy = torch.randn(1, 3, 224, 224)
            out = self.forward(dummy)
            assert out.shape == (1, self.fc2.out_features), (
                f"Output shape mismatch: expected (1, {self.fc2.out_features}), got {out.shape}"
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Conv -> BN -> ReLU -> Pool
        x = self.pool(F.relu(self.bn1(self.conv1(x))))   # [B, 16, 112, 112]
        x = self.pool(F.relu(self.bn2(self.conv2(x))))   # [B, 32, 56, 56]
        x = self.adaptive_pool(x)                         # [B, 32, 7, 7] — fixed output
        x = x.view(x.size(0), -1)                         # [B, 1568]
        x = self.dropout(F.relu(self.fc1(x)))             # [B, 256]
        x = self.fc2(x)                                   # [B, num_classes]
        return x


# Instantiation — dimension verification runs automatically
model = ForgeCNN(num_classes=10)
print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")

# io.thecodeforge: Production Vision Service Environment
# Pin the full version triple: PyTorch + CUDA + cuDNN
# Never use 'latest' — it will silently change and break reproducibility
FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Install Python dependencies first — this layer caches independently of source code
# Avoids re-installing packages on every source change
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy model weights and inference source separately
# Model weights change less frequently than source code
# Keeping them in separate COPY instructions preserves Docker layer cache
COPY ./models /app/models
COPY ./src /app/src

# Run as a non-root user — required by most enterprise security policies
# and good practice regardless
USER 1000

# Health check: verify the inference service starts and the model loads correctly
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
    CMD curl -f http://localhost:8080/health || exit 1

ENTRYPOINT ["python", "src/inference_service.py"]

-- io.thecodeforge: Schema for managing training image references
-- Store paths and labels here. Store actual image bytes on disk or S3.
CREATE TABLE IF NOT EXISTS io.thecodeforge.vision_assets (
    asset_id          UUID        PRIMARY KEY DEFAULT gen_random_uuid(),
    file_path         TEXT        NOT NULL UNIQUE,  -- relative path or S3 key, never absolute
    label_id          INT         NOT NULL,
    label_name        VARCHAR(128),                 -- denormalized for readability in queries
    split_type        VARCHAR(10) CHECK (split_type IN ('train', 'val', 'test', 'held_out')),
    source_dataset    VARCHAR(64),                  -- track provenance when merging multiple datasets
    is_augmented      BOOLEAN     DEFAULT FALSE,    -- flag synthetically generated images
    last_augmented_at TIMESTAMP,
    created_at        TIMESTAMP   DEFAULT NOW()
);

-- Index on split_type for fast DataLoader queries
CREATE INDEX IF NOT EXISTS idx_vision_assets_split
    ON io.thecodeforge.vision_assets (split_type, label_id);

-- Fetch a random training batch — used by the PyTorch DataLoader during epoch iteration
SELECT file_path, label_id
FROM   io.thecodeforge.vision_assets
WHERE  split_type = 'train'
  AND  is_augmented = FALSE           -- exclude pre-generated augmented copies if present
ORDER  BY RANDOM()
LIMIT  32;

-- Audit class balance before training — imbalanced classes need weighted sampling
SELECT label_name, split_type, COUNT(*) AS image_count
FROM   io.thecodeforge.vision_assets
GROUP  BY label_name, split_type
ORDER  BY label_name, split_type;

# io.thecodeforge: Common mistake patterns and their correct alternatives
import torch
import torch.nn as nn
from torchvision import transforms, models


# ─── MISTAKE 1: Hardcoded flattened dimension ────────────────────────────────
# Bad: fails silently until forward pass if you change upstream layers
class BadCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # This hardcoded number is wrong if you add another pooling layer
        self.classifier = nn.Linear(32 * 112 * 112, 10)


# Good: use AdaptiveAvgPool2d — flattened size is always channels * H * W
class GoodCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))  # fixed output: 32 * 7 * 7 = 1568
        self.classifier = nn.Linear(32 * 7 * 7, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.adaptive_pool(x)
        return self.classifier(x.view(x.size(0), -1))


# ─── MISTAKE 2: Forgetting to store normalization with the model ─────────────
# Bad: normalization parameters live only in the training script
# train_transforms = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# ...six months later, the inference service ships without normalization

# Good: save normalization params in the checkpoint
NORM_MEAN = [0.485, 0.456, 0.406]
NORM_STD = [0.229, 0.224, 0.225]

def save_model_with_norm(model, path):
    torch.save({
        'state_dict': model.state_dict(),
        'norm_mean': NORM_MEAN,
        'norm_std': NORM_STD,
    }, path)

def load_model_with_norm(model_class, path):
    checkpoint = torch.load(path, map_location='cpu')
    model = model_class()
    model.load_state_dict(checkpoint['state_dict'])
    inference_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=checkpoint['norm_mean'], std=checkpoint['norm_std']),
    ])
    return model, inference_transform


# ─── MISTAKE 3: Calling model.forward() directly ────────────────────────────
# Bad: bypasses all registered forward hooks
# output = model.forward(input_tensor)

# Good: call the model as a callable
# output = model(input_tensor)


# ─── MISTAKE 4: Using augmentation during validation ────────────────────────
# Bad: validation results are non-deterministic, benchmark comparisons are meaningless
bad_val_transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # should never be in a validation pipeline
    transforms.ToTensor(),
])

# Good: validation pipeline is fully deterministic — no random transforms
good_val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=NORM_MEAN, std=NORM_STD),
])

# io.thecodeforge: Production transfer learning pattern with two-stage fine-tuning
import torch
import torch.nn as nn
from torchvision import models

NUM_CLASSES = 10
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# ─── Stage 0: Load pre-trained weights ──────────────────────────────────────
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# ResNet-18 architecture breakdown:
# - 17 convolutional layers learning features from edges up to object parts
# - Final fc layer: Linear(512, 1000) for ImageNet 1000 classes
print(f"Original output classes: {model.fc.out_features}")


# ─── Stage 1: Freeze backbone, train head only (warm-up) ────────────────────
for param in model.parameters():
    param.requires_grad = False

# Replace the classifier head — this new layer has requires_grad=True by default
model.fc = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, NUM_CLASSES),
)

model = model.to(DEVICE)

# Only the new head parameters are optimized during warm-up
warmup_optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3
)

print(f"Warm-up trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")

# Train for 5-10 warm-up epochs here (training loop omitted for brevity)
# ...after warm-up epochs...


# ─── Stage 2: Unfreeze all layers and fine-tune with lower LR ───────────────
for param in model.parameters():
    param.requires_grad = True

# Lower learning rate prevents large gradients from destroying pre-trained features
# Layer-wise learning rate decay: backbone gets 10x smaller LR than the head
backbone_params = [p for n, p in model.named_parameters() if 'fc' not in n]
head_params     = [p for n, p in model.named_parameters() if 'fc' in n]

finetune_optimizer = torch.optim.Adam([
    {'params': backbone_params, 'lr': 1e-5},   # backbone: small LR, preserve features
    {'params': head_params,     'lr': 1e-4},   # head: larger LR, learn class boundaries
])

# Cosine annealing reduces LR smoothly over fine-tuning epochs
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    finetune_optimizer, T_max=20, eta_min=1e-7
)

print(f"Fine-tuning trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

# io.thecodeforge: Domain-aware augmentation pipelines for 2026 production use
from torchvision import transforms
import torchvision.transforms.v2 as transforms_v2  # v2 API: faster, supports bounding boxes

NORM_MEAN = [0.485, 0.456, 0.406]
NORM_STD  = [0.229, 0.224, 0.225]


# ─── Natural images (objects, products, vehicles) ────────────────────────────
train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.7, 1.0), ratio=(0.75, 1.33)),
    transforms_v2.RandomHorizontalFlip(p=0.5),
    transforms_v2.RandomRotation(degrees=15),
    transforms_v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05),
    transforms_v2.RandomGrayscale(p=0.05),          # occasional grayscale improves robustness
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
    transforms_v2.RandomErasing(p=0.1, scale=(0.02, 0.15)),  # cutout regularization
])


# ─── Validation and inference — always deterministic ─────────────────────────
val_transforms = transforms_v2.Compose([
    transforms_v2.Resize(256),
    transforms_v2.CenterCrop(224),
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])
# Important: val_transforms == inference_transforms
# If they differ, your benchmark does not represent production behavior


# ─── Medical imaging — laterality-aware, no horizontal flip ──────────────────
medical_train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.85, 1.0)),
    transforms_v2.RandomRotation(degrees=10),       # small rotation only
    transforms_v2.ColorJitter(brightness=0.15, contrast=0.15),  # conservative color shift
    # NO RandomHorizontalFlip — left/right matters for anatomical laterality
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])


# ─── Satellite / aerial imagery — rotation-invariant by nature ───────────────
aerial_train_transforms = transforms_v2.Compose([
    transforms_v2.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms_v2.RandomHorizontalFlip(p=0.5),
    transforms_v2.RandomVerticalFlip(p=0.5),        # valid for top-down imagery
    transforms_v2.RandomRotation(degrees=90),       # full rotation invariance for aerial views
    transforms_v2.ColorJitter(brightness=0.2, contrast=0.2),
    transforms_v2.ToImage(),
    transforms_v2.ToDtype(torch.float32, scale=True),
    transforms_v2.Normalize(mean=NORM_MEAN, std=NORM_STD),
])

// io.thecodeforge
import torch.nn as nn

class BaselineCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        return self.classifier(self.features(x))

// io.thecodeforge
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-4)
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3)

best_val_loss = float('inf')
epochs_no_improve = 0
early_stop_patience = 7

for epoch in range(50):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    val_loss = validate(model, test_loader, criterion)
    scheduler.step(val_loss)
    
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_cnn.pth')
        epochs_no_improve = 0
    else:
        epochs_no_improve += 1
        if epochs_no_improve >= early_stop_patience:
            print(f'Early stop at epoch {epoch}')
            break

// io.thecodeforge
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

model.eval()
all_preds = []
all_labels = []

with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
        all_preds.extend(preds.numpy())
        all_labels.extend(labels.numpy())

cm = confusion_matrix(all_labels, all_preds, normalize='true')
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=test_dataset.classes)
disp.plot(cmap='Blues', values_format='.2f')
plt.title('Normalized Confusion Matrix')
plt.savefig('confusion_matrix.png')

# Top-5 accuracy
top5_preds = torch.topk(outputs, 5, dim=1).indices
correct_top5 = top5_preds.eq(labels.view(-1, 1)).sum().item()
print(f'Top-5 accuracy: {correct_top5 / len(labels):.3f}')