Update 01-convolutional-nn

Author: yakhyo

tutorials/02-intermediate/01-convolutional-nn/main.py

@@ -0,0 +1,204 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+import torchvision
+import torchvision.transforms as transforms
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+# ========================================= #
+#         Load and normalize the data       #
+# ========================================= #
+
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+])
+
+batch_size = 4
+train_data = torchvision.datasets.CIFAR10(
+    root='data',
+    train=True,
+    download=True,
+    transform=transform
+)
+
+test_data = torchvision.datasets.CIFAR10(
+    root='data',
+    train=False,
+    download=True,
+    transform=transform
+)
+
+train_loader = torch.utils.data.DataLoader(
+    dataset=train_data,
+    batch_size=batch_size,
+    shuffle=True
+)
+
+test_loader = torch.utils.data.DataLoader(
+    dataset=test_data,
+    batch_size=batch_size,
+    shuffle=False
+)
+
+classes = ('plane', 'car', 'bird', 'cat',
+           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+
+
+# ========================================= #
+#         Visualize Training Images         #
+# ========================================= #
+
+def imshow(img):
+    img = img / 2 + 0.5  # unnormalize
+    npimg = img.numpy()
+    plt.imshow(np.transpose(npimg, (1, 2, 0)))
+    plt.show()
+
+
+# Get some random training images
+images, labels = next(iter(train_loader))
+imshow(torchvision.utils.make_grid(images))
+
+# Print labels
+print(' '.join('%5s' % classes[labels[j]] for j in range(batch_size)))
+
+
+# ========================================= #
+#    Define Convolutional Neural Network    #
+# ========================================= #
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
+        self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
+        self.fc2 = nn.Linear(in_features=120, out_features=84)
+        self.fc3 = nn.Linear(in_features=84, out_features=10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = torch.relu(x)
+        x = self.pool(x)
+
+        x = self.conv2(x)
+        x = torch.relu(x)
+        x = self.pool(x)
+
+        x = x.view(-1, 16 * 5 * 5)
+
+        x = self.fc1(x)
+        x = torch.relu(x)
+
+        x = self.fc2(x)
+        x = torch.relu(x)
+
+        x = self.fc3(x)
+
+        return x
+
+
+net = Net()
+
+# ========================================= #
+#   Define a Loss function and Optimizer    #
+# ========================================= #
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(params=net.parameters(), lr=0.001, momentum=0.9)
+
+# ========================================= #
+#              Train the network            #
+# ========================================= #
+
+for epoch in range(2):  # loop over the dataset multiple times
+
+    running_loss = 0.0
+    for i, data in enumerate(train_loader, 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()
+
+        # print statistics
+        running_loss += loss.item()
+        if i % 2000 == 1999:  # print every 2000 mini-batches
+            print('[%d, %5d] loss: %.3f' %
+                  (epoch + 1, i + 1, running_loss / 2000))
+            running_loss = 0.0
+
+print('Finished Training')
+
+# Save the trained model
+
+PATH = './cifar_net.pth'
+torch.save(net.state_dict(), PATH)
+
+# ========================================= #
+#               Test the network            #
+# ========================================= #
+
+# Show test images
+images, labels = next(iter(test_loader))
+imshow(torchvision.utils.make_grid(images))
+print('Ground Truth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
+
+# Load the model
+net = Net()
+net.load_state_dict(torch.load(PATH))
+
+outputs = net(images)
+_, predicted = torch.max(outputs, 1)
+print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))
+
+# How network performs on whole test data
+correct = 0
+total = 0
+# For testing we don't need calculate gradients
+with torch.no_grad():
+    for data in test_loader:
+        images, labels = data
+        outputs = net(images)
+        _, predicted = torch.max(outputs.data, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+print(f'Accuracy of the network on test data: {100 * correct / total}')
+
+# ========================================= #
+#            Class-based Accuracy           #
+# ========================================= #
+
+# prepare to count predictions for each class
+correct_pred = {classname: 0 for classname in classes}
+total_pred = {classname: 0 for classname in classes}
+
+# again no gradients needed
+with torch.no_grad():
+    for data in test_loader:
+        images, labels = data
+        outputs = net(images)
+        _, predictions = torch.max(outputs, 1)
+        # collect the correct predictions for each class
+        for label, prediction in zip(labels, predictions):
+            if label == prediction:
+                correct_pred[classes[label]] += 1
+            total_pred[classes[label]] += 1
+
+# print accuracy for each class
+for classname, correct_count in correct_pred.items():
+    accuracy = 100 * float(correct_count) / total_pred[classname]
+    print("Accuracy for class {:5s} is: {:.1f} %".format(classname,
+                                                         accuracy))