{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Data loading performance comparison\n\nThe massive dataset in this competition means that data loading is a significant portion of training time. In this notebook, I've compared 3 different approaches for loading the data:\n\n* Just load the provided PNG files without any preprocessing\n* Significantly compress the data by rendering images as RGB jpeg's (while maintaining image dimensions.)\n* Converting the raw PNGs to preprocessed numpy arrays\n\nConverting the images to JPEGs as described speeds up data loading by a factor of ~3. That's obviously great, but I'm concerned that it signficantly reduces model performance because so much data is lost by the compression. The JPEG images are 13 times smaller on disk than the originals, and while the JPEG algorithm throws out high frequencies that aren't perceptible to humans, I don't know how losing those frequencies affects a CNN.\n\nAs an alternative, we can convert the images to numpy arrays. That leads to a much smaller speed up - about 25% - but it doesn't lead to any data loss. In that sense, it's a free lunch. 25% isn't huge, but with training times on the order of 10 hours, it's also nothing to sneeze at.\n\nAs a final note, I've stitched together code from three authors with my own additions, so apologies for the variety of styles and naming conventions."},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom os import sys\n\nimport imageio\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm_notebook as tqdm\n\nimport torch\nimport torch.utils.data as D\n\nimport torchvision\nfrom torchvision import transforms as T\n\n# Create a dataframe including one experiment's worth of samples\ndf = pd.read_csv('../input/train.csv').head(1106)  # HEPG-01","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# First case: don't do any preprocessing\n# Modeled on https://www.kaggle.com/leighplt/densenet121-pytorch\n# This also serves as a base class for the datasets in cases 2 and 3.\nclass ImageDS(D.Dataset):\n    def __init__(self, df, img_dir, channels=[1,2,3,4,5,6]):\n        self.records = df.to_records(index=False)\n        self.channels = channels\n        self.img_dir = img_dir\n        self.len = df.shape[0] * 2\n        \n    @staticmethod\n    def _load_img_as_tensor(file_name):\n        with Image.open(file_name) as img:\n            return T.ToTensor()(img)\n\n    def _get_img_path(self, index, site, channel):\n        experiment, well, plate = self.records[index].experiment, self.records[index].well, self.records[index].plate\n        return '/'.join([self.img_dir,experiment,f'Plate{plate}',f'{well}_s{site}_w{channel}.png'])\n    \n    def _load_data(self, index, site):\n        paths = [self._get_img_path(index, site, ch) for ch in self.channels]\n        # Although we're normalizing here, the computational cost is insignificant\n        normalize = T.Normalize(\n            mean=[0.5] * 6,\n            std=[0.5] * 6\n        )\n        return normalize(torch.cat([self._load_img_as_tensor(img_path) for img_path in paths]))\n    \n    def __getitem__(self, index):\n        site = (index % 2) + 1\n        index = index // 2\n        return self._load_data(index, site)\n\n    def __len__(self):\n        \"\"\"\n        Total number of samples in the dataset\n        \"\"\"\n        return self.len\n    \ndef loop(data_loader):\n    for _ in data_loader:\n        pass\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Note that we're loading images directly from the input folder.\nds = ImageDS(df, '../input/train')\nloader = D.DataLoader(ds, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)\n%timeit -r 3 loop(loader) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Loading and preprocessing a whole folder of images takes about a minute. That's not-so-great. Let's compare that to converting our 6-channel image sets to 3-channel JPEGs of the same dimension:"},{"metadata":{"trusted":true},"cell_type":"code","source":"!git clone https://github.com/recursionpharma/rxrx1-utils\n!mkdir -p processed-data/jpg\nsys.path.append('rxrx1-utils')\nimport rxrx.io as rio","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Option two: Convert to an rgb jpg\n# Modified version of https://www.kaggle.com/xhlulu/recursion-2019-load-size-and-resize-images\n\ndef convert_to_rgb(df, img_dir='processed-data/jpg/', resize=False, new_size=224, extension='jpeg'):\n    N = df.shape[0]\n    for i in tqdm(range(N)):\n        code = df['id_code'][i]\n        experiment = df['experiment'][i]\n        plate = df['plate'][i]\n        well = df['well'][i]\n        for site in [1, 2]:\n            save_path = f'{img_dir}{code}_s{site}.{extension}'\n\n            im = rio.load_site_as_rgb(\n                'train', experiment, plate, well, site, \n                base_path='../input/'\n            )\n            im = im.astype(np.uint8)\n            im = Image.fromarray(im)\n            \n            if resize:\n                im = im.resize((new_size, new_size), resample=Image.BILINEAR)\n            im.save(save_path)\n\nclass JpgImageDS(ImageDS):\n    def __init__(self, df, img_dir):\n        super().__init__(df, img_dir)\n        \n    def _get_img_path(self, index, site):\n        code = self.records[index].id_code\n        return f'{self.img_dir}{code}_s{site}.jpeg'\n    \n    def _load_data(self, index, site):\n        return self._load_img_as_tensor(self._get_img_path(index, site))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"convert_to_rgb(df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Option two: Load jpegs\nds = JpgImageDS(df, 'processed-data/jpg/')\nloader = D.DataLoader(ds, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)\n%timeit -r 3 loop(loader)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"As I mentioned, this loop is **a lot** faster:  ~20 seconds down from ~60, for a 3x improvement. But it's not a faster way to load data per se. If anything, the need to uncompress the jpeg makes it slower. Instead, we've just reduced the data size by a factor of ~13, from 1.6GB to 126MB:"},{"metadata":{"trusted":true},"cell_type":"code","source":"!du -sh ../input/train/HEPG2-01/\n!du -sh processed-data/jpg/","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"(The 1.6GB size actually understates things: The images are ~3.6gb unzipped.)\n\nLet's see if we can improve performance without throwing out data:"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Option 3: Create preprocessed numpy files.\n# Code mostly from https://www.kaggle.com/gidutz/starter-kernel-recursion-pharmaceuticals\n\nBASE_DIR = '../input'\nOUTPUT_DIR = 'processed-data/npy/'\nDATA_PATH_FORMAT = os.path.join(BASE_DIR, 'train/{experiment}/Plate{plate}/{well}_s{sample}_w{channel}.png')\n\ndf_pixel_stats = pd.read_csv(os.path.join(BASE_DIR, 'pixel_stats.csv')).set_index(['id_code','site', 'channel'])\n\ndef transform_image(sample_data, pixel_data, site):\n    x=[]\n    for channel in [1,2,3,4,5,6]:\n        impath = DATA_PATH_FORMAT.format(experiment=sample.experiment,\n                                        plate=sample_data.plate,\n                                        well=sample_data.well,\n                                        sample=site,\n                                        channel=channel)\n        # normalize the channel\n        img = np.array(imageio.imread(impath)).astype(np.float64)\n        img -= pixel_data.loc[channel]['mean']\n        img /= pixel_data.loc[channel]['std']\n        img *= 255 # To keep MSB\n\n        x.append(img)\n\n    return np.stack(x).T.astype(np.byte)\n\n\n!mkdir -p {OUTPUT_DIR}\nfor _, sample in tqdm(df.iterrows(), total=len(df)):\n    for site in [1, 2]:\n        pixel_data = df_pixel_stats.loc[sample.id_code, site, :].reset_index().set_index('channel')\n        x = transform_image(sample, pixel_data, site)\n        np.save(os.path.join(OUTPUT_DIR, '{sample_id}_s{site}.npy').format(sample_id=sample.id_code, site=site), x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class NpyImageDS(ImageDS):\n    def __init__(self, df, img_dir):\n        super().__init__(df, img_dir)\n        \n    def _get_img_path(self, index, site):\n        sample_id = self.records[index].id_code\n        return f'{self.img_dir}{sample_id}_s{site}.npy'\n    \n    def _load_data(self, index, site):\n        return torch.Tensor(np.load(self._get_img_path(index, site)).astype(np.float32)/ 255.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ds = NpyImageDS(df, OUTPUT_DIR)\nloader = D.DataLoader(ds, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)\n%timeit -r 3 loop(loader) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Loading the .npy files is about 25% faster, although that number can vary from run to run. And although there's nothing in our code that would cause data loss, we can confirm the .npy files are the size we'd expect:"},{"metadata":{"trusted":true},"cell_type":"code","source":"!du -sh ../input/train/HEPG2-01/\n!du -sh processed-data/npy/","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"So why are the numpy files faster? I was worried that the answer would be something anticlimactic, like \"We've already unzipped them\" (unlike the PNG files, which Kaggle stores in a zip folder.) But testing that, I saw no speed up from pre-unzipping the files:"},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -r processed-data\n!mkdir -p processed-data/raw/HEPG2-01\n!cp -r ../input/train/HEPG2-01/ processed-data/raw/HEPG2-01/\nds = ImageDS(df, 'processed-data/raw/HEPG2-01')\nloader = D.DataLoader(ds, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)\n%timeit -r 1 loop(loader) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Loading numpy files appears to just be faster. I don't know why, so if anyone has insight into that please comment."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Apparently these directories need to be removed to avoid an error.\n!rm -r  rxrx1-utils\n!rm -r processed-data","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}