From 8f4f24523235965cfa2041ed00cc40fc0b4bd367 Mon Sep 17 00:00:00 2001 From: ericmarin Date: Fri, 10 Apr 2026 15:06:56 +0200 Subject: added MNIST, changed cache and parser --- examples/mnist/mnist.py | 45 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 examples/mnist/mnist.py (limited to 'examples/mnist/mnist.py') diff --git a/examples/mnist/mnist.py b/examples/mnist/mnist.py new file mode 100644 index 0000000..0a81878 --- /dev/null +++ b/examples/mnist/mnist.py @@ -0,0 +1,45 @@ +import torch, torch.nn as nn +from torchvision.datasets import MNIST +from torch.utils.data import DataLoader +from torchvision import transforms + +class MNIST_MLP(nn.Module): + def __init__(self, hidden_dim): + super().__init__() + self.layers = nn.Sequential( + nn.Flatten(), + nn.Linear(28 * 28, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, 10), + ) + def forward(self, x): + return self.layers(x) + +train_dataset = MNIST('./', download=True, transform=transforms.ToTensor(), train=True) +trainloader = DataLoader(train_dataset, batch_size=128, shuffle=True) + +def train_model(name: str, dim): + net = MNIST_MLP(hidden_dim=dim) + loss_fn = nn.CrossEntropyLoss() + optimizer = torch.optim.Adam(net.parameters(), lr=0.5e-4) + + print(f"Training {name} ({dim} neurons)...") + for epoch in range(100): + global loss + for data in trainloader: + inputs, targets = data + optimizer.zero_grad() + outputs = net(inputs) + loss = loss_fn(outputs, targets) + loss.backward() + optimizer.step() + if (epoch + 1) % 10 == 0: + print(f" Epoch {epoch+1}, Loss: {loss.item():.4f}") + return net + +if __name__ == "__main__": + torch_net_a = train_model("Network A", 6).eval() + torch_net_b = train_model("Network B", 12).eval() + + torch.onnx.export(torch_net_a, (torch.randn(1, 28, 28),), "mnist_a.onnx") + torch.onnx.export(torch_net_b, (torch.randn(1, 28, 28),), "mnist_b.onnx") -- cgit v1.2.3