In [1]:
import torch
import numpy as np
In [2]:
import torch.nn.functional as F
from torch import nn, optim
In [3]:
from torchvision import datasets
import torchvision.transforms as transforms
In [4]:
num_workers = 2
batch_size = 20
In [6]:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_data = datasets.CIFAR10(root="CIFAR10_data", train=True, download=True, transform=transform)
test_data = datasets.CIFAR10(root="CIFAR10_data", train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers=num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)
In [7]:
import matplotlib.pyplot as plt
%matplotlib inline
# helper function to un-normalize and display an image
def imshow(img):
img = img / 2 + 0.5 # unnormalize
plt.imshow(np.transpose(img, (1, 2, 0))) # convert from Tensor image
In [8]:
classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck']
In [9]:
import matplotlib.pyplot as plt
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy()
In [13]:
fig = plt.figure(figsize=(25, 4))
for idx in np.arange(20):
ax = fig.add_subplot(2, 20/2, idx + 1, xticks=[], yticks=[])
imshow(images[idx])
ax.set_title(str(classes[labels[idx].item()]))
In [14]:
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(64 * 4 * 4, 512)
self.fc2 = nn.Linear(512, 10)
self.dropout = nn.Dropout(0.2)
self.log_softmax = nn.LogSoftmax(dim=1)
def forward(self, x):
# flatten
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
# flatten
x = x.view(-1, 64 * 4 * 4)
x = self.dropout(x)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.log_softmax(self.fc2(x))
return x
In [15]:
model = Net()
print(model)
In [16]:
criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
In [17]:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
In [26]:
n_epochs = 30
model.to(device)
for epoch in range(n_epochs):
train_loss = 0.0
for data, target in train_loader:
optimizer.zero_grad()
data, target = data.to(device), target.to(device)
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
train_loss += loss.item() * data.size(0)
train_loss = train_loss / len(train_loader.dataset)
print("Epoch {}: Training loss {:.6f}".format(epoch, train_loss))
In [27]:
test_loss = 0.0
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
#!Important: Change model to evaluation mode to deactivate dropout
model.eval()
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
loss = criterion(output, target)
test_loss += loss.item() * data.size(0)
output = torch.exp(output)
_, pred = torch.max(output, 1)
correct = np.squeeze(pred.eq(target.data.view_as(pred)))
for i in range(batch_size):
label = target.data[i]
class_correct[label] += correct[i].item()
class_total[label] += 1
#!Important: Change model to training mode to activate dropout
model.train()
In [28]:
test_loss = test_loss / len(test_loader.dataset)
print("Test loss: {:.3f}\n".format(test_loss))
for i in range(10):
if class_total[i] > 0:
print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
str(i), 100 * class_correct[i] / class_total[i],
np.sum(class_correct[i]), np.sum(class_total[i])))
else:
print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))
print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (
100. * np.sum(class_correct) / np.sum(class_total),
np.sum(class_correct), np.sum(class_total)))
In [29]:
# obtain one batch of test images
dataiter = iter(test_loader)
images, labels = dataiter.next()
images, labels = images.to(device), labels.to(device)
# get sample outputs
output = model(images)
# convert output probabilities to predicted class
_, preds = torch.max(output, 1)
# prep images for display
images = images.cpu().numpy()
# plot the images in the batch, along with predicted and true labels
fig = plt.figure(figsize=(25, 4))
for idx in np.arange(20):
ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
imshow(images[idx])
ax.set_title("{} ({})".format(str(preds[idx].item()), str(labels[idx].item())),
color=("green" if preds[idx]==labels[idx] else "red"))
In [ ]: