Upload train_digit_classifier.py

train_digit_classifier.py

@@ -0,0 +1,300 @@
+"""
+train_digit_classifier.py
+
+A fully documented training script for a convolutional neural network (CNN)
+classifier trained on MNIST + EMNIST digits + blank images.
+
+Author: Deep Shah
+License: GPL-3.0
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, Dataset, TensorDataset
+from sklearn.model_selection import train_test_split
+import os
+
+# ----------------------------------------------------------------------
+# 1. Reproducibility Setup
+# ----------------------------------------------------------------------
+
+# Set fixed seeds to make results deterministic (important for debugging and reproducibility)
+torch.manual_seed(42)
+np.random.seed(42)
+
+# ----------------------------------------------------------------------
+# 2. Device Selection
+# ----------------------------------------------------------------------
+
+# Automatically use GPU if available; fallback to CPU otherwise
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"[INFO] Using device: {device}")
+
+# ----------------------------------------------------------------------
+# 3. EMNIST Loader (Custom Dataset class)
+# ----------------------------------------------------------------------
+
+class EMNISTDigitsDataset(Dataset):
+    """
+    A PyTorch-compatible wrapper for the EMNIST digits dataset loaded via TensorFlow Datasets.
+    Ensures data is shaped correctly and optionally transformed.
+    """
+
+    def __init__(self, split="train", transform=None):
+        import tensorflow_datasets as tfds
+        ds = tfds.load("emnist/digits", split=split, as_supervised=True)
+        self.images = []
+        self.labels = []
+        for image, label in tfds.as_numpy(ds):
+            if image.ndim == 2:
+                image = image[..., np.newaxis]
+            elif image.ndim == 4 and image.shape[0] == 1:
+                image = image[0]
+            self.images.append(image)
+            self.labels.append(label)
+        self.images = np.array(self.images, dtype=np.float32) / 255.0  # Normalize to [0,1]
+        self.labels = np.array(self.labels, dtype=np.int64)
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image = self.images[idx]
+        label = self.labels[idx]
+        if self.transform:
+            image = self.transform(torch.tensor(image.transpose(2, 0, 1))).transpose(1, 2).numpy()
+        return torch.tensor(image.transpose(2, 0, 1), dtype=torch.float32), torch.tensor(label, dtype=torch.long)
+
+# ----------------------------------------------------------------------
+# 4. Data Augmentation Strategy
+# ----------------------------------------------------------------------
+
+# We use a modest augmentation strategy to improve generalization
+train_transform = transforms.Compose([
+    transforms.ToPILImage(),
+    transforms.RandomRotation(10),                        # Handle slanted handwriting
+    transforms.RandomAffine(degrees=0, scale=(0.9, 1.1), translate=(0.1, 0.1)),  # Simulate slight distortions
+    transforms.ToTensor()
+])
+
+# ----------------------------------------------------------------------
+# 5. Load Datasets (MNIST + EMNIST + Blank)
+# ----------------------------------------------------------------------
+
+# Load MNIST
+mnist_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True)
+mnist_images = mnist_dataset.data.numpy().astype(np.float32) / 255.0
+mnist_images = mnist_images[..., np.newaxis]
+mnist_labels = mnist_dataset.targets.numpy()
+
+# Load EMNIST
+emnist_dataset = EMNISTDigitsDataset(split="train", transform=None)
+emnist_images = emnist_dataset.images
+emnist_labels = emnist_dataset.labels
+
+# Create blank (all-black) 28x28 images, labeled with class 10
+x_blank = np.zeros((5000, 28, 28, 1), dtype=np.float32)
+y_blank = np.full((5000,), 10, dtype=np.int64)
+
+# Combine all datasets
+x_combined = np.concatenate([mnist_images, emnist_images, x_blank], axis=0)
+y_combined = np.concatenate([mnist_labels, emnist_labels, y_blank], axis=0)
+
+# Shuffle for randomness
+indices = np.random.permutation(len(x_combined))
+x_combined = x_combined[indices]
+y_combined = y_combined[indices]
+
+# ----------------------------------------------------------------------
+# 6. Train/Validation Split
+# ----------------------------------------------------------------------
+
+x_train, x_val, y_train, y_val = train_test_split(
+    x_combined, y_combined, test_size=0.1, random_state=42
+)
+
+# Convert to PyTorch format
+train_dataset = TensorDataset(
+    torch.tensor(x_train.transpose(0, 3, 1, 2), dtype=torch.float32),
+    torch.tensor(y_train, dtype=torch.long)
+)
+val_dataset = TensorDataset(
+    torch.tensor(x_val.transpose(0, 3, 1, 2), dtype=torch.float32),
+    torch.tensor(y_val, dtype=torch.long)
+)
+
+train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+# ----------------------------------------------------------------------
+# 7. CNN Architecture
+# ----------------------------------------------------------------------
+
+class CNN(nn.Module):
+    """
+    This CNN is designed to:
+    - Use 3 convolutional blocks with increasing depth (32 -> 64 -> 128)
+    - Use BatchNorm to stabilize training
+    - Use Dropout to prevent overfitting
+    - Flatten and use 2 dense layers to classify
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1),    # Small receptive field
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, 3, padding=1),   # Slightly deeper
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Dropout(0.1),                  # Helps regularize
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Dropout(0.1)
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128 * 7 * 7, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 11)  # 0-9 digits + blank (class 10)
+        )
+
+    def forward(self, x):
+        return self.classifier(self.features(x))
+
+model = CNN().to(device)
+
+# ----------------------------------------------------------------------
+# 8. Training Configuration
+# ----------------------------------------------------------------------
+
+# CrossEntropyLoss is standard for multi-class classification
+criterion = nn.CrossEntropyLoss()
+
+# Adam is used because it's efficient for noisy gradients & fast convergence
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+# ReduceLROnPlateau reduces LR when validation loss plateaus (adaptive control)
+scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.2, patience=2, min_lr=1e-6)
+
+# Early stopping is used to prevent overfitting and wasted training
+patience = 5
+patience_counter = 0
+best_val_loss = float("inf")
+best_model_state = None
+
+# ----------------------------------------------------------------------
+# 9. Training Loop
+# ----------------------------------------------------------------------
+
+for epoch in range(1, 51):
+    model.train()
+    running_loss = 0
+    correct = 0
+    total = 0
+
+    for images, labels in train_loader:
+        images, labels = images.to(device), labels.to(device)
+
+        # Apply data augmentation on CPU
+        for i in range(len(images)):
+            images[i] = train_transform(images[i].cpu()).to(device)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+        _, predicted = torch.max(outputs, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+    train_acc = 100 * correct / total
+    train_loss = running_loss / len(train_loader)
+
+    # ----------------
+    # Validation phase
+    # ----------------
+    model.eval()
+    val_loss = 0
+    val_correct = 0
+    val_total = 0
+    with torch.no_grad():
+        for images, labels in val_loader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            val_loss += loss.item()
+            _, predicted = torch.max(outputs, 1)
+            val_total += labels.size(0)
+            val_correct += (predicted == labels).sum().item()
+
+    val_acc = 100 * val_correct / val_total
+    val_loss /= len(val_loader)
+
+    print(f"Epoch {epoch:02d}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%, "
+          f"Val Loss={val_loss:.4f}, Val Acc={val_acc:.2f}%")
+
+    # Adjust learning rate if plateau
+    scheduler.step(val_loss)
+
+    # Save best model
+    if val_loss < best_val_loss:
+        best_val_loss = val_loss
+        best_model_state = model.state_dict()
+        patience_counter = 0
+    else:
+        patience_counter += 1
+        if patience_counter >= patience:
+            print("[INFO] Early stopping triggered.")
+            break
+
+# Load best model
+model.load_state_dict(best_model_state)
+
+# Save PyTorch weights
+torch.save(model.state_dict(), "mnist_emnist_blank_cnn_v1.pth")
+print("[INFO] Model weights saved as mnist_emnist_blank_cnn_v1.pth")
+
+# Convert to TorchScript for deployment (required by Hugging Face Inference API)
+model.eval()
+example_input = torch.randn(1, 1, 28, 28).to(device)
+scripted_model = torch.jit.trace(model, example_input)
+scripted_model.save("mnist_emnist_blank_cnn_v1.pt")
+print("[INFO] TorchScript model saved as mnist_emnist_blank_cnn_v1.pt")
+
+# ONNX export 
+# We move to CPU just for export (then restore the device).
+prev_device = next(model.parameters()).device
+try:
+    model_cpu = model.to("cpu").eval()
+    dummy = torch.randn(1, 1, 28, 28)  # match input shape
+
+    onnx_path = "mnist_emnist_blank_cnn_v1.onnx"
+    torch.onnx.export(
+        model_cpu,
+        dummy,
+        onnx_path,
+        export_params=True,
+        opset_version=13,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["logits"],
+        dynamic_axes={"input": {0: "batch_size"}, "logits": {0: "batch_size"}},
+    )
+    print(f"[INFO] ONNX model saved as {onnx_path}")
+finally:
+    model.to(prev_device).eval()  # restore original device