From: Francois Fleuret Date: Tue, 11 May 2021 18:25:15 +0000 (+0200) Subject: Initial commit. X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=d74d7be5abef26c78d014bd179f2c52f81aca65b;p=pytorch.git Initial commit. --- diff --git a/tinymnist.py b/tinymnist.py new file mode 100755 index 0000000..8642b22 --- /dev/null +++ b/tinymnist.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python + +# Any copyright is dedicated to the Public Domain. +# https://creativecommons.org/publicdomain/zero/1.0/ + +# Written by Francois Fleuret + +import time, os +import torch, torchvision +from torch import nn +from torch.nn import functional as F + +lr, nb_epochs, batch_size = 1e-1, 10, 100 + +data_dir = os.environ.get('PYTORCH_DATA_DIR') or './data/' + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +###################################################################### + +train_set = torchvision.datasets.MNIST(root = data_dir, train = True, download = True) +train_input = train_set.data.view(-1, 1, 28, 28).float() +train_targets = train_set.targets + +test_set = torchvision.datasets.MNIST(root = data_dir, train = False, download = True) +test_input = test_set.data.view(-1, 1, 28, 28).float() +test_targets = test_set.targets + +###################################################################### + +class SomeLeNet(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size = 5) + self.conv2 = nn.Conv2d(32, 64, kernel_size = 5) + self.fc1 = nn.Linear(256, 200) + self.fc2 = nn.Linear(200, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), kernel_size = 3)) + x = F.relu(F.max_pool2d(self.conv2(x), kernel_size = 2)) + x = x.view(x.size(0), -1) + x = F.relu(self.fc1(x)) + x = self.fc2(x) + return x + +###################################################################### + +model = SomeLeNet() + +nb_parameters = sum(p.numel() for p in model.parameters()) + +print(f'nb_parameters {nb_parameters}') + +optimizer = torch.optim.SGD(model.parameters(), lr = lr) +criterion = nn.CrossEntropyLoss() + +model.to(device) +criterion.to(device) + +train_input, train_targets = train_input.to(device), train_targets.to(device) +test_input, test_targets = test_input.to(device), test_targets.to(device) + +mu, std = train_input.mean(), train_input.std() +train_input.sub_(mu).div_(std) +test_input.sub_(mu).div_(std) + +start_time = time.perf_counter() + +for k in range(nb_epochs): + acc_loss = 0. + + for input, targets in zip(train_input.split(batch_size), + train_targets.split(batch_size)): + output = model(input) + loss = criterion(output, targets) + acc_loss += loss.item() + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + nb_test_errors = 0 + for input, targets in zip(test_input.split(batch_size), + test_targets.split(batch_size)): + wta = model(input).argmax(1) + nb_test_errors += (wta != targets).long().sum() + test_error = nb_test_errors / test_input.size(0) + duration = time.perf_counter() - start_time + + print(f'loss {k} {duration:.02f}s {acc_loss:.02f} {test_error*100:.02f}%') + +######################################################################