Create lenet_mnist_torch.py

lenet_mnist_torch.py

@@ -0,0 +1,102 @@
+"""
+This file creates a simple lenet network using the MNIST dataset.
+"""
+
+import random
+
+import torch
+from torchvision import datasets, transforms
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Download the MNIST Dataset
+
+def get_mnist_dataset():
+    transform = transforms.ToTensor()
+    train_set = datasets.MNIST(root='./data', train=True,  transform=transform, download=True)
+    test_set  = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
+    return train_set, test_set
+
+# Create the lenet model
+
+class Classifier(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Conv2d(1, 32, 5),  # 28 -> 24
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),   # 24 -> 12
+            nn.Conv2d(32, 32, 5), # 12 ->  8
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),   #  8 ->  4
+            nn.Flatten(),
+            nn.Linear(32*4*4, 100),
+            nn.ReLU(),
+            nn.Linear(100, 100),
+            nn.ReLU(),
+            nn.Linear(100, 10)
+        )
+
+    def forward(self, x):
+        return self.network(x)
+
+# Compute accuracy function
+
+def compute_accuracy(model, data_set, nb_samples):
+    nb_valid = 0
+    for it in range(nb_samples):
+        # get a sample
+        sample_idx = torch.randint(len(data_set), size=(1,)).item()
+        img, label = data_set[sample_idx]
+        # compute the output
+        x = torch.reshape(img, (1,1,28,28))
+        y_h = model.forward(x)
+        pred_label = torch.argmax(y_h).item()
+        if label == pred_label :
+            nb_valid = nb_valid + 1
+    return nb_valid / nb_samples
+
+# Train the model
+
+def train_model(NB_ITERATION, CHECK_PERIOD, train_set, test_set, classifier):
+    accuracy_history = []
+    for it in range(NB_ITERATION):
+        sample_idx = random.randint(0, len(train_set)-1)
+        img, label = train_set[sample_idx]
+        x = torch.flatten(img)
+        x = torch.reshape(x, (1,1,28,28))
+        y = torch.zeros(1,10)
+        y[0][label] = 1
+        y_h = classifier.forward(x)
+        #print(y_h.shape, 'test')
+        l = F.mse_loss(y, y_h)
+        l.backward()
+        for p in classifier.parameters():
+            with torch.no_grad():
+                p -= 0.01 * p.grad
+            p.grad.zero_()
+
+        if it % CHECK_PERIOD == 0:
+            accuracy = compute_accuracy(classifier, test_set, CHECK_PERIOD)
+            accuracy_history.append(accuracy)
+            print(f'it {it}: accuracy = {accuracy:.8f} ')
+
+
+def create_lenet():
+    # Get Dataset
+    train_set, test_set = get_mnist_dataset()
+
+    # Create model
+    classifier = Classifier()
+
+    # Train model
+    NB_ITERATION = 50000
+    CHECK_PERIOD = 3000
+    print("NB_ITERATIONS = ", NB_ITERATION)
+    print("CHECK_PERIOD  = ", CHECK_PERIOD)
+    print("\nTraining LeNet...")
+    train_model(NB_ITERATION, CHECK_PERIOD, train_set, test_set, classifier)
+
+    # Export as ONNX
+    x = torch.Tensor(1,1,28,28)
+    torch.onnx.export(classifier.network, x, 'lenet.onnx', verbose=False, input_names=[ "input" ], output_names=[ "output" ])