Image Classifier
This section provides an example of model training using a PyTorch convolutional neural net. We will be using the model described in the PyTorch documentation here for model training.
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import markov
# Get and load the dataset (CIFAR10 dataset)
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 4
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# Define the model
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = torch.flatten(x, 1) # flatten all dimensions except batch
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# Create the ExperimentRecorder
recorder = markov.ExperimentRecorder(
name="Image classifier using CNN for CIFAR10",
note="We are using a convolutional neural net to "
"train a classifier with the CIFAR10 dataset. "
"CIFAR10 has the classes: ‘airplane’, ‘automobile’, "
"‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, "
"‘ship’, ‘truck’. The images in CIFAR-10 are of "
"size 3x32x32, i.e. 3-channel color images of 32x32 pixels in size.",
hyper_parameters={
"lr": 0.001,
"momentum": 0.9,
"batch_size": 4,
"optimizer": "SGD"
}
)
recorder.register()
# Start model training
with recorder:
avg_running_loss = 10000 # INF
for epoch in range(10): # Loop over the dataset multiple times
running_loss = 0.0
data_points = 0
for i, data in enumerate(trainloader, 0):
# Get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# Zero the parameter gradients
optimizer.zero_grad()
# Forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
data_points += 1
avg_running_loss = running_loss / data_points
# Record running loss
recorder.add_record({"average running loss": avg_running_loss})
recorder.summary.add_training_loss(value=str(avg_running_loss))
Last updated
Was this helpful?