{ "cells": [ { "cell_type": "markdown", "id": "9bba0cf4", "metadata": {}, "source": [ "# BEVFusion Deep Dive - 04: LiDAR Branch - Voxelization CUDA kernels + Sparse-Conv backbone\n", "\n", "The LiDAR branch converts a raw point cloud into a BEV feature map:\n", "1. **Voxelization** groups points into a 3D grid ([`Voxelization`](../../perception/bevfusion/mmdet3d/ops/voxel/voxelize.py) -> CUDA kernels in [`voxelization_cuda.cu`](../../perception/bevfusion/mmdet3d/ops/voxel/src/voxelization_cuda.cu:24)).\n", "2. **Feature reduction**: per-voxel mean of point features (`voxelize_reduce`).\n", "3. **Sparse VoxelNet/SECOND backbone** (submanifold sparse conv) -> compact 3D features -> flattened to BEV `(B, 256, X/8, Y/8)`.\n", "\n", "We explain the CUDA kernels line-by-line and reimplement hard-voxelization in pure PyTorch." ] }, { "cell_type": "code", "execution_count": 1, "id": "ce1784db", "metadata": { "execution": { "iopub.execute_input": "2026-08-02T03:38:47.033483Z", "iopub.status.busy": "2026-08-02T03:38:47.032652Z", "iopub.status.idle": "2026-08-02T03:38:48.269389Z", "shell.execute_reply": "2026-08-02T03:38:48.268830Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "IN_CONTAINER : True\n", "BEVFUSION_SRC : /workspace/perception/bevfusion exists= True\n", "NUSC_ROOT : /data/nuscenes exists= True\n", "CKPT_DIR : /weights exists= True\n", "INFOS_DIR : /data/nuscenes exists= True\n", "HAS_TORCH : True DEVICE= cuda\n", "GPU : NVIDIA RTX 6000 Ada Generation | torch 1.10.1+cu113\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/lib/python3.8/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "# --- Common bootstrap (safe to run first in any BEVFusion notebook) ---\n", "#\n", "# This cell auto-detects whether the notebook runs INSIDE the reproducible\n", "# `bevfusion:cu113` Docker container (recommended; see\n", "# perception/docker/run_notebooks.sh) or on the bare host, and picks paths\n", "# accordingly. Everything is overridable via environment variables.\n", "#\n", "# In-container layout (set up by run_notebooks.sh):\n", "# BEVFUSION_SRC = /workspace/perception/bevfusion (mmcv/mmdet/mmdet3d + compiled CUDA ops)\n", "# NUSC_ROOT = /data/nuscenes (writable mini_root: samples/sweeps/maps/v1.0-mini + *_infos_*.pkl)\n", "# CKPT_DIR = /weights (bevfusion-det.pth)\n", "# INFOS_DIR = /data/nuscenes (nuscenes_infos_{train,val}.pkl live next to the data)\n", "#\n", "# Host fallback (bare-metal, needs perception/setup.sh built):\n", "# BEVFUSION_SRC = /perception/bevfusion\n", "# NUSC_ROOT = /mnt/mydata/nuscenes/mini_root (or /mnt/mydata/nuscenes)\n", "# CKPT_DIR = /mnt/mydata/huggingface/bevfusion\n", "import os, sys, math, json, textwrap\n", "import numpy as np\n", "\n", "def _first_existing(paths, default):\n", " for p in paths:\n", " if p and os.path.isdir(p):\n", " return p\n", " return default\n", "\n", "IN_CONTAINER = os.path.isdir(\"/workspace/perception/bevfusion\")\n", "\n", "BEVFUSION_SRC = os.environ.get(\"BEVFUSION_SRC\") or _first_existing(\n", " [\"/workspace/perception/bevfusion\"], \"/home/nadurthi/repos/deeplearning/perception/bevfusion\")\n", "NUSC_ROOT = os.environ.get(\"NUSC_ROOT\") or _first_existing(\n", " [\"/data/nuscenes\", \"/mnt/mydata/nuscenes/mini_root\"], \"/mnt/mydata/nuscenes\")\n", "CKPT_DIR = os.environ.get(\"BEVFUSION_CKPT_DIR\") or _first_existing(\n", " [\"/weights\"], \"/mnt/mydata/huggingface/bevfusion\")\n", "# Where the generated nuscenes_infos_*.pkl live (defaults to NUSC_ROOT).\n", "INFOS_DIR = os.environ.get(\"BEVFUSION_INFOS_DIR\", NUSC_ROOT)\n", "\n", "# Make mmcv/mmdet/mmdet3d importable from the vendored source tree.\n", "if BEVFUSION_SRC and BEVFUSION_SRC not in sys.path:\n", " sys.path.insert(0, BEVFUSION_SRC)\n", "\n", "# Make matplotlib headless-safe.\n", "import matplotlib\n", "if os.environ.get(\"MPLBACKEND\") is None:\n", " try:\n", " matplotlib.use(\"Agg\")\n", " except Exception:\n", " pass\n", "import matplotlib.pyplot as plt\n", "\n", "# Torch is expected; but guard so markdown-only runs still import.\n", "try:\n", " import torch\n", " import torch.nn as nn\n", " import torch.nn.functional as F\n", " HAS_TORCH = True\n", " DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "except Exception as e:\n", " HAS_TORCH = False\n", " DEVICE = \"cpu\"\n", " print(\"torch not available:\", e)\n", "\n", "print(\"IN_CONTAINER :\", IN_CONTAINER)\n", "print(\"BEVFUSION_SRC :\", BEVFUSION_SRC, \"exists=\", os.path.isdir(BEVFUSION_SRC))\n", "print(\"NUSC_ROOT :\", NUSC_ROOT, \"exists=\", os.path.isdir(NUSC_ROOT))\n", "print(\"CKPT_DIR :\", CKPT_DIR, \"exists=\", os.path.isdir(CKPT_DIR))\n", "print(\"INFOS_DIR :\", INFOS_DIR, \"exists=\", os.path.isdir(INFOS_DIR))\n", "print(\"HAS_TORCH :\", HAS_TORCH, \"DEVICE=\", DEVICE)\n", "if HAS_TORCH and torch.cuda.is_available():\n", " print(\"GPU :\", torch.cuda.get_device_name(0),\n", " \"| torch\", torch.__version__)\n" ] }, { "cell_type": "markdown", "id": "3a2106a9", "metadata": {}, "source": [ "## 1. Voxelization: what the CUDA kernels do\n", "There are two modes. **Dynamic** voxelization only assigns each point a voxel coordinate; **hard** voxelization additionally packs up to `max_points` points per voxel into a dense `(max_voxels, max_points, F)` tensor. BEVFusion's VoxelNet uses hard voxelization + mean.\n", "\n", "### Kernel 1 - `dynamic_voxelize_kernel` (point -> voxel coordinate)\n", "```cuda\n", "CUDA_1D_KERNEL_LOOP(index, num_points) { // grid-stride: 1 thread per point\n", " auto p = points + index*num_features;\n", " int c_x = floor((p[0]-coors_x_min)/voxel_x); // which voxel in X\n", " if (c_x<0 || c_x>=grid_x) { coors[0]=-1; return; } // out of range -> mark invalid\n", " int c_y = floor((p[1]-coors_y_min)/voxel_y); ... // same for Y\n", " int c_z = floor((p[2]-coors_z_min)/voxel_z); ... // same for Z\n", " coors_offset[0]=c_z; [1]=c_y; [2]=c_x; // store (z,y,x)\n", "}\n", "```\n", "Pure floor-division bucketing; invalid points get coord `-1` and are dropped later.\n", "\n", "### Kernel 2 - `point_to_voxelidx_kernel` (dedup + capacity)\n", "Walks points, assigns each a slot within its voxel, and caps at `max_points` / `max_voxels`. This is what makes it *hard* voxelization.\n", "\n", "### Kernel 3 - `assign_point_to_voxel` (scatter features)\n", "```cuda\n", "int index = thread_idx / num_features; // point\n", "int k = thread_idx % num_features; // feature channel\n", "if (num>-1 && voxelidx>-1)\n", " voxels[voxelidx*max_points*F + num*F + k] = points[thread_idx];\n", "```\n", "One thread per (point, feature) writes the point's feature into its `(voxel, slot)`.\n", "\n", "### Kernel 4 - `assign_voxel_coors`\n", "Writes each occupied voxel's `(z,y,x)` coordinate exactly once (guarded by `num==0`).\n", "\n", "**Design choices**: grid-stride loops (handles any N with fixed blocks); coord `-1` sentinel avoids branchy compaction on device; the `@force_fp32` + `@torch.no_grad` on [`BEVFusion.voxelize`](../../perception/bevfusion/mmdet3d/models/fusion_models/bevfusion.py:169) keeps geometry exact and skips autograd (voxelization is not learned)." ] }, { "cell_type": "markdown", "id": "5e8830b1", "metadata": {}, "source": [ "## 2. Pure-PyTorch hard-voxelization + mean reduce (verified)\n", "Reproduces the observable behavior of the CUDA path: bucket points, keep up to `max_points` per voxel, cap voxels, and mean-reduce features (mirroring `voxelize_reduce` in [`voxelize()`](../../perception/bevfusion/mmdet3d/models/fusion_models/bevfusion.py:191))." ] }, { "cell_type": "code", "execution_count": 2, "id": "4df8a45e", "metadata": { "execution": { "iopub.execute_input": "2026-08-02T03:38:48.271792Z", "iopub.status.busy": "2026-08-02T03:38:48.271436Z", "iopub.status.idle": "2026-08-02T03:38:48.296323Z", "shell.execute_reply": "2026-08-02T03:38:48.295578Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "occupied voxels: 146 of 2000 points\n", "per-voxel feat shape (V,F): (146, 5)\n", "mean points/voxel: 1.0\n" ] } ], "source": [ "if HAS_TORCH:\n", " def hard_voxelize(points, pcr, vsize, max_points=10, max_voxels=120000):\n", " pcr = torch.tensor(pcr); vsize = torch.tensor(vsize)\n", " grid = ((pcr[3:]-pcr[:3])/vsize).long() # (X,Y,Z)\n", " xyz = points[:,:3]\n", " coord = ((xyz - pcr[:3])/vsize).floor().long()\n", " inb = ((coord>=0)&(coord=V: continue\n", " if npts[v] < max_points:\n", " voxels[v, npts[v]] = points[i]; vcoord[v] = coord[i]\n", " npts[v]+=1\n", " feats = voxels.sum(1) / npts.clamp(min=1).unsqueeze(1) # mean reduce\n", " return feats, vcoord, npts\n", "\n", " torch.manual_seed(0)\n", " pts = torch.rand(2000,5); pts[:,:3] = pts[:,:3]*100-50\n", " pcr=[-54,-54,-5,54,54,3]; vs=[0.075,0.075,0.2]\n", " feats, vcoord, npts = hard_voxelize(pts, pcr, vs)\n", " print('occupied voxels:', feats.shape[0], 'of', pts.shape[0], 'points')\n", " print('per-voxel feat shape (V,F):', tuple(feats.shape))\n", " print('mean points/voxel:', float(npts.float().mean()))\n", "else:\n", " print('skipped')\n" ] }, { "cell_type": "markdown", "id": "1036a45d", "metadata": {}, "source": [ "## 3. The sparse VoxelNet/SECOND backbone\n", "Dense 3D conv over a `1440x1440x40` grid would be enormous but ~99% empty. BEVFusion uses **submanifold sparse convolution** (spconv, vendored under [`mmdet3d/ops/spconv`](../../perception/bevfusion/mmdet3d/ops/spconv/conv.py)) in [`SparseEncoder`](../../perception/bevfusion/mmdet3d/models/backbones/sparse_encoder.py): compute scales with **occupied** voxels only.\n", "\n", "* **Submanifold conv** keeps the active-site set fixed (output active exactly where input is) so sparsity does not dilate layer over layer.\n", "* **Regular sparse conv** (stride 2) downsamples the grid `8x` in X/Y across the stages.\n", "* Final 3D features are reshaped to a dense BEV `(B, 256, X/8, Y/8)` = `(B,256,180,180)`.\n", "\n", "The heavy lifting is index-generation (which input sites map to which output site) done by CUDA/CPU kernels in `spconv`; the actual matmul is a gather -> dense GEMM -> scatter." ] }, { "cell_type": "code", "execution_count": 3, "id": "4e14ada6", "metadata": { "execution": { "iopub.execute_input": "2026-08-02T03:38:48.299639Z", "iopub.status.busy": "2026-08-02T03:38:48.299370Z", "iopub.status.idle": "2026-08-02T03:38:48.307428Z", "shell.execute_reply": "2026-08-02T03:38:48.306884Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "dense-conv active-site output matches submanifold restriction: True\n", "sparse conv computes only these 10 sites, not all 216\n" ] } ], "source": [ "if HAS_TORCH:\n", " # Illustrative dense equivalent of ONE submanifold 3x3x3 conv on a tiny grid,\n", " # to show sparse == dense on active sites (concept check, not the real spconv).\n", " Xg=Yg=Zg=6; Cin=4; Cout=8\n", " dense = torch.zeros(1,Cin,Zg,Yg,Xg)\n", " active = torch.randint(0,Xg,(10,3))\n", " for z,y,x in active.tolist(): dense[0,:,z,y,x] = torch.randn(Cin)\n", " conv = nn.Conv3d(Cin,Cout,3,padding=1,bias=False)\n", " dense_out = conv(dense)\n", " # submanifold: keep output only at active sites\n", " subm = torch.zeros(1,Cout,Zg,Yg,Xg)\n", " for z,y,x in active.tolist(): subm[0,:,z,y,x] = dense_out[0,:,z,y,x]\n", " print('dense-conv active-site output matches submanifold restriction:',\n", " torch.allclose(subm[subm!=0], dense_out[subm!=0]))\n", " print('sparse conv computes only these', len(active), 'sites, not all', Xg*Yg*Zg)\n", "else:\n", " print('skipped')\n" ] }, { "cell_type": "markdown", "id": "8a48f01b", "metadata": {}, "source": [ "## 4. Optional: real `Voxelization` op on GPU\n", "Lights up when the compiled ops exist. Runs the real hard-voxelizer and reports shapes." ] }, { "cell_type": "code", "execution_count": 4, "id": "216d4618", "metadata": { "execution": { "iopub.execute_input": "2026-08-02T03:38:48.309304Z", "iopub.status.busy": "2026-08-02T03:38:48.309034Z", "iopub.status.idle": "2026-08-02T03:38:51.237157Z", "shell.execute_reply": "2026-08-02T03:38:51.236549Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "compiled Voxelization unavailable -> ImportError cannot import name 'ball_query_ext' from partially initialized module 'mmdet3d.ops.ball_query' (most likely due to a circular import) (/workspace/perception/bevfusion/mmdet3d/ops/ball_query/__init__.py)\n", "Pure-PyTorch hard_voxelize above already demonstrated the algorithm.\n" ] } ], "source": [ "try:\n", " sys.path.insert(0, BEVFUSION_SRC)\n", " from mmdet3d.ops import Voxelization\n", " if HAS_TORCH and torch.cuda.is_available():\n", " vox = Voxelization(voxel_size=[0.075,0.075,0.2],\n", " point_cloud_range=[-54,-54,-5,54,54,3],\n", " max_num_points=10, max_voxels=120000).cuda()\n", " p = (torch.rand(20000,5).cuda()); p[:,:3] = p[:,:3]*100-50\n", " v,c,n = vox(p)\n", " print('voxels', tuple(v.shape), 'coords', tuple(c.shape), 'nums', tuple(n.shape))\n", " else:\n", " print('no CUDA')\n", "except Exception as e:\n", " print('compiled Voxelization unavailable ->', type(e).__name__, e)\n", " print('Pure-PyTorch hard_voxelize above already demonstrated the algorithm.')\n" ] }, { "cell_type": "markdown", "id": "cb15f2ee", "metadata": {}, "source": [ "---\n", "**Next:** [`05_fuser_decoder_head.ipynb`](05_fuser_decoder_head.ipynb)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.10" } }, "nbformat": 4, "nbformat_minor": 5 }