lr, nb_epochs, batch_size = 1e-1, 10, 100
-data_dir = os.environ.get('PYTORCH_DATA_DIR') or './data/'
+data_dir = os.environ.get("PYTORCH_DATA_DIR") or "./data/"
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
######################################################################
-train_set = torchvision.datasets.MNIST(root = data_dir, train = True, download = True)
+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_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.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 = 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}')
+print(f"nb_parameters {nb_parameters}")
-optimizer = torch.optim.SGD(model.parameters(), lr = lr)
+optimizer = torch.optim.SGD(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
model.to(device)
start_time = time.perf_counter()
for k in range(nb_epochs):
- acc_loss = 0.
+ acc_loss = 0.0
- for input, targets in zip(train_input.split(batch_size),
- train_targets.split(batch_size)):
+ 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.step()
nb_test_errors = 0
- for input, targets in zip(test_input.split(batch_size),
- test_targets.split(batch_size)):
+ 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}%')
+ print(f"loss {k} {duration:.02f}s {acc_loss:.02f} {test_error*100:.02f}%")
######################################################################