emnist training script

frontend/train_emnist.py

@@ -0,0 +1,133 @@
+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
+
+# EMNIST Letters: labels 1-26 (A-Z), but torchvision returns 1-26
+# We remap to 0-25 for a clean 26-class output
+
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+
+print("Downloading EMNIST Letters dataset...")
+train_set = torchvision.datasets.EMNIST(
+    root='./emnist_data', split='letters', train=True,
+    download=True, transform=transform
+)
+test_set = torchvision.datasets.EMNIST(
+    root='./emnist_data', split='letters', train=False,
+    download=True, transform=transform
+)
+
+train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=0)
+test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=0)
+
+class LetterCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 32, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Dropout2d(0.25),
+
+            nn.Conv2d(32, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Dropout2d(0.25),
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(64 * 7 * 7, 128),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(128, 26),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        return x
+
+model = LetterCNN()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Training on {device}")
+model.to(device)
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
+
+# EMNIST letters labels are 1-26, remap to 0-25
+for epoch in range(10):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    for images, labels in train_loader:
+        labels = labels - 1  # remap 1-26 to 0-25
+        images, labels = images.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+
+    scheduler.step()
+    acc = 100.0 * correct / total
+    print(f"Epoch {epoch+1}/10 - loss: {running_loss/len(train_loader):.4f}, train acc: {acc:.2f}%")
+
+# Test accuracy
+model.eval()
+correct = 0
+total = 0
+with torch.no_grad():
+    for images, labels in test_loader:
+        labels = labels - 1
+        images, labels = images.to(device), labels.to(device)
+        outputs = model(images)
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+
+print(f"Test accuracy: {100.0 * correct / total:.2f}%")
+
+# Export to ONNX
+model.eval()
+model.to('cpu')
+dummy = torch.randn(1, 1, 28, 28)
+onnx_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "public", "models", "emnist-letters.onnx")
+
+torch.onnx.export(
+    model, dummy, onnx_path,
+    input_names=['input'],
+    output_names=['output'],
+    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
+    opset_version=13
+)
+
+import os
+size_kb = os.path.getsize(onnx_path) / 1024
+print(f"Saved ONNX model to {onnx_path} ({size_kb:.1f} KB)")
+
+# Verify with onnx
+import onnx
+m = onnx.load(onnx_path)
+onnx.checker.check_model(m)
+print("ONNX model verified OK")
+print(f"Input: {m.graph.input[0].type.tensor_type.shape}")
+print(f"Output: {m.graph.output[0].type.tensor_type.shape}")