{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Custom C++ and CUDA Extensions with PyTorch\n\nIn this notebook, I'd like to show you how to write custom C++ extensions. I'll show you a simple \"Hello World\" example, and my GPU/CUDA Rasterizer solution for the Lyft Competition.\n\n\n1. [Hello World!](#hello_world)\n1. [Lyft GPU Rasterizer](#cpp_cuda_codes)\n1. [Benchmarks](#benchmarks)\n\n**Tools you'll need for building and using a custom C++ code in Python:**\n\n- **pybind11** is a lightweight header-only library that exposes C++ types in Python and vice versa, mainly to create Python bindings of existing C++ code. [More details here](https://github.com/pybind/pybind11)\n- **Ninja** is a small build system with a focus on speed. It differs from other build systems in two major respects: it is designed to have its input files generated by a higher-level build system, and it is designed to run builds as fast as possible. [More details here](https://ninja-build.org/)\n- **GCC** GNU Compiler Collection (v7.5.0) or an alternative C++ compiler\n- **CUDA driver + CUDA Toolkit**\n- **PyTorch**\n\nFortunately, PyTorch has lots of useful utility tools, so we don't have to deal with any of the libraries mentioned above directly.\n"},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# l5kit dependencies\n!pip install pymap3d==2.1.0\n!pip install protobuf==3.12.2\n!pip install transforms3d\n!pip install zarr\n!pip install ptable\n\n!pip uninstall -y typing\n\n# The modified l5kit from my github repo (gpu branch)\n!pip install --no-dependencies git+https://github.com/pestipeti/lyft-l5kit.git@gpu#subdirectory=l5kit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import cv2\nimport math\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport torch\nimport l5kit\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.resnet import resnet50\nfrom tqdm.notebook import tqdm\nfrom typing import Dict\n\nfrom torch.utils.cpp_extension import load, load_inline, is_ninja_available\nfrom torch.utils.data._utils.collate import default_collate\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.data.map_api import MapAPI\nfrom l5kit.data.filter import (filter_agents_by_labels,\n                               filter_agents_by_track_id,\n                               filter_tl_faces_by_status)\nfrom l5kit.dataset import AgentDataset\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.geometry import transform_points, rotation33_as_yaw, transform_point\nfrom l5kit.rasterization import build_rasterizer, Rasterizer, RenderContext\nfrom l5kit.rasterization.box_rasterizer import get_ego_as_agent\nfrom l5kit.rasterization.semantic_rasterizer import elements_within_bounds\nfrom l5kit.rasterization.rasterizer_builder import (_load_metadata,\n                                                    get_hardcoded_world_to_ecef)\n\nfrom typing import List, Optional, Tuple, Union","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"print(f\"PyTorch version: {torch.__version__}\")\nprint(f\"Is Ninja build available: {is_ninja_available()}\")\nprint(f\"L5Kit version: {l5kit.__version__}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<a id=\"hellow_world\"></a>\n# Hello World!\n\n\n- **[torch.utils.cpp_extension.load](https://pytorch.org/docs/stable/_modules/torch/utils/cpp_extension.html#load)** Loads a PyTorch C++ extension just-in-time (JIT).\n- **[torch.utils.cpp_extension.load_inline](https://pytorch.org/docs/stable/_modules/torch/utils/cpp_extension.html#load_inline)** Loads a PyTorch C++ extension just-in-time (JIT) from string sources.\n\n\nWith these two methods, you can load and build your C++ code. Under the hood, PyTorch will create a Ninja build file and it will call Ninja to compile your source into a dynamic library.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"cpp_source = \"\"\"\nstd::string hello_world() {\n  return \"Hello World!\";\n}\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"my_module = load_inline(\n    name='my_module',\n    cpp_sources=[cpp_source],\n    functions=['hello_world'],\n    build_directory='/kaggle/working',\n    verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Final C++ source\nWhen you execute the cell above (`load` or `load_inline` calls) PyTorch will add some extra C++ code, and the final source will look something like this\n\n```c++\n#include <torch/extension.h>\n\nstd::string hello_world() {\n  return \"Hello World!\";\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\nm.def(\"hello_world\", torch::wrap_pybind_function(hello_world), \"hello_world\");\n}\n```\n\n### Ninja Build File\nPyTorch also will generate the Ninja Build File\n\n```\nninja_required_version = 1.3\ncxx = c++\n\ncflags = -DTORCH_EXTENSION_NAME=my_module -DTORCH_API_INCLUDE_EXTENSION_H\\\n         -isystem ... -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14\npost_cflags = \nldflags = -shared -L/opt/conda/lib/python3.7/site-packages/torch/lib\\\n          -lc10 -ltorch_cpu -ltorch -ltorch_python\n\nrule compile\n  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags\n  depfile = $out.d\n  deps = gcc\n\nrule link\n  command = $cxx $in $ldflags -o $out\n\nbuild main.o: compile /kaggle/working/main.cpp\n\nbuild my_module.so: link main.o\n\ndefault my_module.so\n\n```\n\n### Compiling errors\n\nPyTorch does not forward the console's output, so if something goes wrong, you'll only get a \"Ninja compile failed\" general error message. For more details, you have to run ninja compile directly.\n\n```\n$ ninja -v -C [/dir/to/ninja.build]\n```"},{"metadata":{"trusted":true},"cell_type":"code","source":"my_module.hello_world()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<a id=\"cpp_cuda_codes\"></a>\n# C++ and CUDA codes\n\n### C++ Wrapper\nThis is just a simple C++ wrapper function"},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"lyft_gpu_cpp_src = \"\"\"\n#include <torch/extension.h>\n#include <vector>\n#include <iostream>\n\nusing namespace std;\n\n/**\n * CUDA rasterize\n *\n * @param target torch.cuda.FloatTensor, shape: BxCxHxW (Batch, channels, height, width)\n * @param lines torch.cuda.FloatTensor List of lines to be drawn, shape: Nx7\n * @param agents torch.cuda.FloatTensor List of agents to be drawn, shape: Nx9\n * @param w2i torch.cuda.FloatTensor World to Image transformation matrices, shape: Bx3x3\n *\n * @return void\n */\nvoid cuda_rasterize(torch::Tensor &target, torch::Tensor &lines, torch::Tensor &agents, torch::Tensor &w2i);\n\n// Basic checks\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)\n\nvoid rasterize(torch::Tensor &target, torch::Tensor &lines, torch::Tensor &agents, torch::Tensor &w2i) {\n\n    // Run some checks.\n    CHECK_INPUT(target);\n    CHECK_INPUT(agents);\n    CHECK_INPUT(lines);\n    CHECK_INPUT(w2i);\n\n    cuda_rasterize(target, lines, agents, w2i);\n}\n\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"*NOTE: The implementation is an early version. It is far from optimal, and it is not production-ready. I did not have enought time to clean up the code. Maybe later I will.*\n\n### CUDA Rasterizer\nYou can find the CUDA rasterizer code in the cell below. You pass the necessary data, and it will draw the lines and the agents to the target. \n\n*The transformations for drawing the lines is calculated on CPU (It was fast enough this way, so I did not finished)*\n\n#### `target` torch.cuda.FloatTensor (BxCxHxW)\nThis is the target image (batch of images); before calling the rasterizer, you should clear the previous drawings.\nDimensions: batch, channels, height, width\n\n#### `lines` torch.cuda.FloatTensor (Nx7)\nThis tensor contains the data about all of the lines you'd like to draw.\n- N - Number of lines\n- [:, 0] - batch index. When the DataLoader loads the data (in parallel), this is unknown. Later, the `collate_fn` will fill this (see the code).\n- [:, 1:5] - x1, y1, x2, y2 Line coordinates\n- [:, 5] - Channel number\n- [:, 6] - Channel value (0..255; technically it can be any number)\n\n#### `agents` torch.cuda.FloatTensor (Nx9)\n- N - Number of agents (on all batches)\n- [:, 0] - batch index\n- [:, 1:3] - Centroix x, y\n- [:, 3:6] - Extent (length, width, height)\n- [:, 6] - Yaw\n- [:, 7] - Channel number\n- [:, 8] - Channel value\n\n#### `w2i` torch.cuda.FloatTensor (Bx3x3)\nWorld to image (raster_from_world) transformation matrices"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"lyft_gpu_cuda_src = \"\"\"\n#include <torch/extension.h>\n\n#include <stdio.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n\nusing namespace std;\nusing namespace torch::indexing;\n\nnamespace {\n\ntemplate <typename scalar_t>\n__global__ void cuda_draw_lines_kernel(\n    torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> target,\n    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> lines,\n    size_t lines_size) {\n\n    int batch_size = target.size(0);\n    int channels = target.size(1);\n    int height = target.size(2);\n    int width = target.size(3);\n\n    int index = blockIdx.x * blockDim.x + threadIdx.x;\n    int stride = blockDim.x * gridDim.x;\n\n    int dx, dy, dx1, dy1, px, py, xe, ye, x, y;\n    int bx, x1, x2, y1, y2, ch, vl;\n\n    for (int i = index; i < lines_size; i += stride) {\n\n        bx = lines[i][0]; // batch idx\n        x1 = lines[i][1]; // x1\n        y1 = lines[i][2]; // y1\n        x2 = lines[i][3]; // x2\n        y2 = lines[i][4]; // y2\n        ch = lines[i][5]; // channel idx (0-2: RGB)\n        vl = lines[i][6]; // pixel value (0..255)\n\n        /*\n            The line drawing algorithm is from this post:\n            https://jstutorial.medium.com/how-to-code-your-first-algorithm-draw-a-line-ca121f9a1395\n        */\n\n        dx = x2 - x1;\n        dy = y2 - y1;\n\n        dx1 = std::abs(dx);\n        dy1 = std::abs(dy);\n\n        px = 2 * dy1 - dx1;\n        py = 2 * dx1 - dy1;\n\n        // The line is X-axis dominant\n        if (dy1 <= dx1) {\n\n            if (dx >= 0) {\n                // Line is drawn left to right\n                x = x1; y = y1; xe = x2;\n            } else {\n                // Line is drawn right to left (swap ends)\n                x = x2; y = y2; xe = x1;\n            }\n\n            if (x >= 0 && x < width && y >= 0 && y < height) {\n                target[bx][ch][y][x] = vl;\n            }\n\n            while (x < xe) {\n                x = x + 1;\n\n                if (px < 0) {\n                    px = px + 2 * dy1;\n                } else {\n                    if ((dx < 0 && dy < 0) || (dx > 0 && dy > 0)) {\n                        y = y + 1;\n                    } else {\n                        y = y - 1;\n                    }\n                    px = px + 2 * (dy1 - dx1);\n                }\n\n                if (x >= 0 && x < width && y >= 0 && y < height) {\n                    target[bx][ch][y][x] = vl;\n                }\n            }\n\n        } else {\n\n            if (dy >= 0) {\n                // Line is drawn bottom to top\n                x = x1; y = y1; ye = y2;\n            } else {\n                // Line is drawn top to bottom\n                x = x2; y = y2; ye = y1;\n            }\n\n            if (x >= 0 && x < width && y >= 0 && y < height) {\n                target[bx][ch][y][x] = vl;\n            }\n\n            while(y < ye) {\n                y = y + 1;\n\n                if (py <= 0) {\n                    py = py + 2 * dx1;\n                } else {\n                    if ((dx < 0 && dy<0) || (dx > 0 && dy > 0)) {\n                        x = x + 1;\n                    } else {\n                        x = x - 1;\n                    }\n                    py = py + 2 * (dx1 - dy1);\n                }\n\n                if (x >= 0 && x < width && y >= 0 && y < height) {\n                    target[bx][ch][y][x] = vl;\n                }\n\n            }\n        }//_else\n\n    }//_for\n\n}\n\ntemplate <typename scalar_t>\n__global__ void cuda_generate_box_coords(\n    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> agents,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> box_coords,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> transformations,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> centroids,\n    size_t agents_size) {\n    \n    int index = blockIdx.x * blockDim.x + threadIdx.x;\n    int stride = blockDim.x * gridDim.x;\n    \n    float sy, cy, extent_l, extent_w;\n    \n    for (int i = index; i < agents_size; i += stride) {\n        \n        extent_l = agents[i][3] / 2;\n        extent_w = agents[i][4] / 2;\n        \n        // Top-left (p1)\n        box_coords[i][0][0] = -extent_l;\n        box_coords[i][1][0] = -extent_w;\n        box_coords[i][2][0] = 0.0;\n        // Bottom-left (p2)\n        box_coords[i][0][1] = -extent_l;\n        box_coords[i][1][1] = extent_w;\n        box_coords[i][2][1] = 0.0;\n        // Bottom-right (p3)\n        box_coords[i][0][2] = extent_l;\n        box_coords[i][1][2] = extent_w;\n        box_coords[i][2][2] = 0.0;\n        // Top-right (p4)\n        box_coords[i][0][3] = extent_l;\n        box_coords[i][1][3] = -extent_w;\n        box_coords[i][2][3] = 0.0;\n        \n        // agnet's yaw\n        cy = cos(agents[i][6]);\n        sy = sin(agents[i][6]);        \n        \n        transformations[i][0][0] = cy;\n        transformations[i][0][1] = -sy;\n        transformations[i][1][0] = sy;\n        transformations[i][1][1] = cy;\n        transformations[i][2][2] = 1.0;\n\n        // agent's centroid xy (world space)\n        centroids[i][0][0] = agents[i][1];\n        centroids[i][1][0] = agents[i][2];\n        centroids[i][2][0] = 1.0;\n        centroids[i][0][1] = agents[i][1];\n        centroids[i][1][1] = agents[i][2];\n        centroids[i][2][1] = 1.0;\n        centroids[i][0][2] = agents[i][1];\n        centroids[i][1][2] = agents[i][2];\n        centroids[i][2][2] = 1.0;\n        centroids[i][0][3] = agents[i][1];\n        centroids[i][1][3] = agents[i][2];\n        centroids[i][2][3] = 1.0;\n\n    }\n\n}\n    \ntemplate <typename scalar_t>\n__global__ void cuda_world_to_image_transform(\n    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> agents,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> w2i,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> transformations,\n    size_t agents_size){\n \n    int index = blockIdx.x * blockDim.x + threadIdx.x;\n    int stride = blockDim.x * gridDim.x;\n    int bidx = 0;  // batch idx\n\n    for (int i = index; i < agents_size; i += stride) {\n        bidx = agents[i][0];\n\n        transformations[i][0][0] = w2i[bidx][0][0];\n        transformations[i][0][1] = w2i[bidx][0][1];\n        transformations[i][0][2] = w2i[bidx][0][2];\n        transformations[i][1][0] = w2i[bidx][1][0];\n        transformations[i][1][1] = w2i[bidx][1][1];\n        transformations[i][1][2] = w2i[bidx][1][2];\n        transformations[i][2][0] = w2i[bidx][2][0];\n        transformations[i][2][1] = w2i[bidx][2][1];\n        transformations[i][2][2] = w2i[bidx][2][2];\n    }\n\n}\n    \ntemplate <typename scalar_t>\n__global__ void cuda_generate_rectangles(\n    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> agents,\n    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> box_coords,\n    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> lines,\n    size_t agents_size) {\n\n    int index = blockIdx.x * blockDim.x + threadIdx.x;\n    int stride = blockDim.x * gridDim.x;\n    \n    int bi, ch, vl, p1, p2;\n\n    for (int i = index; i < agents_size; i += stride) {\n        bi = agents[i][0];\n        ch = agents[i][7];\n        vl = agents[i][8];\n        \n        for (int j = 0; j < 4; j++) {\n            \n            p1 = j;\n            p2 = j + 1;\n            \n            if (p2 > 3) {\n                p2 = 0;\n            }\n            \n            lines[i * 4 + j][0] = bi;\n            lines[i * 4 + j][1] = (int)box_coords[i][0][p1];\n            lines[i * 4 + j][2] = (int)box_coords[i][1][p1];\n            lines[i * 4 + j][3] = (int)box_coords[i][0][p2];\n            lines[i * 4 + j][4] = (int)box_coords[i][1][p2];\n            lines[i * 4 + j][5] = ch;  // channel\n            lines[i * 4 + j][6] = vl;  // color\n        }\n        \n    }\n};\n    \n\n}//_namespace\n\n// CUDA declarations\nvoid cuda_rasterize(torch::Tensor &target, torch::Tensor &lines, torch::Tensor &agents, torch::Tensor &w2i) {\n\n    const int blocks_size = 1024;\n    const int line_blocks = (lines.size(0) + blocks_size - 1) / blocks_size;\n    const int lines_size = lines.size(0);\n\n    // 1, Draw lines\n    AT_DISPATCH_FLOATING_TYPES(target.type(), \"cuda_draw_lines\", ([&] {\n        cuda_draw_lines_kernel<scalar_t><<<line_blocks, blocks_size>>>(\n            target.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),\n            lines.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            lines_size\n            );\n      }));\n\n    const int agents_size = agents.size(0);\n    const int agent_blocks = (agents_size + blocks_size - 1) / blocks_size;\n\n    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0).requires_grad(false);\n\n    // 2, Convert box coordinates\n\n    // 2.1. - Generate corners\n    torch::Tensor boxes = torch::zeros({agents_size, 3, 4}, options);\n    torch::Tensor transformations = torch::zeros({agents_size, 3, 3}, options);\n    torch::Tensor centroids = torch::zeros({agents_size, 3, 4}, options);\n    torch::Tensor borders = torch::zeros({agents_size * 4, 7}, options);\n\n    AT_DISPATCH_FLOATING_TYPES(target.type(), \"cuda_generate_box_coords\", ([&] {\n        cuda_generate_box_coords<scalar_t><<<agent_blocks, blocks_size>>>(\n            agents.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            boxes.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            transformations.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            centroids.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            agents_size\n        );\n    }));\n\n    // 2.2. - Transform boxes (rotate by agent's yaw)\n    boxes = torch::matmul(transformations, boxes);\n\n    // 2.3. - Move agents to its position in the world space.\n    boxes = torch::add(boxes, centroids);\n\n    // 2.4. - Transform from world to image coordinates\n    transformations *= 0.0;\n    AT_DISPATCH_FLOATING_TYPES(target.type(), \"cuda_world_to_image_transform\", ([&] {\n        cuda_world_to_image_transform<scalar_t><<<agent_blocks, blocks_size>>>(\n            agents.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            w2i.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            transformations.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            agents_size\n        );\n    }));\n\n    boxes = torch::matmul(transformations, boxes);\n\n    // 2.5. - Generate border lines.\n    borders.index_put_({Slice(None, None), 0}, -1);\n    AT_DISPATCH_FLOATING_TYPES(target.type(), \"cuda_generate_rectangles\", ([&] {\n        cuda_generate_rectangles<scalar_t><<<agent_blocks, blocks_size>>>(\n            agents.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            boxes.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),\n            borders.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            agents_size\n        );\n    }));\n\n    // 2.6. - Draw box border\n    AT_DISPATCH_FLOATING_TYPES(target.type(), \"cuda_draw_lines\", ([&] {\n        cuda_draw_lines_kernel<scalar_t><<<line_blocks, blocks_size>>>(\n            target.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),\n            borders.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),\n            borders.size(0)\n            );\n    }));\n\n    // 2.7. - Fill boxes.\n    cudaDeviceSynchronize();\n}\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lyft_gpu = load_inline(\n    name='lyft_gpu_rasterizer',\n    cpp_sources=[lyft_gpu_cpp_src],\n    cuda_sources=[lyft_gpu_cuda_src],\n    functions=['rasterize'],\n    build_directory='/kaggle/working',\n    with_cuda=True,\n    verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Rasterizer (l5kit extension)\n\nThe code below is more or less a copy-pasted version of the l5kit rasterizers (+ the modification for the GPU).\nThere are a few changes in the l5kit as well (see my github repo - gpu branch)"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# Lane-line colors\nDIVIDER_TYPE_CODES = [128, 128, 16, 40, 64, 88, 112, 136, 160, 184, 208, 232, 255]\n\nclass GPURasterizer(Rasterizer):\n\n    def __init__(self, raster_size: Tuple[int, int], pixel_size: Union[np.ndarray, list, float],\n                 ego_center: np.ndarray, filter_agents_threshold: float, history_num_frames: int,\n                 semantic_map_path: str, world_to_ecef: np.ndarray):\n        super().__init__()\n\n        self.raster_size = np.array(raster_size)\n        self.pixel_size = np.array(pixel_size)\n        self.ego_center = ego_center\n        self.filter_agents_threshold = filter_agents_threshold\n        self.history_num_frames = history_num_frames\n        self.world_to_ecef = world_to_ecef\n        self.proto_API = MapAPI(semantic_map_path, world_to_ecef)\n        \n        self.render_context = RenderContext(\n            raster_size_px=self.raster_size,\n            pixel_size_m=self.pixel_size,\n            center_in_raster_ratio=self.ego_center,\n        )\n\n        self.bounds_info = self.get_bounds()\n\n    def get_bounds(self) -> dict:\n        \"\"\"\n        For each elements of interest returns bounds [[min_x, min_y],[max_x, max_y]] and proto ids\n        Coords are computed by the MapAPI and, as such, are in the world ref system.\n\n        Returns:\n            dict: keys are classes of elements, values are dict with `bounds` and `ids` keys\n        \"\"\"\n        lanes_ids = []\n        crosswalks_ids = []\n        speedbumps_ids = []\n        speedhumps_ids = []\n\n        lanes_bounds = np.empty((0, 2, 2), dtype=np.float)  # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]\n        crosswalks_bounds = np.empty((0, 2, 2), dtype=np.float)  # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]\n        speedbumps_bounds = np.empty((0, 2, 2), dtype=np.float)  # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]\n        speedhumps_bounds = np.empty((0, 2, 2), dtype=np.float)  # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]\n\n        for element in self.proto_API:\n            element_id = MapAPI.id_as_str(element.id)\n\n            if self.proto_API.is_lane(element):\n                lane = self.proto_API.get_lane_coords(element_id)\n                x_min = np.min(lane[\"xyz\"][0, :])\n                y_min = np.min(lane[\"xyz\"][1, :])\n                x_max = np.max(lane[\"xyz\"][0, :])\n                y_max = np.max(lane[\"xyz\"][1, :])\n\n                lanes_bounds = np.append(lanes_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0)\n                lanes_ids.append(element_id)\n\n            if self.proto_API.is_crosswalk(element):\n                crosswalk = self.proto_API.get_crosswalk_coords(element_id)\n                x_min = np.min(crosswalk[\"xyz\"][0, :])\n                y_min = np.min(crosswalk[\"xyz\"][1, :])\n                x_max = np.max(crosswalk[\"xyz\"][0, :])\n                y_max = np.max(crosswalk[\"xyz\"][1, :])\n\n                crosswalks_bounds = np.append(\n                    crosswalks_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0,\n                )\n                crosswalks_ids.append(element_id)\n\n            if self.proto_API.is_speedbump(element):\n                speedbump = self.proto_API.get_speedbump_coords(element_id)\n                x_min = np.min(speedbump[\"xyz\"][0, :])\n                y_min = np.min(speedbump[\"xyz\"][1, :])\n                x_max = np.max(speedbump[\"xyz\"][0, :])\n                y_max = np.max(speedbump[\"xyz\"][1, :])\n\n                speedbumps_bounds = np.append(\n                    speedbumps_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0,\n                )\n                speedbumps_ids.append(element_id)\n\n            if self.proto_API.is_speedhump(element):\n                speedhump = self.proto_API.get_speedhump_coords(element_id)\n                x_min = np.min(speedhump[\"xyz\"][0, :])\n                y_min = np.min(speedhump[\"xyz\"][1, :])\n                x_max = np.max(speedhump[\"xyz\"][0, :])\n                y_max = np.max(speedhump[\"xyz\"][1, :])\n\n                speedhumps_bounds = np.append(\n                    speedhumps_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0,\n                )\n                speedhumps_ids.append(element_id)\n\n        return {\n            \"lanes\": {\"bounds\": lanes_bounds, \"ids\": lanes_ids},\n            \"crosswalks\": {\"bounds\": crosswalks_bounds, \"ids\": crosswalks_ids},\n            \"speedbumps\": {\"bounds\": speedbumps_bounds, \"ids\": speedbumps_ids},\n            \"speedhumps\": {\"bounds\": speedhumps_bounds, \"ids\": speedhumps_ids},\n        }\n\n    def rasterize(self,\n                  history_frames: np.ndarray,\n                  history_agents: List[np.ndarray],\n                  history_tl_faces: List[np.ndarray],\n                  agent: Optional[np.ndarray] = None\n                  ) -> Union[np.ndarray, dict]:\n\n        if agent is None:\n            ego_translation_m = history_frames[0][\"ego_translation\"]\n            ego_yaw_rad = rotation33_as_yaw(history_frames[0][\"ego_rotation\"])\n        else:\n            ego_translation_m = np.append(agent[\"centroid\"], history_frames[0][\"ego_translation\"][-1])\n            ego_yaw_rad = agent[\"yaw\"]\n\n        raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)\n        world_from_raster = np.linalg.inv(raster_from_world)\n\n        # get XY of center pixel in world coordinates\n        center_in_raster_px = np.asarray(self.raster_size) * (0.5, 0.5)\n        center_in_world_m = transform_point(center_in_raster_px, world_from_raster)\n        \n        # box_channels = (self.history_num_frames + 1) * 2\n        # channels = box_channels + 3\n\n        res = {\n            \"lines\": self.generate_map_data(center_in_world_m, raster_from_world, history_tl_faces[0]),\n            \"agents\": self.generate_agents(history_frames, history_agents, history_tl_faces,\n                                           raster_from_world, agent),\n            \"raster_from_world\": raster_from_world\n        }\n\n        return res\n\n    def to_rgb(self, in_im: np.ndarray, **kwargs: dict) -> np.ndarray:\n        out_im = in_im[0:3, ...].copy() * 255.0\n        out_im = np.transpose(out_im, (1, 2, 0))\n\n        out_im_agent = np.zeros((self.raster_size[1], self.raster_size[0], 3), dtype=np.float32)\n        agent_chs = in_im[3:3+self.history_num_frames+1][::-1]\n        agent_color = (0, 0, 255)\n        for ch in agent_chs:\n            out_im_agent *= 0.85\n            out_im_agent[ch > 0] = agent_color\n\n        out_im_ego = np.zeros((self.raster_size[1], self.raster_size[0], 3), dtype=np.float32)\n        ego_chs = in_im[3+self.history_num_frames+1:][::-1]\n        ego_color = (0, 255, 0)\n\n        for ch in ego_chs:\n            out_im_ego *= 0.85\n            out_im_ego[ch > 0] = ego_color\n\n        out_box_im = (np.clip(out_im_agent + out_im_ego, 0, 255)).astype(np.uint8)\n\n        mask_box = np.any(out_box_im > 0, -1)\n        out_im[mask_box] = out_box_im[mask_box]\n        out_im = out_im.astype(np.uint8)\n\n        return out_im\n\n    def generate_map_data(self, center_world: np.ndarray, world_to_image_space: np.ndarray,\n                          tl_faces: np.ndarray) -> Union[np.ndarray, dict]:\n\n        # filter using half a radius from the center\n        raster_radius = float(np.linalg.norm(self.raster_size * self.pixel_size)) / 2\n\n        # get active traffic light faces\n        active_tl_ids = set(filter_tl_faces_by_status(tl_faces, \"ACTIVE\")[\"face_id\"].tolist())\n\n        res = []\n        lane_ids = []\n        tls = {}\n        for idx in elements_within_bounds(center_world, self.bounds_info[\"lanes\"][\"bounds\"], raster_radius):\n            lane_id = self.bounds_info[\"lanes\"][\"ids\"][idx]\n            lane = self.proto_API[lane_id].element.lane\n\n            # no traffic light face is controlling\n            # this lane (blue channel 128 pixel value)\n            lane_type = [2, 128]\n            tl = 0\n\n            lane_tl_ids = set([MapAPI.id_as_str(la_tc) for la_tc in lane.traffic_controls])\n\n            for tl_id in lane_tl_ids.intersection(active_tl_ids):\n                if self.proto_API.is_traffic_face_colour(tl_id, \"red\"):\n                    lane_type = [0, 255]\n                    tl = 1\n                elif self.proto_API.is_traffic_face_colour(tl_id, \"green\"):\n                    lane_type = [1, 255]\n                    tl = 2\n                elif self.proto_API.is_traffic_face_colour(tl_id, \"yellow\"):\n                    lane_type = [1, 128]\n                    tl = 3\n\n            lane_coords = self.proto_API.get_lane_coords(lane_id)\n            lanes_xy = world_to_image_space.dot(lane_coords[\"xyz\"]).T[:, :2]\n\n            lanes = np.zeros(shape=(len(lanes_xy) - 1, 7), dtype=np.float32)\n            lanes[:, 0] = 0\n\n            lanes[:, 5] = lane_type[0]  # channel\n            lanes[:, 6] = lane_type[1]  # value\n\n            if lane_type[0] == 2:\n                # no traffic light -> lane boundary type encode\n                ln = lane_coords[\"left_point_num\"]\n                left_type = lane_coords[\"left_divider_type\"]\n                right_type = lane_coords[\"right_divider_type\"]\n\n                if left_type > 12:\n                    left_type = 0\n\n                if right_type > 12:\n                    right_type = 0\n\n                lanes[:ln, 6] = DIVIDER_TYPE_CODES[left_type]\n                lanes[ln:, 6] = DIVIDER_TYPE_CODES[right_type]\n\n            lanes[:, 1:3] = lanes_xy[:-1, 0:2]\n            lanes[:, 3:5] = lanes_xy[1:, 0:2]\n\n            res.append(lanes)\n\n        for idx in elements_within_bounds(center_world, self.bounds_info[\"crosswalks\"][\"bounds\"], raster_radius):\n            crosswalk = self.proto_API.get_crosswalk_coords(self.bounds_info[\"crosswalks\"][\"ids\"][idx])\n\n            cross_xy = world_to_image_space.dot(crosswalk[\"xyz_t\"]).T[:, :2]\n\n            crosswalks = np.zeros(shape=(len(cross_xy), 7), dtype=np.float32)\n            crosswalks[:, 0] = 0\n            crosswalks[:, 5] = 0\n            crosswalks[:, 6] = 128\n\n            crosswalks[:-1, 1:3] = cross_xy[:-1, 0:2]\n            crosswalks[:-1, 3:5] = cross_xy[1:, 0:2]\n\n            # \"Polygon\" closing line\n            crosswalks[-1, 1:3] = cross_xy[-1, 0:2]\n            crosswalks[-1, 3:5] = cross_xy[0, 0:2]\n\n            res.append(crosswalks)\n\n        for idx in elements_within_bounds(center_world, self.bounds_info[\"speedbumps\"][\"bounds\"], raster_radius):\n            speedbumps = self.proto_API.get_speedbump_coords(self.bounds_info[\"speedbumps\"][\"ids\"][idx])\n\n            cross_xy = world_to_image_space.dot(speedbumps[\"xyz\"]).T[:, :2]\n\n            speedbumps = np.zeros(shape=(len(cross_xy), 7), dtype=np.float32)\n            speedbumps[:, 0] = 0\n            speedbumps[:, 5] = 1\n            speedbumps[:, 6] = 128\n\n            speedbumps[:-1, 1:3] = cross_xy[:-1, 0:2]\n            speedbumps[:-1, 3:5] = cross_xy[1:, 0:2]\n\n            # \"Polygon\" closing line\n            speedbumps[-1, 1:3] = cross_xy[-1, 0:2]\n            speedbumps[-1, 3:5] = cross_xy[0, 0:2]\n\n            res.append(speedbumps)\n            \n        for idx in elements_within_bounds(center_world, self.bounds_info[\"speedhumps\"][\"bounds\"], raster_radius):\n            speedhumps = self.proto_API.get_speedhump_coords(self.bounds_info[\"speedhumps\"][\"ids\"][idx])\n\n            cross_xy = world_to_image_space.dot(speedhumps[\"xyz\"]).T[:, :2]\n\n            speedhumps = np.zeros(shape=(len(cross_xy), 7), dtype=np.float32)\n            speedhumps[:, 0] = 0\n            speedhumps[:, 5] = 1\n            speedhumps[:, 6] = 64\n\n            speedhumps[:-1, 1:3] = cross_xy[:-1, 0:2]\n            speedhumps[:-1, 3:5] = cross_xy[1:, 0:2]\n\n            # \"Polygon\" closing line\n            speedhumps[-1, 1:3] = cross_xy[-1, 0:2]\n            speedhumps[-1, 3:5] = cross_xy[0, 0:2]\n\n            res.append(speedhumps)\n\n        return np.concatenate(res) if len(res) > 0 else np.array([])\n\n    def generate_agents(self, history_frames: np.ndarray, history_agents: List[np.ndarray],\n                        history_tl_faces: List[np.ndarray], world_to_image_space: np.ndarray,\n                        agent: Optional[np.ndarray] = None):\n\n        nframes = len(history_frames)\n        boxes = []\n\n        for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):\n            agents = filter_agents_by_labels(agents, self.filter_agents_threshold)\n            av_agent = get_ego_as_agent(frame).astype(agents.dtype)\n\n            if agent is None:\n                # generate agents\n                self.generate_box_data(boxes, agents, 3 + i, 255)\n\n                # generate av_agent (ego)\n                self.generate_box_data(boxes, av_agent, 3 + nframes + i, 255)\n            else:\n                agent_ego = filter_agents_by_track_id(agents, agent[\"track_id\"])\n\n                if len(agent_ego) == 0:  # Agent not in this history frame\n                    # generate agents\n                    self.generate_box_data(boxes, np.append(agents, av_agent), 3 + i, 255)\n\n                else:  # Add AV to agents and remove the agent from agents\n                    agents = agents[agents != agent_ego[0]]\n\n                    # generate agents\n                    self.generate_box_data(boxes, np.append(agents, av_agent), 3 + i, 255)\n\n                    # generate agent (ego)\n                    self.generate_box_data(boxes, agent_ego, 3 + nframes + i, 255)\n\n        return np.array(boxes)\n\n    def generate_box_data(self, boxes, agents, channel_id, value):\n\n        for idx, agent in enumerate(agents):\n            extent = agent[\"extent\"]\n\n            if agent[\"label_probabilities\"][3] == 1:\n                # Car - minimum extent..\n                # EGO_EXTENT_WIDTH = 1.85\n                # EGO_EXTENT_LENGTH = 4.87\n                extent[0] = max(3, extent[0])\n                extent[1] = max(1.4, extent[1])\n\n            elif agent[\"label_probabilities\"][12] == 1:\n                # Cyclists\n                extent[0] = max(1.3, extent[0])\n                extent[1] = max(0.5, extent[1])\n\n            if agent[\"label_probabilities\"][14] == 1:\n                # Increase the pedestrians \"size\"\n                extent = extent * 1.8\n            else:\n                extent = extent\n\n            boxes.append((\n                0,  # Batch index (filled later)\n\n                agent[\"centroid\"][0],  # c_x\n                agent[\"centroid\"][1],  # c_y\n\n                extent[0],  # Length\n                extent[1],  # Width\n                extent[2],  # Height\n                agent[\"yaw\"],  # yaw\n\n                channel_id,\n                value\n            ))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<a id=\"benchmarks\"></a>\n# Benchmarks\n\n*Note: The code below is from the official [baseline notebook](https://www.kaggle.com/lucabergamini/lyft-baseline-09-02)*"},{"metadata":{"trusted":true},"cell_type":"code","source":"os.environ[\"L5KIT_DATA_FOLDER\"] = \"/kaggle/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': [350, 220],\n        'pixel_size': [0.4, 0.4],\n        'ego_center': [0.25, 0.5],\n        'map_type': 'py_semantic',\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/sample.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 4\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Official L5Kit Rasterizer"},{"metadata":{"trusted":true},"cell_type":"code","source":"def build_model(cfg: Dict) -> torch.nn.Module:\n    # load pre-trained Conv2D model\n    model = resnet50(pretrained=False)\n\n    # change input channels number to match the rasterizer's output\n    num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n    num_in_channels = 3 + num_history_channels\n    model.conv1 = nn.Conv2d(\n        num_in_channels,\n        model.conv1.out_channels,\n        kernel_size=model.conv1.kernel_size,\n        stride=model.conv1.stride,\n        padding=model.conv1.padding,\n        bias=False,\n    )\n    # change output size to (X, Y) * number of future states\n    num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n    model.fc = nn.Linear(in_features=2048, out_features=num_targets)\n\n    return model\n\ndef forward(data, model, device):\n    inputs = data[\"image\"].permute(0, 3, 1, 2).to(device)\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    # Forward pass\n    outputs = model(inputs).reshape(targets.shape)\n    return outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntrain_cfg = cfg[\"train_data_loader\"]\n\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n\nrasterizer = build_rasterizer(cfg, dm)\ntrain_dataset = AgentDataset(cfg, train_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\n# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = build_model(cfg).to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"progress_bar = tqdm(train_dataloader, total=100)\nmodel.train()\n\nfor itr, data in enumerate(progress_bar):\n    \n    outputs = forward(data, model, device)\n    # Skip the rest of the training...\n    \n    if itr == 100:\n        break\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## GPU Rasterizer"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"def lyft_gpu_collate_fn(batch):\n    element = batch[0]\n\n    rasterizer_args = [\n        \"lines\",\n        \"agents\"\n    ]\n\n    result = {\n        key: default_collate([d[key] for d in batch]) for key in element if key not in rasterizer_args\n    }\n\n    for element in batch:\n\n        for rasterizer_arg in rasterizer_args:\n            if rasterizer_arg not in result:\n                result[rasterizer_arg] = []\n\n            result[rasterizer_arg].append(element[rasterizer_arg])\n\n    for bidx in range(len(result[\"lines\"])):\n        result[\"lines\"][bidx][:, 0] = bidx\n\n    for bidx in range(len(result[\"agents\"])):\n        result[\"agents\"][bidx][:, 0] = bidx\n\n    result[\"lines\"] = torch.from_numpy(np.concatenate(result[\"lines\"]))\n    result[\"agents\"] = torch.from_numpy(np.concatenate(result[\"agents\"]))\n\n    return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"raster_cfg = cfg[\"raster_params\"]\nmodel_cfg = cfg[\"model_params\"]\n\nsemantic_map_filepath = dm.require(raster_cfg[\"semantic_map_key\"])\n\ntry:\n    dataset_meta = _load_metadata(raster_cfg[\"dataset_meta_key\"], dm)\n    world_to_ecef = np.array(dataset_meta[\"world_to_ecef\"], dtype=np.float64)\nexcept (KeyError, FileNotFoundError):\n    world_to_ecef = get_hardcoded_world_to_ecef()\n\ndata_rasterizer = GPURasterizer(\n    raster_size=raster_cfg[\"raster_size\"],\n    pixel_size=raster_cfg[\"pixel_size\"],\n    ego_center=raster_cfg[\"ego_center\"],\n    filter_agents_threshold=0.5,\n    history_num_frames=cfg[\"model_params\"][\"history_num_frames\"],\n    semantic_map_path=semantic_map_filepath,\n    world_to_ecef=world_to_ecef,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def forward_gpu(data, model, device, cuda_generator):\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    matrix = data[\"raster_from_world\"].to(device).to(torch.float)\n    \n    data[\"lines\"] = data[\"lines\"].to(device).type(torch.cuda.FloatTensor)\n    data[\"agents\"] = data[\"agents\"].to(device).type(torch.cuda.FloatTensor)\n    \n    cuda_generator.rasterize(data[\"image\"], data[\"lines\"], data[\"agents\"], matrix)\n    \n    data[\"image\"] /= 255.\n    \n    last_batch = False\n    if len(data[\"image\"]) != len(targets):\n        # Last batch\n        last_batch = len(targets)    \n    \n    # Forward pass\n    outputs = model(data[\"image\"] if not last_batch else data[\"image\"][:last_batch, ...]).reshape(targets.shape)\n    \n    return outputs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open(cached=True)\ntrain_dataset = AgentDataset(cfg, train_zarr, data_rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                              collate_fn=lyft_gpu_collate_fn,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"],\n                              drop_last=True)\n\nchannels = ((cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2) + 3\n\nimages = torch.zeros((train_cfg[\"batch_size\"],\n                      channels,\n                      raster_cfg[\"raster_size\"][1],\n                      raster_cfg[\"raster_size\"][0]))\n\nimages = images.to(device).type(torch.cuda.FloatTensor)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"progress_bar = tqdm(train_dataloader, total=100)\nmodel.train()\n\nfor itr, data in enumerate(progress_bar):\n    \n    data[\"image\"] = images\n    data[\"image\"] *= 0    \n    \n    outputs = forward_gpu(data, model, device, lyft_gpu)\n    # Skip the rest of the training...\n    \n    \n    if itr == 100:\n        break\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img = data_rasterizer.to_rgb(data[\"image\"][0].data.cpu().numpy())\n\nplt.imshow(img[::-1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"--------------------\n**Thank you for reading**"}],"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}