{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch k-fold cross validation DataLoader.\n\nThe purpose of this kernel is to provide an easy to use system for stratified k-fold cross validation with pytorch models.\n***"},{"metadata":{},"cell_type":"markdown","source":"Features:\n* Stratified K-fold cross validation support.\n* Data augmentation support.\n* Test time augmentation support (TTA).\n\n**_If you have any comments or suggestions please don't hesitate to let me know, any feedback is appreciated._**"},{"metadata":{"_uuid":"4dc73910-841f-49fa-bbb2-bcae36b8e237","_cell_guid":"dfb5155a-5b77-4fcd-b139-7bf9c8e36908","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nimport torch\nfrom torch.utils.data import Sampler\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data import BatchSampler\nfrom torchvision.transforms import transforms\nfrom sklearn.model_selection import StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaLeafDataset(Dataset):\n    def __init__(self, \n                 root_dir='../input/cassava-leaf-disease-classification/', \n                 image_dir_name='train_images',\n                 label_csv_name='train.csv',\n                 transform=transforms.ToTensor(),\n                 device='cpu'):\n        super().__init__()\n        self.root_dir = root_dir\n        self.image_dir = os.path.join(self.root_dir, image_dir_name)\n        self.df_labels = pd.read_csv(os.path.join(self.root_dir, label_csv_name))\n        self.transform = transform\n        self.device=device\n        \n    def __getitem__(self, image_id):\n        path = os.path.join(self.image_dir, image_id)\n        label = np.array(self.df_labels[self.df_labels['image_id'] == image_id]['label'].values[0])\n        label = torch.from_numpy(label).to(device)\n        image = Image.open(path)\n        image = self.transform(image).to(device)\n        return (image, label)\n    \n    def __len__(self):\n        return len(self.df_labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageSampler(Sampler):\n    def __init__(self, \n                 sample_idx,\n                 data_source='../input/cassava-leaf-disease-classification/train.csv'):\n        super().__init__(data_source)\n        self.sample_idx = sample_idx\n        self.df_images = pd.read_csv(data_source)\n        \n    def __iter__(self):\n        image_ids = self.df_images['image_id'].loc[self.sample_idx]\n        return iter(image_ids)\n    \n    def __len__(self):\n        return len(self.sample_idx)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageBatchSampler(BatchSampler):\n    def __init__(self, \n                 sampler,\n                 aug_count=5,\n                 batch_size=30,\n                 drop_last=True):\n        super().__init__(sampler, batch_size, drop_last)\n        self.aug_count = aug_count\n        assert self.batch_size % self.aug_count == 0, 'Batch size must be an integer multiple of the aug_count.'\n        \n    def __iter__(self):\n        batch = []\n        \n        for image_id in self.sampler:\n            for i in range(self.aug_count):\n                batch.append(image_id)\n            if len(batch) == self.batch_size:\n                yield batch\n                batch = []\n        if len(batch) > 0 and not self.drop_last:\n            yield batch\n    \n    def __len__(self):\n        if self.drop_last:\n            return len(self.sampler) // self.batch_size\n        else:\n            return (len(self.sampler) + self.batch_size - 1) // self.batch_size","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The create_split_loaders function is used to create the two DataLoaders required for each split.\n\nThe 'aug_count' parameter will be explained later, ignore it for now.\n\nEach split has (k-1) training folds and 1 validation fold. Two DataLoaders are provided for each split. \n\nThe _train_loader_ is used to iterate through all batches in the training folds of a given split. The _valid_loader_ is used to iterate through all samples in the validation fold of a given split.\n\nThe function returns the two DataLoaders for the split in a tuple."},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_split_loaders(dataset, split, aug_count, batch_size):\n    train_folds_idx = split[0]\n    valid_folds_idx = split[1]\n    train_sampler = ImageSampler(train_folds_idx)\n    valid_sampler = ImageSampler(valid_folds_idx)\n    train_batch_sampler = ImageBatchSampler(train_sampler, \n                                            aug_count, \n                                            batch_size)\n    valid_batch_sampler = ImageBatchSampler(valid_sampler, \n                                            aug_count=1, \n                                            batch_size=batch_size,\n                                            drop_last=False)\n    train_loader = DataLoader(dataset, batch_sampler=train_batch_sampler)\n    valid_loader = DataLoader(dataset, batch_sampler=valid_batch_sampler)\n    return (train_loader, valid_loader)    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The get_all_split_loaders function is used to create the DataLoaders for each split.\n\nThe 'aug_count' parameter specifies how many variations of each image in the dataset must be provided. \n\nRandom transformations are applied to each image when it is sampled from the dataset. If you would like your model to be trained seeing 5 different variations of each image in the training set, use 'aug_count=5'. \n\n_Note: Regardless of the value of aug_count all images in the dataset will be seen by the model. The dataset size is N_samples = N_images * aug_count._\n\nThe create_split_loaders function is called for each split, this creates a two DataLoaders for each split. One DataLoader for the samples in the training folds and one DataLoader for the samples in the validation fold.\n\nThe DataLoaders for each split are stored in a tuple. All tuples are stored in a list which is returned by the function."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_all_split_loaders(dataset, cv_splits, aug_count=5, batch_size=30):\n    \"\"\"Create DataLoaders for each split.\n\n    Keyword arguments:\n    dataset -- Dataset to sample from.\n    cv_splits -- Array containing indices of samples to \n                 be used in each fold for each split.\n    aug_count -- Number of variations for each sample in dataset.\n    batch_size -- batch size.\n    \n    \"\"\"\n    split_samplers = []\n    \n    for i in range(len(cv_splits)):\n        split_samplers.append(\n            create_split_loaders(dataset,\n                                 cv_splits[i], \n                                 aug_count, \n                                 batch_size)\n        )\n    return split_samplers","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### The system can be used as follows:\n\nWe begin by loading in the image_ids and their corresponding class labels. This information is stored in the train.csv file."},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Use sklearn to split the dataset into stratified folds."},{"metadata":{"trusted":true},"cell_type":"code","source":"splitter = StratifiedKFold(n_splits=4, shuffle=True, random_state=0)\n\nsplits = []\nfor train_idx, test_idx in splitter.split(df_train['image_id'], df_train['label']):\n    splits.append((train_idx, test_idx))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Specify the image transformations to be used for data augmentation. \n\nCreate an instance of the dataset."},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.RandomAffine(degrees=45, \n                           translate=(0.05,0.05),\n                           scale=(0.8, 1.0)),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomVerticalFlip(p=0.5),\n    transforms.CenterCrop((400, 500)),\n    transforms.ToTensor()\n])\n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\ndataset = CassavaLeafDataset(transform=transform, device=device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Create a list containing the DataLoaders for each split:\n\n[(train_loader_split_1, val_loader_split_1), \n (train_loader_split_2, val_loader_split_2), ... ]"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloaders = get_all_split_loaders(dataset, splits, aug_count=5, batch_size=10)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is the training loop for the model."},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_batch_loader -- batches all samples in training folds.\n# valid_batch_loader -- batches all samples in validation fold.\nfor train_batch_loader, valid_batch_loader in dataloaders:\n    # Loop through all batches in training folds for a given split.\n    for batch in train_batch_loader:\n        # Train model on the training folds in the split.\n        break\n    \n    # Loop through all batches in validation fold for a given split.\n    for batch in valid_batch_loader:\n        # Test model on the validation fold in the split.   \n        break\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Example training batch:\n\nbatch_size = 10\n\naug_count = 5"},{"metadata":{"trusted":true,"_kg_hide-input":false},"cell_type":"code","source":"for train_batch_loader, valid_batch_loader in dataloaders:\n    for batch in train_batch_loader:\n        train_batch = batch\n        break\n    break\n\nfig, ax = plt.subplots(2,5, figsize=(20, 10))\nimages = train_batch[0]\nfor i,image in enumerate(images):\n    image = image.transpose(0,2)\n    if i < 5:\n        ax[0, i].imshow(image)\n    else:\n        ax[1, i%5].imshow(image)\nplt.show()","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}