{"cells":[{"metadata":{},"cell_type":"markdown","source":"# L5Kit Rasterizer Performance Improvement\n\nThis improves the L5Kit rasterizer performance (by up to 40% single process) using Numba.\n\nProfiling L5Kit shows that the `l5kit.rasterizer.SemanticRasterizer.transform_points` is a major contributor to overall processing time. It takes ~30% of rasterizing time using the base settings (with larger raster sizes the time to load and process semantic map data makes up a smaller proportion so the improvement will be less). Re-implementing this function in Numba can improve performance significantly.","execution_count":null},{"metadata":{"collapsed":"true"},"cell_type":"markdown","source":"## Setup","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"Fork of l5kit that just refactors slightly to allow replacing a single function. Otherwise identical to the released l5kit v1.0.6.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip uninstall -y typing\n!pip install git+https://github.com/thomasbrandon/l5kit@v1.0.6-perf#subdirectory=l5kit\n!pip install omegaconf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from typing import List, Optional, Tuple\n\nfrom l5kit.data import DataManager, LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.rasterization import build_rasterizer, Rasterizer, SemanticRasterizer, BoxRasterizer, SemBoxRasterizer, StubRasterizer\nfrom l5kit.rasterization.rasterizer_builder import _load_metadata\n\nfrom omegaconf import OmegaConf\nimport numpy as np\nimport pandas as pd\nimport numba as nb\nfrom time import perf_counter\nfrom tqdm.auto import tqdm\nfrom contextlib import contextmanager\nimport importlib","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CONFIG_STR = \"\"\"\n# Config format schema number\nformat_version: 4\n\n###################\n## Model options\nmodel_params:\n  history_num_frames: 0\n  history_step_size: 1\n  history_delta_time: 0.1\n\n  future_num_frames: 50\n  future_step_size: 1\n  future_delta_time: 0.1\n\n###################\n## Input raster parameters\nraster_params:\n  # raster image size [pixels]\n  raster_size:\n    - 224\n    - 224\n  # raster's spatial resolution [meters per pixel]: the size in the real world one pixel corresponds to.\n  pixel_size:\n    - 0.5\n    - 0.5\n  # From 0 to 1 per axis, [0.5,0.5] would show the ego centered in the image.\n  ego_center:\n    - 0.25\n    - 0.5\n  map_type: \"py_semantic\"\n\n  # the keys are relative to the dataset environment variable\n  satellite_map_key: \"aerial_map/aerial_map.png\"\n  semantic_map_key: \"semantic_map/semantic_map.pb\"\n  dataset_meta_key: \"meta.json\"\n\n  # e.g. 0.0 include every obstacle, 0.5 show those obstacles with >0.5 probability of being\n  # one of the classes we care about (cars, bikes, peds, etc.), >=1.0 filter all other agents.\n  filter_agents_threshold: 0.5\n\n###################\n## Data loader options\ntrain_data_loader:\n  key: \"scenes/train.zarr\"\n\"\"\"\nCONFIG =  OmegaConf.create(CONFIG_STR)\n\nDATA_DIR=\"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(DATA_DIR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_dataset(cfg, zarr, map_type=\"SemBox\", raster_size=None, pixel_size=None, rast_class=None):\n    if not map_type is None: cfg.raster_params['map_type'] = map_type\n    if not raster_size is None: cfg.raster_params['raster_size'] = [raster_size, raster_size]\n    if not pixel_size is None: cfg.raster_params['pixel_size'] = [pixel_size, pixel_size]\n        \n    # Build rasterizer\n    raster_size = tuple(cfg.raster_params.raster_size)\n    pixel_size = np.array(cfg.raster_params.pixel_size)\n    ego_center = np.array(cfg.raster_params.ego_center)\n    semantic_map_filepath = dm.require(cfg.raster_params.semantic_map_key)\n    dataset_meta = _load_metadata(cfg.raster_params.dataset_meta_key, dm)\n    world_to_ecef = np.array(dataset_meta[\"world_to_ecef\"], dtype=np.float64)\n    filter_agents_threshold = cfg.raster_params.filter_agents_threshold\n    history_num_frames = cfg.model_params.history_num_frames\n    if map_type == \"SemBox\":\n        if rast_class is None: rast_class = SemBoxRasterizer\n        args = (raster_size, pixel_size, ego_center, filter_agents_threshold, history_num_frames, semantic_map_filepath, world_to_ecef,)\n    elif map_type == \"Semantic\":\n        if rast_class is None: rast_class = SemanticRasterizer\n        args = (raster_size, pixel_size, ego_center, semantic_map_filepath, world_to_ecef,)\n    elif map_type == \"Box\":\n        if rast_class is None: rast_class = BoxRasterizer\n        args = (raster_size, pixel_size, ego_center, filter_agents_threshold, history_num_frames)\n    else:\n        raise ValueError(\"Unknown rasterizer type: \" + map_type)\n    rast = rast_class(*args)\n    # Build dataset\n    ds = AgentDataset(cfg, zarr, rast)\n    return ds\n\ndef fmt_time(val):\n  units = ('p','n','µ','m','','k')\n  scale = int(np.floor(np.log10(np.abs(val))/3)) if val != 0 else 0\n  if scale < 2:\n    val *= 1000**(-scale)\n    unit = units[scale+4]\n    return f'{val:.1f}{unit}s'\n  return f\"{val:.1g}s\"\n\ndef summarise_times(times: np.array, iter_sec:bool=False):\n    summ = f\"{fmt_time(times.mean())} +/- {fmt_time(times.std())}\"\n    if iter_sec: summ += f\"; {len(times)/times.sum():.1f}it/s\"\n    return summ\n\ndef dataset_perf(ds, num=500, summary=False, report=True, progress=True):\n    n = len(ds)\n    times = [0] * num\n    it = range(num)\n    if progress: it = tqdm(it)\n    for i in it:\n        start = perf_counter()\n        # Don't just get sequential items to pull from different scenes\n        _ = ds[i * 9773 % n]\n        times[i] = perf_counter() - start\n    times = np.array(times)\n    summ = summarise_times(times, iter_sec=True)\n    if report:\n        print(f\"Retrieved {num} items in {fmt_time(times.sum())}.\\n  {summ}\")\n    if summary:\n        return summ\n    else:\n        return times","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_path = dm.require(CONFIG.train_data_loader.key)\ntrain_zarr = ChunkedDataset(data_path).open()\ntrain_ds = create_dataset(CONFIG, train_zarr, map_type=\"Semantic\")\nprint(train_ds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Numba Implementation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Transform an array of points with a coordinate system transformation matrix\n# For d dimensional points the transformation matrix should be of size d+1 (i.e. 2D points use a 3x3 matrix).\n@nb.guvectorize([(nb.float64[:,:], nb.float64[:,:], nb.float64[:,:])],\n                \"(p,d),(t,t)->(p,d)\", nopython=True)\ndef transform_points_nb(points, transf_matrix, res):\n    n_dim = transf_matrix.shape[0] - 1\n    #assert points.shape[1] == n_dim, \"Mismatched dimensions\"\n    # For each point compute a dot product with the transformation matrix fixing the Z coord to 1.\n    for p in range(points.shape[0]):\n        for out_dim in range(n_dim):\n            val = 0\n            for dim in range(n_dim):\n                val += points[p, dim] * transf_matrix[out_dim, dim]\n            val += transf_matrix[out_dim, n_dim] # *1 - Fixed Z\n            res[p,out_dim] = val","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Unfortunately I couldn't find a performant way of sharing code between the subpixel and non-subpixel variants.\nBest solutions were adding ~5-10% to execution time. Will investigate further. For now just duplicate.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"CV2_SHIFT_VALUE = 256\n# Transform an array of points with a coordinate system transformation matrix\n# For d dimensional points the transformation matrix should be of size d+1 (i.e. 2D points use a 3x3 matrix).\n@nb.guvectorize([(nb.float64[:,:], nb.float64[:,:], nb.int32[:,:])],\n                \"(p,d),(t,t)->(p,d)\", nopython=True)\ndef transform_points_subpixel_nb(points, transf_matrix, res):\n    n_dim = transf_matrix.shape[0] - 1\n    #assert points.shape[1] == n_dim, \"Mismatched dimensions\"\n    # For each point compute a dot product with the transformation matrix fixing the Z coord to 1.\n    for p in range(points.shape[0]):\n        for out_dim in range(n_dim):\n            val = 0\n            for dim in range(n_dim):\n                val += points[p, dim] * transf_matrix[out_dim, dim]\n            val += transf_matrix[out_dim, n_dim] # *1 - Fixed Z\n            res[p,out_dim] = int(val * CV2_SHIFT_VALUE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Verify code","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from l5kit.geometry import transform_points\nfrom l5kit.rasterization.semantic_rasterizer import transform_points_subpixel","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_item = train_ds[0]\ntransf_matrix = test_item[\"world_to_image\"]\ntransf_matrix","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lane_id = train_ds.rasterizer.bounds_info[\"lanes\"][\"ids\"][0]\nlane_coords = train_ds.rasterizer.proto_API.get_lane_coords(lane_id)\ntest_points = lane_coords['xyz_left'][:,:2]\ntest_points[:5]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.testing.assert_equal(transform_points(test_points, transf_matrix),\n                        transform_points_nb(test_points, transf_matrix))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.testing.assert_equal(transform_points_subpixel(test_points, transf_matrix),\n                        transform_points_subpixel_nb(test_points, transf_matrix))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"So seems to match output of original (only conversion of 2D points tested here as that's all the SemanticRasterizer uses, transform of 3D points not well tested but hopefully should work).","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Synthetic Benchmark","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def transform_perf(func, transf_matrix, num_items=100, num_repeats=1000, num_warmup=100, report=False):\n    inps = np.random.randn(num_repeats+num_warmup, num_items, 2)\n    times = [0] * inps.shape[0]\n    for i in range(inps.shape[0]):\n        it = inps[i, :, :]\n        start = perf_counter()\n        res = func(it, transf_matrix)\n        times[i] = perf_counter() - start\n    times = np.array(times[num_warmup:])\n    if report: print(summarise_times(times))\n    return times","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"times = transform_perf(transform_points, transf_matrix, report=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results = []\nfuncs = {\"transform_points\": (transform_points, transform_points_nb),\n         \"transform_points_subpixel\": (transform_points_subpixel, transform_points_subpixel_nb)}\nfor func_name, (orig_func,numba_func) in funcs.items():\n    print(\"Testing \" + func_name)\n    for num_items in (10, 1000, 10000, 100000):\n        orig_times = transform_perf(orig_func, transf_matrix, num_items)\n        numba_times = transform_perf(numba_func, transf_matrix, num_items)\n        results.append({\"Function\": func_name, \"Points\": num_items,\n                        \"Original\": summarise_times(orig_times),\n                        \"Numba\": summarise_times(numba_times),\n                        \"orig_mean\": orig_times.mean(), \"numba_mean\": numba_times.mean()})\nresults = pd.DataFrame(results)\nresults[\"Improvement\"] = (1-(results.numba_mean/results.orig_mean)).apply(\"{:.0%}\".format)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results[[\"Function\",\"Points\",\"Original\",\"Numba\",\"Improvement\"]]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset Benchmark","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"@contextmanager\ndef create_patched_dataset(*args, **kwargs):\n    mod = importlib.import_module(\"l5kit.rasterization.semantic_rasterizer\")\n    orig_func = mod.transform_points_subpixel\n    mod.transform_points_subpixel = transform_points_subpixel_nb\n    yield create_dataset(*args, **kwargs, rast_class=mod.SemanticRasterizer)\n    mod.transform_points_subpixel = orig_func","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"orig_times = dataset_perf(train_ds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with create_patched_dataset(CONFIG, train_zarr, map_type=\"Semantic\") as numba_ds:\n    numba_times = dataset_perf(numba_ds)\nprint(f\"\\n{1-(numba_times.mean() / orig_times.mean()):.1%} improvement\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python [conda env:.conda-lyft]","language":"python","name":"conda-env-.conda-lyft-py"},"language_info":{"name":"python","version":"3.7.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":4}