105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
'''
|
|
put the last few chapters together to classify digits from MNIST dataset
|
|
-> MNIST
|
|
-> DataLoader, Transformation
|
|
-> Loss + Optimizer
|
|
-> Training Loop with batch training
|
|
-> Model Evaluation
|
|
-> GPU Support
|
|
'''
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision
|
|
import torchvision.transforms as transforms
|
|
import matplotlib.pyplot as plt
|
|
|
|
# device config
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
print(f'Device is {device}')
|
|
|
|
# hyper parameters
|
|
input_size = 784 # 28x28 pixel images
|
|
hidden_size = 100 # PLAY WITH THIS
|
|
num_classes = 10 # digits 0..9
|
|
num_epochs = 2 # PLAY WITH THIS
|
|
batch_size = 100
|
|
learning_rate = .001
|
|
|
|
# MNIST
|
|
train_dataset = torchvision.datasets.MNIST(root='./data', train=True,
|
|
download=True, transform=transforms.ToTensor())
|
|
|
|
test_dataset = torchvision.datasets.MNIST(root='./data', train=False,
|
|
transform=transforms.ToTensor())
|
|
|
|
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)
|
|
|
|
examples = iter(train_loader)
|
|
samples, labels = examples.next()
|
|
print(samples.shape, labels.shape)
|
|
|
|
for i in range(6):
|
|
plt.subplot(2, 3, i+1)
|
|
plt.imshow(samples[i][0], cmap='gray')
|
|
# plt.show()
|
|
|
|
class NeuralNet(nn.Module):
|
|
def __init__(self, input_size, hidden_size, num_classes):
|
|
super(NeuralNet, self).__init__()
|
|
self.layers = nn.Sequential(
|
|
nn.Linear(input_size, hidden_size),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_size, num_classes)
|
|
# no softmax because its included in the CE loss function
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.layers(x)
|
|
|
|
model = NeuralNet(input_size, hidden_size, num_classes).to(device)
|
|
print(model)
|
|
|
|
# loss and optimizer
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
|
|
|
# training loop
|
|
num_total_steps = len(train_loader)
|
|
for epoch in range(num_epochs):
|
|
for batch, (images, labels) in enumerate(train_loader):
|
|
# reshape 100, 1, 28, 28 -> 100, 784
|
|
images = images.reshape(-1, 28*28).to(device) # reshape and send to gpu if available
|
|
labels = labels.to(device)
|
|
|
|
# forward
|
|
outputs = model(images)
|
|
loss = criterion(outputs, labels)
|
|
|
|
# backward + update
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if (batch+1) % 100 == 0:
|
|
print(f'Epoch {epoch+1}/{num_epochs}, step {batch+1}/{num_total_steps}, loss = {loss.item():.4f}')
|
|
|
|
# test
|
|
with torch.no_grad():
|
|
n_correct = 0
|
|
n_samples = 0
|
|
for images, labels in test_loader:
|
|
images = images.reshape(-1, 28*28).to(device)
|
|
labels = labels.to(device)
|
|
outputs = model(images)
|
|
|
|
# value, index (index is class label)
|
|
_, predictions = torch.max(outputs, 1)
|
|
n_samples += labels.shape[0]
|
|
n_correct += (predictions==labels).sum().item()
|
|
|
|
acc = 100.*n_correct/n_samples
|
|
print(f'Accuracy = {acc}%') |