{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Lyft/Zarr cache, data sampler benchmark\n\n\n| Zarr cache | Sampler                   | max it/s | time /1000 iterations |\n| ---------- | ------------------------- | -------- | --------------------- |\n| Disabled   | RandomSampler             | 3.94     | 5:01                  |\n| Enabled    | RandomSampler             | 4.16     | 4:02                  |\n| Enabled    | LyftZarrCacheFixedSampler | 4.62     | 3:37                  |\n\n*SequentialSampler would solve the caching issue, but it would cause an overfit soon because of the similar consecutive images.*\n\n## Notes\nIf you want to get the most speed out of your dataloader, you'll need to adjust a few settings.\n\nYou should take a look at these:\n- **num_workers** At most the number of CPU cores you have\n- **cache_size_bytes** ChunkDataset.open(...) it should something like (I haven't tested): available RAM * 0.8 / num_workers\n- **chunk_size** LyftZarrCacheFixedSampler(....) it should be the max that fits in the cache\n"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import l5kit\nassert l5kit.__version__ == \"1.1.0\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport torch\nimport random\nimport plotly.express as px\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.sampler import Sampler\nfrom tqdm.notebook import tqdm\nfrom typing import Dict\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset\nfrom l5kit.rasterization import build_rasterizer\n\nfrom numcodecs import blosc\n\nblosc.set_nthreads(1)\nblosc.use_threads = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# os.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-full-training-set\"\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(None)\n\ncfg = {\n    'format_version': 4,\n    'model_params': {\n        'history_num_frames': 10,\n        'history_step_size': 1,\n        'history_delta_time': 0.1,\n        'future_num_frames': 50,\n        'future_step_size': 1,\n        'future_delta_time': 0.1,\n    },\n\n    'raster_params': {\n        'raster_size': [224, 224],\n        'pixel_size': [0.5, 0.5],\n        'ego_center': [0.25, 0.5],\n        'map_type': 'stub_debug',\n        # 'semantic_map_key': 'semantic_map/semantic_map.pb',\n        'dataset_meta_key': 'meta.json',\n        'filter_agents_threshold': 0.5,\n        'disable_traffic_light_faces': False,\n    },\n\n    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 4\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TqdmExtra(tqdm):\n    \n    @staticmethod\n    def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False,\n                     unit='it', unit_scale=False, rate=None, bar_format=None,\n                     postfix=None, unit_divisor=1000, initial=0, colour=None,\n                     **extra_kwargs):\n        r = tqdm.format_meter(n, total, elapsed, ncols=ncols, prefix=prefix, ascii=ascii,\n                              unit=unit, unit_scale=unit_scale, rate=rate, bar_format=bar_format,\n                              postfix=postfix, unit_divisor=unit_divisor, initial=initial,\n                              colour=colour, **extra_kwargs)\n        \n        global itps\n        itps.append(n/elapsed)\n        \n        return r","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Without zarr cache"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_cfg = cfg[\"train_data_loader\"]\nds_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open(cached=False)\n\nrasterizer = build_rasterizer(cfg, dm)\n\ntrain_dataset = AgentDataset(cfg, ds_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                             shuffle=train_cfg[\"shuffle\"],\n                             batch_size=train_cfg[\"batch_size\"],\n                             num_workers=train_cfg[\"num_workers\"])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"itps = []\nprogress_bar = TqdmExtra(train_dataloader, total=1000)\n\nfor i, data in enumerate(progress_bar):\n\n    if i == 1000:\n        break\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = px.line(x=range(len(itps)), y=itps, title='it/s without zarr cache')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del ds_zarr\ndel rasterizer\ndel train_dataset\ndel train_dataloader\ndel progress_bar\ndel data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# With cache (random sampler)"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_cfg = cfg[\"train_data_loader\"]\nds_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open(cached=True, cache_size_bytes=int(5e9))\n\nrasterizer = build_rasterizer(cfg, dm)\n\ntrain_dataset = AgentDataset(cfg, ds_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                             shuffle=train_cfg[\"shuffle\"],\n                             batch_size=train_cfg[\"batch_size\"],\n                             num_workers=train_cfg[\"num_workers\"])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"itps = []\nprogress_bar = TqdmExtra(train_dataloader, total=1000)\n\nfor i, data in enumerate(progress_bar):\n\n    if i == 1000:\n        break\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = px.line(x=range(len(itps)), y=itps, title='it/s with zarr cahce + random sampler')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del progress_bar\ndel train_dataloader\ndel data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# With cache (custom sampler)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftZarrCacheFixedSampler(Sampler):\n\n    def __init__(self, datasource: AgentDataset, chunk_size=1000000):\n        super(LyftZarrCacheFixedSampler, self).__init__(datasource)\n\n        self.chunk_size = chunk_size\n        self.datasource = datasource\n        self.datasource.agents_indices.sort(kind='stable')\n\n        self.epoch = 0\n        self.n_chunks = len(self.datasource.agents_indices) // chunk_size\n        self.n_last_chunk = len(self.datasource.agents_indices) % chunk_size\n\n        if self.n_last_chunk > 0:\n            self.n_chunks += 1\n            \n        print(f\"Number of chunks: {self.n_chunks}\")\n        print(f\"Number of agents: {len(self.datasource.agents_indices)}\")\n\n    def __len__(self) -> int:\n        return len(self.datasource.agents_indices) // 4\n\n    def __iter__(self):\n        indices = np.array([x for x in range(len(self.datasource.agents_indices))])\n\n        res_idx = []\n\n        for chunk in range(self.n_chunks):\n            from_idx = self.chunk_size * chunk\n            to_idx = min([self.chunk_size * (chunk + 1), len(indices)])\n\n            x = indices[from_idx:to_idx]\n            x = x[self.epoch::4]\n\n            np.random.shuffle(x)\n            res_idx.append(x)\n\n        self.epoch += 1\n        if self.epoch == 4:\n            self.epoch = 0\n\n        random.shuffle(res_idx)\n        indices = np.hstack(res_idx)\n\n        return iter(indices)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset,\n                             sampler=LyftZarrCacheFixedSampler(train_dataset, chunk_size=500000),\n                             batch_size=train_cfg[\"batch_size\"],\n                             num_workers=train_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"itps = []\nprogress_bar = TqdmExtra(train_dataloader, total=1000)\n\nfor i, data in enumerate(progress_bar):\n\n    if i == 1000:\n        break\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = px.line(x=range(len(itps)), y=itps, title='it/s with zarr cache + custom sampler')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}