{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Fine-Tuning with PyTorch\n\nThroughout this notebook, we use the <a href='https://www.kaggle.com/c/cifar-10/overview'>CIFAR-10</a> dataset from Kaggle, a popular computer-vision dataset of 60,000 32x32 color images to be classified in one of ten classes, with 6000 images per class. This dataset is complex enough to give a good idea of the benefits of fine-tuning, and why this process can be used to achieve high accuracy scores without spending hours and money to re-train complex models from scratch. With **default hyperparameters** and a medium-sized model (ResNet34), we are able to achieve around **96% accuracy** on Kaggle, 3% off the state-of-the-art for this dataset.","metadata":{}},{"cell_type":"markdown","source":"# Table of Contents\n1. [Extracting Data](#extraction)\n2. [Datasets and DataLoaders](#data)  \n3. [Training with Validation](#validation)  \n4. [Full Training](#training)\n5. [Generating Predictions](#testing) ","metadata":{}},{"cell_type":"markdown","source":"## Extracting Data <a name=\"extraction\"></a>\nThe original CIFAR-10 dataset provided by Kaggle is composed of two *.7z* files (*train.7z*, *test.7z*), the labels (*trainLabels.csv*) and an example submission (*sampleSubmission.csv*). We provide these files in the *data* folder.  \n\nFirst, we extract the two zipped files in two folers called *original_train* and *original_test*. Then, we move these files to be in the structure required by the `ImageDataset` class of PyTorch, where every image is stored in a folder named as its label (e.g., the nth airplane image will be stored under *airplane/n.png*). We will create four subfolders within *data*:\n* *train*: contains Training data (excluding Validation set), used to train the model during hyperparameter search\n* *valid*: contains Validation data, used to train the model during hyperparameter search\n* *train_valid*: contains Training+Validation data together, used for full re-training of the model\n* *test* contains Test data, i.e., all the unlabelled data we must submit to Kaggle together with a predicted label","metadata":{}},{"cell_type":"code","source":"# py7zr is required to extract the .7tz files\n\n!pip install -q py7zr","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:34:41.443660Z","iopub.execute_input":"2022-09-19T09:34:41.443988Z","iopub.status.idle":"2022-09-19T09:34:50.344570Z","shell.execute_reply.started":"2022-09-19T09:34:41.443963Z","shell.execute_reply":"2022-09-19T09:34:50.343752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import py7zr\nimport shutil\nfrom pathlib import Path\n\nroot = Path('./data')\ninput_path = Path('../input/cifar-10')\n\nwith py7zr.SevenZipFile(input_path/'train.7z', mode='r') as z:\n    z.extractall(root)\n\nwith py7zr.SevenZipFile(input_path/'test.7z', mode='r') as z:\n    z.extractall(root)\n\nshutil.copy(input_path/'trainLabels.csv', root/'trainLabels.csv')\n\n(root/'train').rename(root/'original_train')\n(root/'test').rename(root/'original_test')","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:34:50.346963Z","iopub.execute_input":"2022-09-19T09:34:50.347343Z","iopub.status.idle":"2022-09-19T09:49:12.507813Z","shell.execute_reply.started":"2022-09-19T09:34:50.347306Z","shell.execute_reply":"2022-09-19T09:49:12.505714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from random import random\nimport os\n\ndef copy_file(source_directory, destination_directory, filename):\n    \"\"\"\n    Utility function used to copy a file from a source_directory to a destination_directory\n    \"\"\"\n    destination_directory.mkdir(parents=True, exist_ok=True)\n    shutil.copy(source_directory/filename, destination_directory/filename)\n    \ndef organize_train_valid_dataset(root, labels, valid_probability=0.1):\n    \"\"\"\n    Creates the train, train_valid and valid folders respecting PyTorch's ImageDataset structure, performing\n    train/validation split based on the given percentage\n    \"\"\"\n    source_directory = root/'original_train'\n    \n    with os.scandir(source_directory) as it:\n        for entry in it:\n            if entry.is_file():\n                img_index = entry.name.split('.')[0]  # The index is the name of the image except the extension\n                img_class = labels[labels.id==int(img_index)].label.values[0]  # Find the class by looking up the index in the DF\n                \n                # Randomly assign the image to the valid dataset with probability 'valid_probability'\n                channel = Path('train') if random()>valid_probability else Path('valid')\n                destination_directory = root/channel/img_class\n                \n                # Copy the image to either the train or valid folder, and also to the train_valid folder\n                copy_file(source_directory, destination_directory, entry.name)\n                copy_file(source_directory, root/'train_valid'/img_class, entry.name)\n\ndef organize_test_dataset(root):\n    \"\"\"\n    Creates the test folder respecting PyTorch's ImageDataset structure, using a dummy 'undefined' label\n    \"\"\"\n    source_directory = root/'original_test'\n        \n    with os.scandir(source_directory) as it:\n        for entry in it:\n            if entry.is_file():\n                img_index = entry.name.split('.')[0]  # The index is the name of the image except the extension\n\n                channel = Path('test')\n                destination_directory = root/channel/'undefined'\n\n                copy_file(source_directory, destination_directory, entry.name)","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:49:12.512892Z","iopub.execute_input":"2022-09-19T09:49:12.513359Z","iopub.status.idle":"2022-09-19T09:49:12.526227Z","shell.execute_reply.started":"2022-09-19T09:49:12.513322Z","shell.execute_reply":"2022-09-19T09:49:12.524776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\n# Read in the labels DataFrame with a label for each image\nlabels = pd.read_csv(root/'trainLabels.csv')\n\n# Create the train/train_valid/valid folder structure\nvalid_probability = 0.1\norganize_train_valid_dataset(root, labels, valid_probability)\n\n# Create the test folder structure\norganize_test_dataset(root)","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:49:12.527336Z","iopub.execute_input":"2022-09-19T09:49:12.528124Z","iopub.status.idle":"2022-09-19T09:51:27.061308Z","shell.execute_reply.started":"2022-09-19T09:49:12.528097Z","shell.execute_reply":"2022-09-19T09:51:27.060289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Datasets and DataLoaders <a name='data'></a>\n\nAs mentioned above, we rely on the `ImageDataset` class of PyTorch to create the required datasets for training, validation, training+validation and testing. Out of each dataset, we then create a DataLoader to be used in the training/evaluation loops to efficiently fetch images in batches from disk.","metadata":{}},{"cell_type":"markdown","source":"We perform an initial step to load in the train data and compute the mean and standard deviation of the dataset for each channel (R, G, B), across all images and all pixels. We compute a mean and stdev value batch-by-batch to avoid loading the entire dataset in memory, and then compute the mean of the means and of the stdevs.  \n**NOTE**: if you have enough RAM (or memory on the GPU), you can use a batch_size equal to the entire train_dataset length, it will provide a more accurate estimation of the means and stdevs by channel.","metadata":{}},{"cell_type":"code","source":"!pip install -q --upgrade torchvision","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:51:27.063743Z","iopub.execute_input":"2022-09-19T09:51:27.064092Z","iopub.status.idle":"2022-09-19T09:52:45.556307Z","shell.execute_reply.started":"2022-09-19T09:51:27.064068Z","shell.execute_reply":"2022-09-19T09:52:45.555119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torchvision\nimport torch\n\ntrain_dataset = torchvision.datasets.ImageFolder(\n    root/'train', \n    transform=torchvision.transforms.Compose([\n        # Resize step is required as we will use a ResNet model, which accepts at leats 224x224 images\n        torchvision.transforms.Resize((224,224)),  \n        torchvision.transforms.ToTensor(),\n    ])\n)\n\ntrain_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=512, shuffle=False, num_workers=2, pin_memory=True)\n\nmeans = []\nstdevs = []\nfor X, _ in train_dataloader:\n    # Dimensions 0,2,3 are respectively the batch, height and width dimensions\n    means.append(X.mean(dim=(0,2,3)))\n    stdevs.append(X.std(dim=(0,2,3)))\n\nmean = torch.stack(means, dim=0).mean(dim=0)\nstdev = torch.stack(stdevs, dim=0).mean(dim=0)","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:52:45.559395Z","iopub.execute_input":"2022-09-19T09:52:45.559768Z","iopub.status.idle":"2022-09-19T09:53:42.705502Z","shell.execute_reply.started":"2022-09-19T09:52:45.559734Z","shell.execute_reply":"2022-09-19T09:53:42.703788Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The transforms used for the training and training+validation datasets consist of resizing the images to the required resolution by our ResNet model (224x224), using the `AutoAugment` policy learned on the CIFAR10 dataset and finally converting the image from PIL to Tensor. For the validation and test sets we just resize the image and convert it to Tensor format.","metadata":{}},{"cell_type":"code","source":"train_transforms = torchvision.transforms.Compose([\n        torchvision.transforms.Resize((224,224)),\n        torchvision.transforms.AutoAugment(policy=torchvision.transforms.AutoAugmentPolicy.CIFAR10),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize(mean, stdev)\n    ])\n\ntrain_dataset, train_valid_dataset = [torchvision.datasets.ImageFolder(folder, transform=train_transforms) for folder in [root/'train', root/'train_valid']]\n\n\nvalid_transforms = torchvision.transforms.Compose([\n        torchvision.transforms.Resize((224,224)),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize(mean, stdev)\n    ])\n\nvalid_dataset, test_dataset = [torchvision.datasets.ImageFolder(folder, transform=valid_transforms) for folder in [root/'valid', root/'test']]","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:53:42.707286Z","iopub.execute_input":"2022-09-19T09:53:42.707993Z","iopub.status.idle":"2022-09-19T09:53:44.268666Z","shell.execute_reply.started":"2022-09-19T09:53:42.707962Z","shell.execute_reply":"2022-09-19T09:53:44.266797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The train and train+validation DataLoaders use a smaller `batch_size` as we will also need to keep track of gradients in memory. Furthermore, we shuffle the dataset each epoch to avoid loading the batches in the same order.\nThe valid and test DataLoaders use a larger `batch_size` and do not required to shuffle the dataset as we want deterministic results.\n\nThe number of workers is generally set to `2 * num_gpus` as a rule of thumb for Kaggle, with `pin_memory = True` to speed up data transfer to the GPU.","metadata":{}},{"cell_type":"code","source":"num_gpus = torch.cuda.device_count()\n\ntrain_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2*num_gpus, pin_memory=True)\ntrain_valid_dataloader = torch.utils.data.DataLoader(train_valid_dataset, batch_size=128, shuffle=True, num_workers=2*num_gpus, pin_memory=True)\n\nvalid_dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=256, shuffle=False, num_workers=2*num_gpus, pin_memory=True)\ntest_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=2*num_gpus, pin_memory=True)","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:53:44.271129Z","iopub.execute_input":"2022-09-19T09:53:44.271816Z","iopub.status.idle":"2022-09-19T09:53:44.283605Z","shell.execute_reply.started":"2022-09-19T09:53:44.271778Z","shell.execute_reply":"2022-09-19T09:53:44.282179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training with Validation <a name='validation'></a>\n\nThe first step of the process is to evaluate the model performance on our own Validation set, consisting of 10% of the labelled data we get from Kaggle. Ideally, this step would be performed while finding the best model and hyperparameters to improve the final accuracy. Here, we just perform this step to show the expected model accuracy before submitting the results to Kaggle.  \n**NOTE**: when doing proper hyperparameter tuning, depending on the size of the overall labelle data, a k-fold approach might be more appropriate to estimate the generalization capability of the model.","metadata":{}},{"cell_type":"markdown","source":"We fine-tune a ResNet34 model, trained on ImageNet. Other models might be used, but for the purpose of this notebook a ResNet34 is a good trade-off between training time and model accuracy.  \nThe model originally has a 1000-dimensional output layer, but our dataset has only 10 classes, so we remove the output layer and define a new Fully-Connected layer with just 10 neurons, one for each class in CIFAR-10. The parameters of these new neurons are initialized with Xavier initialization.","metadata":{}},{"cell_type":"code","source":"def get_net():\n    resnet = torchvision.models.resnet34(pretrained=True)\n    \n    # Substitute the FC output layer\n    resnet.fc = torch.nn.Linear(resnet.fc.in_features, 10)\n    torch.nn.init.xavier_uniform_(resnet.fc.weight)\n    return resnet","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:53:44.285099Z","iopub.execute_input":"2022-09-19T09:53:44.285398Z","iopub.status.idle":"2022-09-19T09:53:44.305283Z","shell.execute_reply.started":"2022-09-19T09:53:44.285351Z","shell.execute_reply":"2022-09-19T09:53:44.303026Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The training loop is a standard PyTorch loop where for every epoch we perform the following macro steps:\n1. Iterate over the Train DataLoader by making predictions, calculating loss, backpropagating gradients and updating parameters\n2. Iterate over the Valid DataLoader (if present) to compute the validation loss and accuracy\n3. Decrease the learning rate using the scheduler (if present)\n4. Optionally, store the model checkpoint after a given number of `checkpoint_epochs`","metadata":{}},{"cell_type":"code","source":"import time\n\ndef train(net, train_dataloader, valid_dataloader, criterion, optimizer, scheduler=None, epochs=10, device='cpu', checkpoint_epochs=10):\n    start = time.time()\n    print(f'Training for {epochs} epochs on {device}')\n    \n    for epoch in range(1,epochs+1):\n        print(f\"Epoch {epoch}/{epochs}\")\n        \n        net.train()  # put network in train mode for Dropout and Batch Normalization\n        train_loss = torch.tensor(0., device=device)  # loss and accuracy tensors are on the GPU to avoid data transfers\n        train_accuracy = torch.tensor(0., device=device)\n        for X, y in train_dataloader:\n            X = X.to(device)\n            y = y.to(device)\n            preds = net(X)\n            loss = criterion(preds, y)\n            \n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            \n            with torch.no_grad():\n                train_loss += loss * train_dataloader.batch_size\n                train_accuracy += (torch.argmax(preds, dim=1) == y).sum()\n        \n        if valid_dataloader is not None:\n            net.eval()  # put network in train mode for Dropout and Batch Normalization\n            valid_loss = torch.tensor(0., device=device)\n            valid_accuracy = torch.tensor(0., device=device)\n            with torch.no_grad():\n                for X, y in valid_dataloader:\n                    X = X.to(device)\n                    y = y.to(device)\n                    preds = net(X)\n                    loss = criterion(preds, y)\n\n                    valid_loss += loss * valid_dataloader.batch_size\n                    valid_accuracy += (torch.argmax(preds, dim=1) == y).sum()\n        \n        if scheduler is not None: \n            scheduler.step()\n            \n        print(f'Training loss: {train_loss/len(train_dataloader.dataset):.2f}')\n        print(f'Training accuracy: {100*train_accuracy/len(train_dataloader.dataset):.2f}')\n        \n        if valid_dataloader is not None:\n            print(f'Valid loss: {valid_loss/len(valid_dataloader.dataset):.2f}')\n            print(f'Valid accuracy: {100*valid_accuracy/len(valid_dataloader.dataset):.2f}')\n        \n        if epoch%checkpoint_epochs==0:\n            torch.save({\n                'epoch': epoch,\n                'state_dict': net.state_dict(),\n                'optimizer': optimizer.state_dict(),\n            }, './checkpoint.pth.tar')\n        \n        print()\n    \n    end = time.time()\n    print(f'Total training time: {end-start:.1f} seconds')\n    return net","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:53:44.307137Z","iopub.execute_input":"2022-09-19T09:53:44.307615Z","iopub.status.idle":"2022-09-19T09:53:44.326650Z","shell.execute_reply.started":"2022-09-19T09:53:44.307580Z","shell.execute_reply":"2022-09-19T09:53:44.324461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"In this notebook we only use of at most one GPU, you can freely refactor the code to use DistributedDataParallel if you have more GPUs and/or devices.  \n\nWhen fine-tuning, the model parameters of the network body are trained using a lower learning rate than for the head, since for the latter we have to train them from scratch. We rely on Parameter Groups from PyTorch to define two learning rates for the two groups, and use Adam optimizer with `weight_decay = 5e-4` (find via hyperparameter search).","metadata":{}},{"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\nlr, weight_decay, epochs = 1e-5, 5e-4, 20\n\nnet = get_net().to(device)\n\n# Standard CrossEntropy Loss for multi-class classification problems\ncriterion = torch.nn.CrossEntropyLoss()\n\n# params_1x are the parameters of the network body, i.e., of all layers except the FC layers\nparams_1x = [param for name, param in net.named_parameters() if 'fc' not in str(name)]\noptimizer = torch.optim.Adam([{'params':params_1x}, {'params': net.fc.parameters(), 'lr': lr*10}], lr=lr, weight_decay=weight_decay)\n\nnet = train(net, train_dataloader, valid_dataloader, criterion, optimizer, None, epochs, device)","metadata":{"execution":{"iopub.status.busy":"2022-09-19T09:53:44.328487Z","iopub.execute_input":"2022-09-19T09:53:44.328839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Full Training <a name='training'></a>\nAfter assessing the model performance on the Validation set, we want to re-train the model on the full Training + Validation data to squeeze every performance left before submitting our results to Kaggle. As a general rule, the more data we train on, the better the results will be.","metadata":{}},{"cell_type":"code","source":"lr, weight_decay, epochs = 1e-5, 5e-4, 20\n\nnet = get_net().to(device)\n\ncriterion = torch.nn.CrossEntropyLoss()\n\nparams_1x = [param for name, param in net.named_parameters() if 'fc' not in str(name)]\noptimizer = torch.optim.Adam([{'params':params_1x}, {'params': net.fc.parameters(), 'lr': lr*10}], lr=lr, weight_decay=weight_decay)\n\nnet = train(net, train_valid_dataloader, None, criterion, optimizer, None, epochs, device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generating Predictions <a name='testing'></a>\nAfter re-training the network on the full labelled dataset, we are ready to score the Test set and submit out results to Kaggle. We iterate over the test_dataloader to get a predicted label for each Test image, and then create a final DataFrame like the one provided in *sampleSubmission.csv* to use on Kaggle.","metadata":{}},{"cell_type":"code","source":"preds = []\n\nnet.eval()\nwith torch.no_grad():\n    for X, _ in test_dataloader:\n        X = X.to(device)\n        preds.extend(net(X).argmax(dim=1).type(torch.int32).cpu().numpy())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ids = list(range(1, len(test_dataset)+1))\nids.sort(key=lambda x: str(x))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({'id': ids, 'label': preds})\ndf['label'] = df['label'].apply(lambda x: train_dataset.classes[x])\ndf.to_csv('submission.csv', index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}