{"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":"Dealing with large chunked data is one of the challenges of this competition, which may impose an additional obstacle for people willing to join. Since the competition was not very active recently, I decided to share quite a simple way of dealing with the provided data without the time-consuming conversion of data into other formats or trying to load everything into the RAM, which is quite challenging, especially at Kaggle.\n\nThe chunked data format is not very friendly to Pytorch, and some additional tricks are needed. My approach consists of two things: (1) A **dataloader caching the considered chunks** and (2) **Random Chunked Sampler**, which randomly selects a chunk and then goes through all ids in the chunk before selecting a new one. With this sampler, the caching pipeline spends time reading the corresponding chunk only at the first request, and the following requests are finished quickly with reading data from RAM. Meanwhile, when the cache is full, the earliest record is removed, which limits RAM usage. At inference and test, when reading is sequential, no extra modifications are needed to the data sampler. \nIt would be great if the Pytorch team implemented native support of chunked data.","metadata":{}},{"cell_type":"code","source":"#install torch_geometric\n#change to gpu version in gpu kernel\n!pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-1.13.0+cpu.html\n!pip install torch_geometric==2.2\n!pip install torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-1.13.0+cpu.html","metadata":{"scrolled":true,"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-03-30T14:40:46.115794Z","iopub.execute_input":"2023-03-30T14:40:46.116497Z","iopub.status.idle":"2023-03-30T14:41:31.840280Z","shell.execute_reply.started":"2023-03-30T14:40:46.116447Z","shell.execute_reply":"2023-03-30T14:41:31.838453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import polars as pl\nimport pandas as pd\nimport gc, os, random, math\nimport numpy as np\nfrom tqdm.notebook import tqdm\nfrom collections import OrderedDict\nfrom bisect import bisect_right\n\nimport torch\nfrom typing import Iterator, Iterable, Optional, Sequence, List, TypeVar, Generic, Sized, Union\n\nfrom torch_geometric.data import Data, Dataset\nfrom torch_geometric.loader import DataLoader","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-03-30T14:41:39.584925Z","iopub.execute_input":"2023-03-30T14:41:39.585442Z","iopub.status.idle":"2023-03-30T14:41:56.605834Z","shell.execute_reply.started":"2023-03-30T14:41:39.585389Z","shell.execute_reply":"2023-03-30T14:41:56.604357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = '/kaggle/input/icecube-neutrinos-in-deep-ice/'\nMETA = '/kaggle/input/train-meta-parquet/'\n\n\nSEED = 2023\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n\nseed_everything(SEED)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-03-30T14:42:01.571184Z","iopub.execute_input":"2023-03-30T14:42:01.571636Z","iopub.status.idle":"2023-03-30T14:42:01.581745Z","shell.execute_reply.started":"2023-03-30T14:42:01.571596Z","shell.execute_reply":"2023-03-30T14:42:01.580223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here is an example of a dataset with caching. I use torch_geometric for illustration purposes, and you can easily modify it according to the data format used in your model. Be mindful of the number of workers and cache_size (to my understanding each worker has its own cache).","metadata":{}},{"cell_type":"code","source":"class IceCubeCache(Dataset):\n    def __init__(self, path=PATH, mode='test', meta=None, pulse_limit=128, cache_size=1):\n        val_fnames = ['batch_655.parquet','batch_656.parquet','batch_657.parquet',\n                      'batch_658.parquet','batch_659.parquet']\n        chunk_size=200000\n        self.mode, self.path_meta = mode,meta\n\n        if mode == 'train' or mode == 'eval':\n            assert meta is not None, 'Need to provide labels'\n            self.path = os.path.join(path,'train')\n            self.files = [p for p in sorted(os.listdir(self.path)) \\\n                          if p!='batch_660.parquet'] #660 is shorter\n            if mode == 'train':\n                self.files = sorted(set(self.files) - set(val_fnames))\n            else: self.files = val_fnames\n            self.chunks = [chunk_size]*len(self.files)\n        elif mode == 'test':\n            self.path = os.path.join(path,'test')\n            self.files = [p for p in sorted(os.listdir(self.path))]\n            \n            #make sure that all files are considered regardless the number of events\n            self.chunks = []\n            for fname in self.files:\n                ids = pl.read_parquet(os.path.join(self.path,fname)\\\n                        ).select(['event_id']).unique().to_numpy().reshape(-1)\n                self.chunks.append(len(ids))\n            gc.collect()\n        else: raise NotImplementedError \n            \n        self.chunk_cumsum = np.cumsum(self.chunks)\n        self.cache, self.meta = None,None\n        self.pulse_limit,self.cache_size = pulse_limit,cache_size\n        self.geometry = pd.read_csv(os.path.join(path,'sensor_geometry.csv'))\n        self.geometry = (self.geometry[['x','y','z']].values/500.0).astype(np.float32)\n        \n    def __len__(self):\n        return self.chunk_cumsum[-1]\n    \n    def load_data(self, fname):\n        if self.cache is None: self.cache = OrderedDict()\n        if fname not in self.cache:\n            df = pl.read_parquet(os.path.join(self.path,fname))\n            df = df.groupby(\"event_id\").agg([\n                pl.count(),\n                pl.col(\"sensor_id\").list(),\n                pl.col(\"time\").list(),\n                pl.col(\"charge\").list(),\n                pl.col(\"auxiliary\").list(),])\n            self.cache[fname] = df.sort('event_id')\n            if len(self.cache) > self.cache_size: del self.cache[list(self.cache.keys())[0]]\n                \n    def load_meta(self, fname):\n        if self.meta is None: self.meta = OrderedDict()\n        if fname not in self.meta:\n            fidx = fname.split('.')[0].split('_')[-1]\n            self.meta[fname] = pl.read_parquet(os.path.join(self.path_meta,\n                                f'train_meta_{fidx}.parquet')).sort('event_id')\n            if len(self.meta) > self.cache_size: del self.meta[list(self.meta.keys())[0]]\n        \n    def __getitem__(self, idx0):\n        fidx = bisect_right(self.chunk_cumsum, idx0)\n        fname = self.files[fidx]\n        idx = int(idx0 - self.chunk_cumsum[fidx-1]) if fidx > 0 else idx0\n        \n        self.load_data(fname)\n        df = self.cache[fname][idx]\n        sensor_id =  df['sensor_id'][0].item().to_numpy()\n        time =  df['time'][0].item().to_numpy()\n        charge = df['charge'][0].item().to_numpy()\n        auxiliary = df['auxiliary'][0].item().to_numpy()\n        \n        pos = self.geometry[sensor_id]\n        time = (time - 1e4)/3e4\n        charge = np.log10(charge)/3.0\n        auxiliary = auxiliary - 0.5\n        \n        x = np.stack([pos[:,0], pos[:,1], pos[:,2], time, charge, auxiliary],\n                     -1).astype(np.float32)\n        x =  torch.from_numpy(x)\n        data = Data(x=x, n_pulses=torch.tensor(x.shape[0], dtype=torch.int32))\n        \n        # Downsample large events\n        if data.n_pulses > self.pulse_limit:\n            data.x = data.x[torch.randperm(len(data.x)).numpy()[:self.pulse_limit]]\n            data.n_pulses = torch.tensor(self.pulse_limit, dtype=torch.int32)\n        \n        if self.mode != 'test': \n            self.load_meta(fname)\n            meta = self.meta[fname][idx]\n            azimuth = meta['azimuth'].item()\n            zenith = meta['zenith'].item()\n            target = np.array([azimuth,zenith]).astype(np.float32)\n            target = torch.from_numpy(target)\n        else: target = df['event_id'].item()\n            \n        return data, target","metadata":{"execution":{"iopub.status.busy":"2023-03-30T14:42:03.183429Z","iopub.execute_input":"2023-03-30T14:42:03.183855Z","iopub.status.idle":"2023-03-30T14:42:03.216681Z","shell.execute_reply.started":"2023-03-30T14:42:03.183815Z","shell.execute_reply":"2023-03-30T14:42:03.215173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This class performs chunked-based random sampling: selection of a chunk at random and then sampling all ids within the chunk (so the cache may be used effectively). It is nearly completely borrowed from the Pytotch source code with adding a special consideration for chunks.","metadata":{}},{"cell_type":"code","source":"class RandomChunkSampler(torch.utils.data.Sampler[int]):\n    data_source: Sized\n    replacement: bool\n\n    def __init__(self, data_source: Sized, chunks, num_samples: Optional[int] = None,\n                 generator=None, **kwargs) -> None:\n        # chunks - a list of chunk sizes\n        self.data_source = data_source\n        self._num_samples = num_samples\n        self.generator = generator\n        self.chunks = chunks\n\n        if not isinstance(self.num_samples, int) or self.num_samples <= 0:\n            raise ValueError(\"num_samples should be a positive integer \"\n                             \"value, but got num_samples={}\".format(self.num_samples))\n\n    @property\n    def num_samples(self) -> int:\n        # dataset size might change at runtime\n        if self._num_samples is None:\n            return len(self.data_source)\n        return self._num_samples\n\n    def __iter__(self) -> Iterator[int]:\n        n = len(self.data_source)\n        cumsum = np.cumsum(self.chunks)\n        if self.generator is None:\n            seed = int(torch.empty((), dtype=torch.int64).random_().item())\n            generator = torch.Generator()\n            generator.manual_seed(seed)\n        else:\n            generator = self.generator\n\n        chunk_list = torch.randperm(len(self.chunks), generator=generator).tolist()\n        # sample indexes chunk by chunk\n        for i in chunk_list:\n            chunk_len = self.chunks[i]\n            offset = cumsum[i-1] if i > 0 else 0\n            yield from (offset + torch.randperm(chunk_len, generator=generator)).tolist()\n\n    def __len__(self) -> int:\n        return self.num_samples","metadata":{"execution":{"iopub.status.busy":"2023-03-30T14:42:14.920697Z","iopub.execute_input":"2023-03-30T14:42:14.921277Z","iopub.status.idle":"2023-03-30T14:42:14.936970Z","shell.execute_reply.started":"2023-03-30T14:42:14.921219Z","shell.execute_reply":"2023-03-30T14:42:14.935690Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Three examples below illustrate how to use the code for building train, evaluation, and test dataloaders. I use \"num_workers=0\" because at Kaggle the dataloader freezes if workers are created. At my home computer, 1000 batches (256,000 events > 1 data chunk) are sampled approximately within 20 seconds if 4 workers are used.","metadata":{}},{"cell_type":"code","source":"ds_train = IceCubeCache(mode='train',meta=META)\nsampler = RandomChunkSampler(ds_train, chunks=ds_train.chunks)\n#somehow num_workers > 0 freezes at kaggle\ndl_train = DataLoader(ds_train, batch_size=256, sampler=sampler, drop_last=True, num_workers=0)\n\nfor i,(x,y) in enumerate(tqdm(dl_train, total=1000)):\n    if i >= 1000: break","metadata":{"execution":{"iopub.status.busy":"2023-03-30T14:42:16.971571Z","iopub.execute_input":"2023-03-30T14:42:16.972042Z","iopub.status.idle":"2023-03-30T14:42:32.954045Z","shell.execute_reply.started":"2023-03-30T14:42:16.971997Z","shell.execute_reply":"2023-03-30T14:42:32.951913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_val = IceCubeCache(mode='eval',meta=META)\ndl_val = DataLoader(ds_val, batch_size=256, shuffle=False, drop_last=False, num_workers=0)\nfor i,(x,y) in enumerate(tqdm(dl_val, total=1000)):\n    if i >= 1000: break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_test = IceCubeCache(mode='test')\n#use num_workers=0 at inference\ndl_test = DataLoader(ds_test, batch_size=256, shuffle=False, drop_last=False, num_workers=0)\nfor x,y in tqdm(dl_test):\n    pass","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}