{"cells":[{"metadata":{},"cell_type":"markdown","source":"# How to correctly keep class balance using TPU in distributed training\n\nHi, everyone!\n\nMy name is Alex Shonenkov, I am DL/NLP/CV/TS research engineer. Especially I am in Love with NLP & DL.\n\nFirst of all, I would like to say all of you \"Thanks!\" for the fact that you find my ideas interesting and useful!\n\n- [[TPU-Inference] Super Fast XLMRoberta](https://www.kaggle.com/shonenkov/tpu-inference-super-fast-xlmroberta)\n- [NLP Albumentations](https://www.kaggle.com/shonenkov/nlp-albumentations)\n- [Hack with Parallel Corpus](https://www.kaggle.com/shonenkov/hack-with-parallel-corpus)\n\nIn this kernel I would like to demonstrate simple technique for keeping balance class during distributed training with PyTorch / XLA\n"},{"metadata":{},"cell_type":"markdown","source":"# Motivation\n\nI've studied the best of the best provided public training kernels for this competition and understood, that authors didn't use balance class:\n- [I Like Clean TPU Training Kernels & I Can Not Lie](https://www.kaggle.com/abhishek/i-like-clean-tpu-training-kernels-i-can-not-lie) by [@abhishek](https://www.kaggle.com/abhishek)\n- [Jigsaw Multilingual Toxicity : EDA + Models 🤬](https://www.kaggle.com/tarunpaparaju/jigsaw-multilingual-toxicity-eda-models) by [@tarunpaparaju](https://www.kaggle.com/tarunpaparaju)\n- [Deep Learning For NLP: Zero To Transformers & BERT](https://www.kaggle.com/tanulsingh077/deep-learning-for-nlp-zero-to-transformers-bert) by [@tanulsingh077](https://www.kaggle.com/tanulsingh077)\n\n\nThanks a lot all of them for their works!\n\n\n[@xhlulu](https://www.kaggle.com/xhlulu) said in [discussion](https://www.kaggle.com/c/jigsaw-multilingual-toxic-comment-classification/discussion/140254):\n\n> Balance your data: The training data set consists of over 2M data points, but it is heavily unbalanced since it only has around 120k positive labels. In order to limit training time, I selected all of the positive labels, and subsampled the negative labels such that I have in total 400k labels, which is a lot closer to the class frequency of the validation labels.\n\n\nI also think the balance class gives really good boost in this competition!\n\nYou can see how [@xhlulu](https://www.kaggle.com/xhlulu) realized balance class keeping:\n- [Jigsaw TPU: XLM-Roberta](https://www.kaggle.com/xhlulu/jigsaw-tpu-xlm-roberta) by [@xhlulu](https://www.kaggle.com/xhlulu)\n\n\nHe used sampler with pandas DataFrame:\n\n![](https://i.ibb.co/D8QMZQn/2020-05-01-19-09-45.png)\n\nIt has some problems: not all data can participate in training! Also If you used this method in every epoch you can find OOM problems that often meet with XLA / PyTorch"},{"metadata":{},"cell_type":"markdown","source":"# MAIN IDEA\n\nLets use modificated pytorch sampler for our loaders for retrieve balanced dataset in streaming! \nAlso lets don't forget about distributed training!\n\nI would like to demostrate technique how I decided this problem for my research pipelines."},{"metadata":{},"cell_type":"markdown","source":"### I would like to say thanks [Catalyst Team](https://github.com/catalyst-team/catalyst) for creating the great library!\n\nBut this library can't use TPU backend for now... \n\nSo let's start!"},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py > /dev/null\n!python pytorch-xla-env-setup.py --version 20200420 --apt-packages libomp5 libopenblas-dev > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"import pandas as pd\nimport torch\nfrom torch.utils.data import Dataset,DataLoader\nfrom tqdm import tqdm\nfrom datetime import datetime\nimport os\nimport warnings\nfrom glob import glob\nimport time\n\nwarnings.filterwarnings(\"ignore\")\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data.distributed import DistributedSampler\nfrom typing import Iterator, List, Optional\nfrom torch.utils.data.sampler import Sampler\nfrom torch.utils.data.dataset import Dataset\nfrom operator import itemgetter\nimport numpy as np\n\n##################################\n## parts of code from catalyst: ##\n##################################\n\nclass DatasetFromSampler(Dataset):\n    \"\"\"Dataset of indexes from `Sampler`.\"\"\"\n\n    def __init__(self, sampler: Sampler):\n        \"\"\"\n        Args:\n            sampler (Sampler): @TODO: Docs. Contribution is welcome\n        \"\"\"\n        self.sampler = sampler\n        self.sampler_list = None\n\n    def __getitem__(self, index: int):\n        \"\"\"Gets element of the dataset.\n        Args:\n            index (int): index of the element in the dataset\n        Returns:\n            Single element by index\n        \"\"\"\n        if self.sampler_list is None:\n            self.sampler_list = list(self.sampler)\n        return self.sampler_list[index]\n\n    def __len__(self) -> int:\n        \"\"\"\n        Returns:\n            int: length of the dataset\n        \"\"\"\n        return len(self.sampler)\n\n\nclass DistributedSamplerWrapper(DistributedSampler):\n    \"\"\"\n    Wrapper over `Sampler` for distributed training.\n    Allows you to use any sampler in distributed mode.\n    It is especially useful in conjunction with\n    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each\n    process can pass a DistributedSamplerWrapper instance as a DataLoader\n    sampler, and load a subset of subsampled data of the original dataset\n    that is exclusive to it.\n    .. note::\n        Sampler is assumed to be of constant size.\n    \"\"\"\n\n    def __init__(\n        self,\n        sampler,\n        num_replicas: Optional[int] = None,\n        rank: Optional[int] = None,\n        shuffle: bool = True,\n    ):\n        \"\"\"\n        Args:\n            sampler: Sampler used for subsampling\n            num_replicas (int, optional): Number of processes participating in\n                distributed training\n            rank (int, optional): Rank of the current process\n                within ``num_replicas``\n            shuffle (bool, optional): If true (default),\n                sampler will shuffle the indices\n        \"\"\"\n        super(DistributedSamplerWrapper, self).__init__(\n            DatasetFromSampler(sampler),\n            num_replicas=num_replicas,\n            rank=rank,\n            shuffle=shuffle,\n        )\n        self.sampler = sampler\n\n    def __iter__(self):\n        \"\"\"@TODO: Docs. Contribution is welcome.\"\"\"\n        self.dataset = DatasetFromSampler(self.sampler)\n        indexes_of_indexes = super().__iter__()\n        subsampler_indexes = self.dataset\n        return iter(itemgetter(*indexes_of_indexes)(subsampler_indexes))\n\nclass BalanceClassSampler(Sampler):\n    \"\"\"Abstraction over data sampler.\n    Allows you to create stratified sample on unbalanced classes.\n    \"\"\"\n\n    def __init__(self, labels: List[int], mode: str = \"downsampling\"):\n        \"\"\"\n        Args:\n            labels (List[int]): list of class label\n                for each elem in the datasety\n            mode (str): Strategy to balance classes.\n                Must be one of [downsampling, upsampling]\n        \"\"\"\n        super().__init__(labels)\n\n        labels = np.array(labels)\n        samples_per_class = {\n            label: (labels == label).sum() for label in set(labels)\n        }\n\n        self.lbl2idx = {\n            label: np.arange(len(labels))[labels == label].tolist()\n            for label in set(labels)\n        }\n\n        if isinstance(mode, str):\n            assert mode in [\"downsampling\", \"upsampling\"]\n\n        if isinstance(mode, int) or mode == \"upsampling\":\n            samples_per_class = (\n                mode\n                if isinstance(mode, int)\n                else max(samples_per_class.values())\n            )\n        else:\n            samples_per_class = min(samples_per_class.values())\n\n        self.labels = labels\n        self.samples_per_class = samples_per_class\n        self.length = self.samples_per_class * len(set(labels))\n\n    def __iter__(self) -> Iterator[int]:\n        \"\"\"\n        Yields:\n            indices of stratified sample\n        \"\"\"\n        indices = []\n        for key in sorted(self.lbl2idx):\n            replace_ = self.samples_per_class > len(self.lbl2idx[key])\n            indices += np.random.choice(\n                self.lbl2idx[key], self.samples_per_class, replace=replace_\n            ).tolist()\n        assert len(indices) == self.length\n        np.random.shuffle(indices)\n\n        return iter(indices)\n\n    def __len__(self) -> int:\n        \"\"\"\n        Returns:\n             length of result sample\n        \"\"\"\n        return self.length","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Examples for usage and understanding:"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DatasetRetriever(Dataset):\n\n    def __init__(self, df):\n        self.labels = df['toxic'].values\n\n    def __len__(self):\n        return self.labels.shape[0]\n\n    def __getitem__(self, idx):\n        label = self.labels[idx]\n        return label\n    \n    def get_labels(self):\n        return list(self.labels )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ndf_train = pd.read_csv(f'../input/jigsaw-multilingual-toxic-comment-classification/jigsaw-toxic-comment-train.csv', index_col='id')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = DatasetRetriever(df_train)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Usually loader:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=16,\n    pin_memory=False,\n    drop_last=False,\n    num_workers=2\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result = {'toxic': []}\nfor labels in tqdm(train_loader, total=len(train_loader)):\n    result['toxic'].extend(labels.numpy())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.DataFrame(result)['toxic'].hist()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Using balance sampler:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=16,\n    sampler=BalanceClassSampler(labels=train_dataset.get_labels(), mode='downsampling'),  # here 2 modes: downsampling/upsampling\n    pin_memory=False,\n    drop_last=False,\n    num_workers=2\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result = {'toxic': []}\nfor labels in tqdm(train_loader, total=len(train_loader)):\n    result['toxic'].extend(labels.numpy())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.DataFrame(result)['toxic'].hist()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Now lets to say about distributed usage with this sampler!"},{"metadata":{},"cell_type":"markdown","source":"usually loader:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_experiment1(device):\n    if not os.path.exists('experiment1'):\n        os.makedirs('experiment1')\n        \n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=16,\n        sampler=train_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=1\n    )\n    para_loader = pl.ParallelLoader(train_loader, [device])\n\n    result = {'toxic': []}\n    for labels in para_loader.per_device_loader(device):\n        result['toxic'].extend(labels.cpu().numpy())\n    pd.DataFrame(result).to_csv(f'experiment1/result_{datetime.utcnow().microsecond}.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"modificated distributed loader with balance class:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_experiment2(device):\n    if not os.path.exists('experiment2'):\n        os.makedirs('experiment2')\n\n    train_sampler = DistributedSamplerWrapper(\n        sampler=BalanceClassSampler(labels=train_dataset.get_labels(), mode=\"downsampling\"),\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=16,\n        sampler=train_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=1\n    )\n    para_loader = pl.ParallelLoader(train_loader, [device])\n\n    result = {'toxic': []}\n    for labels in para_loader.per_device_loader(device):\n        result['toxic'].extend(labels.cpu().numpy())\n    pd.DataFrame(result).to_csv(f'experiment2/result_{datetime.utcnow().microsecond}.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    device = xm.xla_device()\n    run_experiment1(device)\n    run_experiment2(device)\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.concat([pd.read_csv(path) for path in glob('experiment1/*.csv')])\nsubmission['toxic'].hist()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.concat([pd.read_csv(path) for path in glob('experiment2/*.csv')])\nsubmission['toxic'].hist()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Thank you for reading my kernel!\n\nSo, I have demonstrated good technique for you, my friends! It will help you avoid OOM problems with training on full data with balance class even distributed training with TPU!\n\nIf you like this format of notebooks I would like continue to make kernels with realizations of my ideas."}],"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}