# pylint: disable=E1101,R,C import numpy as np import torch.nn as nn import torch.nn.functional as F import torch import torch.utils.data as data_utils import gzip import pickle import numpy as np from torch.autograd import Variable MNIST_PATH = "s2_mnist.gz" DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") NUM_EPOCHS = 20 BATCH_SIZE = 32 LEARNING_RATE = 5e-4 class ConvNet(nn.Module): def __init__(self): super().__init__() f1 = 32 f2 = 64 self.feature_layer = nn.Sequential( torch.nn.Conv2d(1, f1, kernel_size=5, stride=3), torch.nn.ReLU(), torch.nn.Conv2d(f1, f2, kernel_size=5, stride=3), torch.nn.ReLU() ) self.out_layer = torch.nn.Linear(f2 * 5**2, 10) def forward(self, x): x = self.feature_layer(x) x = x.view(x.shape[0], -1) x = self.out_layer(x) return x def load_data(path, batch_size): with gzip.open(path, 'rb') as f: dataset = pickle.load(f) train_data = torch.from_numpy( dataset["train"]["images"][:, None, :, :].astype(np.float32)) train_labels = torch.from_numpy( dataset["train"]["labels"].astype(np.int64)) # TODO normalize dataset # mean = train_data.mean() # stdv = train_data.std() train_dataset = data_utils.TensorDataset(train_data, train_labels) train_loader = data_utils.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_data = torch.from_numpy( dataset["test"]["images"][:, None, :, :].astype(np.float32)) test_labels = torch.from_numpy( dataset["test"]["labels"].astype(np.int64)) test_dataset = data_utils.TensorDataset(test_data, test_labels) test_loader = data_utils.DataLoader(test_dataset, batch_size=batch_size, shuffle=True) return train_loader, test_loader, train_dataset, test_dataset def main(): train_loader, test_loader, train_dataset, _ = load_data( MNIST_PATH, BATCH_SIZE) classifier = ConvNet() classifier.to(DEVICE) print("#params", sum([x.numel() for x in classifier.parameters()])) criterion = nn.CrossEntropyLoss() criterion = criterion.to(DEVICE) optimizer = torch.optim.Adam( classifier.parameters(), lr=LEARNING_RATE) for epoch in range(NUM_EPOCHS): for i, (images, labels) in enumerate(train_loader): classifier.train() images = images.to(DEVICE) labels = labels.to(DEVICE) optimizer.zero_grad() outputs = classifier(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() print('\rEpoch [{0}/{1}], Iter [{2}/{3}] Loss: {4:.4f}'.format( epoch+1, NUM_EPOCHS, i+1, len(train_dataset)//BATCH_SIZE, loss.item()), end="") print("") correct = 0 total = 0 for i, (images, labels) in enumerate(test_loader): classifier.eval() with torch.no_grad(): images = images.to(DEVICE) labels = labels.to(DEVICE) outputs = classifier(images) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).long().sum().item() print('Test Accuracy: {0}'.format(100 * correct / total)) if __name__ == '__main__': main()