19. Image classification with neural networksΒΆ

Here we use Pytorch and a Resnext neural networks to classify images on the Kaggle fruits dataset.

[1]:
import numpy as np
import scipy.stats as stats
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import pickle
import torchvision
from torchvision import transforms

%matplotlib inline

[2]:
from google.colab import drive
drive.mount("/gdrive")
Mounted at /gdrive
[3]:
# archive.zip downloaded from https://www.kaggle.com/moltean/fruits
!cp /gdrive/MyDrive/datasets/kaggle-fruits/archive.zip /content
[4]:
!unzip -q -d /content/data /content/data.zip
[5]:
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = torchvision.datasets.ImageFolder(
    '/content/data/fruits-360/Training/',
    transform=preprocess
)

test_dataset = torchvision.datasets.ImageFolder(
    '/content/data/fruits-360/Test/',
    transform=preprocess
)

print(len(train_dataset))
print(train_dataset[0][0].shape)
67692
torch.Size([3, 100, 100])
[6]:
class NeuralNetEstimator():
    def __init__(self, lr=0.001, random_state=None,
                 train_on_a_small_subset_of_data=False):
        self.lr = lr
        self.random_state = random_state
        self.train_on_a_small_subset_of_data = train_on_a_small_subset_of_data

    def fit(self, dataset):
        self.dataset = dataset
        if self.random_state is not None:
            torch.manual_seed(self.random_state)

        self.net = torch.hub.load('pytorch/vision:v0.9.1', 'resnext50_32x4d')
        self.nclasses = len(np.unique(dataset.targets))
        self.net.fc = nn.Linear(net.fc.in_features, self.nclasses)
        cuda = torch.cuda.is_available()

        if cuda:
            self.net.cuda()

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(self.net.parameters(), lr=self.lr)

        print("Starting optimization.")
        self.train_losses = []
        self.val_losses = []

        # db split
        train_idx = torch.randperm(len(dataset))
        val_idx = train_idx[:len(dataset)//10]
        train_idx = train_idx[len(dataset)//10:]

        train_dataset = torch.utils.data.Subset(dataset, train_idx)
        val_dataset = torch.utils.data.Subset(dataset, val_idx)

        if self.train_on_a_small_subset_of_data:
            train_dataset = torch.utils.data.Subset(train_dataset, range(200))
            val_dataset = torch.utils.data.Subset(val_dataset, range(100))

        # initial values for early stopping decision state parameters
        last_val_loss = np.inf
        es_tries = 0

        for epoch in range(100_000):
            try:
                # network training step
                self.net.train()
                dataset_loader_train = torch.utils.data.DataLoader(
                      train_dataset, batch_size=100, shuffle=True,
                      pin_memory=cuda, drop_last=True
                )
                batch_losses = []
                for batch_inputv, batch_target in dataset_loader_train:
                    if cuda:
                        batch_inputv = batch_inputv.cuda()
                        batch_target = batch_target.cuda()

                    optimizer.zero_grad()
                    output = self.net(batch_inputv)
                    batch_loss = criterion(output, batch_target)

                    batch_loss.backward()
                    optimizer.step()
                    batch_losses.append(batch_loss.item())
                loss = np.mean(batch_losses)
                self.train_losses.append(loss)
                print('\rTrain loss', np.round(loss.item(), 2), end='')

                # network evaluation step
                self.net.eval()
                with torch.no_grad():
                    dataset_loader_val = torch.utils.data.DataLoader(
                          val_dataset, batch_size=100, shuffle=False,
                          pin_memory=cuda, drop_last=False,
                    )
                    batch_losses = []
                    batch_sizes = []
                    for batch_inputv, batch_target in dataset_loader_val:
                        if cuda:
                            batch_inputv = batch_inputv.cuda()
                            batch_target = batch_target.cuda()

                        optimizer.zero_grad()
                        output = self.net(batch_inputv)
                        batch_loss = criterion(output, batch_target)

                        batch_losses.append(batch_loss.item())
                        batch_sizes.append(len(batch_inputv))
                    loss = np.average(batch_losses, weights=batch_sizes)
                    self.val_losses.append(loss)

                    print(' | Validation loss', np.round(loss.item(), 2),
                          'in epoch', epoch + 1, end='')

                    # Decisions based on the evaluated values
                    if loss < last_val_loss:
                        best_state_dict = self.net.state_dict()
                        best_state_dict = pickle.dumps(best_state_dict)
                        es_tries = 0
                        last_val_loss = loss
                    else:
                        if es_tries in [20, 40]:
                            self.net.load_state_dict(pickle.loads(best_state_dict))

                        if es_tries >= 60:
                            self.net.load_state_dict(pickle.loads(best_state_dict))
                            break

                        es_tries += 1
                    print(' | es_tries', es_tries, end='', flush=True)
            except KeyboardInterrupt:
                if epoch > 0:
                     print("\nKeyboard interrupt detected.",
                           "Switching weights to lowest validation loss",
                           "and exiting")
                     self.net.load_state_dict(pickle.loads(best_state_dict))
                break
        print(f"\nOptimization finished in {epoch+1} epochs.")

    def get_loss(self, dataset):
        criterion = nn.CrossEntropyLoss()
        cuda = torch.cuda.is_available()

        self.net.eval()
        with torch.no_grad():
            dataset_loader_val = torch.utils.data.DataLoader(
                  dataset, batch_size=100, shuffle=False,
                  pin_memory=cuda, drop_last=False,
            )
            cv_batch_losses = []
            zo_batch_losses = []
            batch_sizes = []
            for batch_inputv, batch_target in dataset_loader_val:
                if cuda:
                    batch_inputv = batch_inputv.cuda()
                    batch_target = batch_target.cuda()

                output = self.net(batch_inputv)
                cv_batch_loss = criterion(output, batch_target)
                zo_batch_loss = (torch.argmax(output, 1) != batch_target).cpu()
                zo_batch_loss = np.array(zo_batch_loss).mean()

                cv_batch_losses.append(cv_batch_loss.item())
                zo_batch_losses.append(zo_batch_loss.item())
                batch_sizes.append(len(batch_inputv))
            cv_loss = np.average(cv_batch_losses, weights=batch_sizes)
            zo_loss = np.average(zo_batch_losses, weights=batch_sizes)

            return cv_loss, zo_loss

nn_estimator = NeuralNetEstimator(
    random_state=0,
    #train_on_a_small_subset_of_data=True,
)
nn_estimator.fit(train_dataset)
# save results
with open('/gdrive/MyDrive/datasets/kaggle-fruits/model.pkl', 'wb') as f:
     pickle.dump(nn_estimator, f)

loss_on_test = nn_estimator.get_loss(test_dataset)
print(f"Loss on test dataset for default parameters: {loss_on_test}")
Using cache found in /root/.cache/torch/hub/pytorch_vision_v0.9.1
Starting optimization.
Train loss 0.01 | Validation loss 0.0 in epoch 17 | es_tries 6
Keyboard interrupt detected. Switching weights to lowest validation loss and exiting

Optimization finished in 18 epochs.
Loss on test dataset for default parameters: (0.08706481696378363, 0.01564703808180536)