Spaces:
Sleeping
Sleeping
File size: 4,380 Bytes
c88c8f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.optim as optim
from datetime import datetime
class OCRModel(nn.Module):
def __init__(self):
super().__init__()
self.stack = nn.Sequential(
# === FIRST BLOCK ====
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
# === NEXT BLOCK ====
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Dropout2d(0.25),
# === Final Block ===
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=1), # Padding=1 handles odd 7x7 dimensions cleanly
nn.Dropout2d(0.25),
# === Linear Flatten ===
nn.Flatten(),
nn.Linear(2048, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 128),
nn.BatchNorm1d(128),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(128, 47)
)
def forward(self, x):
return self.stack(x)
if __name__ == "__main__":
print(f"[{datetime.now()}][Info] Setting data transformation Formula... ")
# AI model training retrieval
train_transform = transforms.Compose([
transforms.RandomRotation(degrees=10, fill=0), # Small rotation for variations
transforms.RandomAffine(degrees=0, translate=(0.08, 0.08), fill=0), # Handles poor centering
transforms.ToTensor(),
transforms.Normalize((0.1751,), (0.3332,))
])
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1751,), (0.3332,))
])
print(f"[{datetime.now()}][Status] Downloading Dataset... Pulling from EMNIST dataset..")
train_dataset = torchvision.datasets.EMNIST(
root='./data',
split='balanced',
train=True,
download=True,
transform=train_transform
)
print(f"[{datetime.now()}][Status] Done, preparing data for training...")
train_loader = DataLoader(dataset=train_dataset, batch_size = 64, shuffle=True, num_workers=0, pin_memory=True)
# === Initialize Trainer ===
print(f"[{datetime.now()}][Status] Data retrieved. Initializing Model Trainer...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = OCRModel().to(device)
print(f"[{datetime.now()}][Status] Loaded Model")
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
print(f"[{datetime.now()}][Status] Criterion and Optimizer Formulas have been loaded. ")
epochs = 10
print(f"[{datetime.now()}][Status] Initializing Training Loop for {epochs} epochs. ")
for epoch in range(epochs):
model.train()
running_loss = 0
correct = 0
total = 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
epoch_loss = running_loss / total
epoch_acc = 100.0 * correct/total
print(f"[TRAINING] Epoch {epoch+1}/{epochs}, Loss is {epoch_loss:.4f}, Accuracy, {epoch_acc:.2f}%")
print(f"[{datetime.now()}][Status] Training Done")
print(f"[{datetime.now()}][Info] Saving Model to disk...")
torch.save(model.state_dict(), "OCR_Model.pt")
print(f"[{datetime.now()}][Status] Done! Exiting program...")
exit(0)
# TODO: Retrain model after adding new layer |