baseline_cnn_cifar10.py
CIFAR-10 CNN baseline - 289 lines.
View on GitHub (baseline_cnn_cifar10.py).
Source
Section titled “Source”"""
Non-spiking CNN baseline for CIFAR-10.
Fair comparison with cifar10_conv_stdp_snn.py:
- Same grayscale input (1 channel, 32x32)
- Same data augmentation
- Same optimizer (AdamW) and scheduler (OneCycleLR)
- Batch size 4096 (optimised for 40GB A100), same epochs (250)
- Same conv front-end architecture (Conv2d 1->64->128->256)
- Replaces STDP/LIF layers with standard FC + ReLU
Run: uv run python baseline_cnn_cifar10.py
"""
import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
class BaselineCNN(nn.Module):
"""CNN with the same conv front-end as ConvSTDPNetwork but standard FC classifier."""
def __init__(self, hidden_size=2048, num_outputs=10, dropout=0.2, conv_dropout=0.1):
super().__init__()
# Same conv front-end as the SNN
self.conv = nn.Sequential(
nn.Conv2d(1, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2), # 16x16
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(2), # 8x8
nn.Dropout2d(conv_dropout),
)
conv_out_size = 256 * 8 * 8 # 16384
# Standard FC classifier (replaces STDP/LIF layers)
hidden2 = hidden_size // 2
hidden3 = hidden2 // 2
self.classifier = nn.Sequential(
nn.Linear(conv_out_size, hidden_size),
nn.LayerNorm(hidden_size),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_size, hidden2),
nn.LayerNorm(hidden2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden2, hidden3),
nn.LayerNorm(hidden3),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden3, num_outputs),
)
def forward(self, x):
# x: [batch, 1, 32, 32]
features = self.conv(x).flatten(1)
return self.classifier(features)
def setup_data(batch_size=512, workers=4):
# Same transforms as the SNN script
train_transform = transforms.Compose(
[
transforms.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)),
transforms.RandomPerspective(distortion_scale=0.2, p=0.5),
transforms.ToTensor(),
transforms.Grayscale(num_output_channels=1),
transforms.RandomErasing(p=0.1),
transforms.Normalize((0.5,), (0.5,)),
]
)
test_transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Grayscale(num_output_channels=1),
transforms.Normalize((0.5,), (0.5,)),
]
)
train_dataset = datasets.CIFAR10(root="./assets", train=True, download=True, transform=train_transform)
test_dataset = datasets.CIFAR10(root="./assets", train=False, download=True, transform=test_transform)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True,
persistent_workers=True,
prefetch_factor=4,
)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=batch_size * 2,
shuffle=False,
num_workers=workers,
pin_memory=True,
persistent_workers=True,
prefetch_factor=4,
)
return train_loader, test_loader, train_dataset.classes
def evaluate(net, test_loader, device):
net.eval()
correct = 0
total = 0
per_class_correct = np.zeros(10)
per_class_total = np.zeros(10)
with torch.no_grad():
for data, targets in test_loader:
data, targets = data.to(device), targets.to(device)
logits = net(data)
predicted = logits.argmax(dim=1)
correct += (predicted == targets).sum().item()
total += targets.size(0)
for cls in range(10):
mask = targets == cls
per_class_correct[cls] += (predicted[mask] == targets[mask]).sum().item()
per_class_total[cls] += mask.sum().item()
accuracy = correct / total
per_class_acc = per_class_correct / (per_class_total + 1e-8)
return accuracy, per_class_acc
def train(net, train_loader, test_loader, classes, num_epochs=250, device="cuda"):
# Scale LR with batch size (linear scaling rule: base_lr=2e-3 at bs=512)
base_lr = 2e-3
lr = base_lr * (train_loader.batch_size / 512)
optimizer = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=0.01, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=lr,
epochs=num_epochs,
steps_per_epoch=len(train_loader),
pct_start=0.3,
anneal_strategy="cos",
)
scaler = torch.amp.GradScaler("cuda")
best_acc = 0
results = []
for epoch in range(num_epochs):
net.train()
running_loss = 0.0
t0 = time.time()
for data, targets in train_loader:
data, targets = data.to(device, non_blocking=True), targets.to(device, non_blocking=True)
optimizer.zero_grad()
with torch.amp.autocast("cuda"):
logits = net(data)
loss = F.cross_entropy(logits, targets, label_smoothing=0.1)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
scheduler.step()
running_loss += loss.item()
avg_loss = running_loss / len(train_loader)
test_acc, per_class_acc = evaluate(net, test_loader, device)
elapsed = time.time() - t0
results.append({"epoch": epoch + 1, "loss": avg_loss, "accuracy": test_acc})
if test_acc > best_acc:
best_acc = test_acc
torch.save(net.state_dict(), os.path.join(OUTPUT_DIR, "baseline_cnn_cifar10_best.pth"))
if (epoch + 1) % 10 == 0 or epoch == 0:
print(
f"Epoch {epoch + 1:3d}/{num_epochs} | "
f"Loss: {avg_loss:.4f} | "
f"Test Acc: {test_acc * 100:.2f}% | "
f"Best: {best_acc * 100:.2f}% | "
f"Time: {elapsed:.1f}s"
)
# Final per-class report
test_acc, per_class_acc = evaluate(net, test_loader, device)
print("\n" + "=" * 60)
print(f"FINAL TEST ACCURACY: {test_acc * 100:.2f}%")
print(f"BEST TEST ACCURACY: {best_acc * 100:.2f}%")
print("=" * 60)
print("\nPer-class accuracy:")
for i, cls in enumerate(classes):
print(f" {cls:>12s}: {per_class_acc[i] * 100:.1f}%")
# Save results
np.savetxt(
os.path.join(OUTPUT_DIR, "baseline_cnn_cifar10_accuracies.csv"),
[r["accuracy"] for r in results],
delimiter=",",
)
np.savetxt(
os.path.join(OUTPUT_DIR, "baseline_cnn_cifar10_losses.csv"),
[r["loss"] for r in results],
delimiter=",",
)
return results
def main():
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
np.random.seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.deterministic = False
torch.set_float32_matmul_precision("high")
if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_flash_sdp(True)
# Optimised for 40GB A100
hidden_size = 2048
num_outputs = 10
batch_size = 4096
num_epochs = 250
workers = min(16, os.cpu_count() - 2)
print("=" * 60)
print("BASELINE CNN FOR CIFAR-10 (non-spiking)")
print("=" * 60)
print(f"Device: {device.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"VRAM Available: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
print(f"Architecture: Conv2d(1->64->128->256) + FC({hidden_size}->{hidden_size // 2}->{hidden_size // 4}->{num_outputs})")
print(f"Input: Grayscale CIFAR-10 (1x32x32)")
print(f"Batch size: {batch_size}")
print(f"Epochs: {num_epochs}")
print(f"Mixed precision: Enabled (FP16 + TF32)")
print(f"torch.compile: Enabled")
print("=" * 60)
net = BaselineCNN(hidden_size=hidden_size, num_outputs=num_outputs).to(device)
net = torch.compile(net)
total_params = sum(p.numel() for p in net.parameters())
trainable_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")
train_loader, test_loader, classes = setup_data(batch_size=batch_size, workers=workers)
print(f"\nTrain samples: {len(train_loader.dataset):,}")
print(f"Test samples: {len(test_loader.dataset):,}")
print()
train(net, train_loader, test_loader, classes, num_epochs=num_epochs, device=device)
if __name__ == "__main__":
main()