import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
# Training settings
batch_size = 64
# MNIST Dataset
train_dataset = datasets.MNIST(root='./data/',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = datasets.MNIST(root='./data/',
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)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.mp = nn.MaxPool2d(2)
self.fc = nn.Linear(320, 10)
def forward(self, x):
in_size = x.size(0)
x = F.relu(self.mp(self.conv1(x)))
x = F.relu(self.mp(self.conv2(x)))
x = x.view(in_size, -1) # flatten the tensor
x = self.fc(x)
return F.log_softmax(x,dim=1)
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
criterion = nn.CrossEntropyLoss()
def train(epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data, target
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % 1000 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def test():
model.eval()
test_loss = 0
correct = 0
for data, target in test_loader:
data, target = data, target
output = model(data)
# sum up batch loss
test_loss += criterion(output, target).item()
# get the index of the max log-probability
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(target.data.view_as(pred)).cpu().sum()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
for epoch in range(1, 10):
train(epoch)
test()
Train Epoch: 1 [0/60000 (0%)] Loss: 2.308939 Test set: Average loss: 0.0028, Accuracy: 9469/10000 (94%) Train Epoch: 2 [0/60000 (0%)] Loss: 0.201025 Test set: Average loss: 0.0017, Accuracy: 9676/10000 (96%) Train Epoch: 3 [0/60000 (0%)] Loss: 0.118083 Test set: Average loss: 0.0013, Accuracy: 9774/10000 (97%) Train Epoch: 4 [0/60000 (0%)] Loss: 0.044904 Test set: Average loss: 0.0011, Accuracy: 9786/10000 (97%) Train Epoch: 5 [0/60000 (0%)] Loss: 0.029464 Test set: Average loss: 0.0011, Accuracy: 9801/10000 (98%) Train Epoch: 6 [0/60000 (0%)] Loss: 0.015576 Test set: Average loss: 0.0010, Accuracy: 9802/10000 (98%) Train Epoch: 7 [0/60000 (0%)] Loss: 0.034443 Test set: Average loss: 0.0008, Accuracy: 9825/10000 (98%) Train Epoch: 8 [0/60000 (0%)] Loss: 0.006708 Test set: Average loss: 0.0008, Accuracy: 9832/10000 (98%) Train Epoch: 9 [0/60000 (0%)] Loss: 0.211187 Test set: Average loss: 0.0008, Accuracy: 9821/10000 (98%)