import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.autograd import Variable from torchvision import datasets, transforms # Training settings batch_size = 64 # MNIST Dataset dataset_path = "../data/mnist" train_dataset = datasets.MNIST(root=dataset_path, train=True, transform=transforms.ToTensor(), download=True) test_dataset = datasets.MNIST(root=dataset_path, train=False, transform=transforms.ToTensor()) # Data Loader (Input Pipeline) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # define Network seq_net = nn.Sequential( nn.Linear(28*28, 300), nn.ReLU(), nn.Linear(300, 100), nn.ReLU(), nn.Linear(100, 10) ) # define optimizer & criterion param = seq_net.parameters() optim = torch.optim.Adam(param, 0.01) criterion = nn.CrossEntropyLoss() # train the network for e in range(100): for batch_idx, (data, target) in enumerate(train_loader): data, target = Variable(data), Variable(target) data = data.view(-1, 784) out = seq_net(data) loss = criterion(out, target) optim.zero_grad() loss.backward() optim.step() if batch_idx % 100 == 0: pred = out.data.max(1, keepdim=True)[1] c = float(pred.eq(target.data.view_as(pred)).cpu().sum())/out.size(0) print("epoch: %5d, loss: %f, acc: %f" % (e+1, loss.data[0], c))