Add renderer compilation and gradient checkpointing (#384)

Thanks to @etienne87 for the idea!

Author: eigenvivek

diffdrr/drr.py

@@ -39,6 +39,8 @@ def __init__(
         patch_size: int | None = None,  # Render patches of the DRR in series
         renderer: str = "siddon",  # Rendering backend, either "siddon" or "trilinear"
         persistent: bool = True,  # Set persistent value in `torch.nn.Module.register_buffer`
+        compile_renderer: bool = False,  # Compile the renderer for performance boost
+        checkpoint_gradients: bool = False,  # Checkpoint gradients to improve memory usage
         **renderer_kwargs,  # Kwargs for the renderer
     ):
         super().__init__()
@@ -96,8 +98,11 @@ def __init__(
             raise ValueError(
                 f"renderer must be 'siddon' or 'trilinear', not {renderer}"
             )
+        if compile_renderer:
+            self.renderer = torch.compile(self.renderer, mode="default")
         self.reshape = reshape
         self.patch_size = patch_size
+        self.checkpoint_gradients = checkpoint_gradients
 
     def reshape_transform(self, img, batch_size):
         if self.reshape:
@@ -141,6 +146,8 @@ def reshape_subsampled_drr(img: torch.Tensor, detector: Detector, batch_size: in
     return drr
 
 # %% ../notebooks/api/00_drr.ipynb 10
+from torch.utils.checkpoint import checkpoint
+
 from .pose import RigidTransform, convert
 
 
@@ -163,7 +170,19 @@ def forward(
 
     # Create the source / target points and render the image
     source, target = self.detector(pose, calibration)
-    img = self.render(self.density, source, target, mask_to_channels, **kwargs)
+
+    if self.checkpoint_gradients:
+        img = checkpoint(
+            self.render,
+            self.density,
+            source,
+            target,
+            mask_to_channels,
+            **kwargs,
+            use_reentrant=False,
+        )
+    else:
+        img = self.render(self.density, source, target, mask_to_channels, **kwargs)
     return self.reshape_transform(img, batch_size=len(pose))
 
 

notebooks/api/00_drr.ipynb

@@ -134,6 +134,8 @@
     "        patch_size: int | None = None,  # Render patches of the DRR in series\n",
     "        renderer: str = \"siddon\",  # Rendering backend, either \"siddon\" or \"trilinear\"\n",
     "        persistent: bool = True,  # Set persistent value in `torch.nn.Module.register_buffer`\n",
+    "        compile_renderer: bool = False,  # Compile the renderer for performance boost\n",
+    "        checkpoint_gradients: bool = False,  # Checkpoint gradients to improve memory usage\n",
     "        **renderer_kwargs,  # Kwargs for the renderer\n",
     "    ):\n",
     "        super().__init__()\n",
@@ -191,8 +193,11 @@
     "            raise ValueError(\n",
     "                f\"renderer must be 'siddon' or 'trilinear', not {renderer}\"\n",
     "            )\n",
+    "        if compile_renderer:\n",
+    "            self.renderer = torch.compile(self.renderer, mode=\"default\")\n",
     "        self.reshape = reshape\n",
     "        self.patch_size = patch_size\n",
+    "        self.checkpoint_gradients = checkpoint_gradients\n",
     "\n",
     "    def reshape_transform(self, img, batch_size):\n",
     "        if self.reshape:\n",
@@ -260,6 +265,8 @@
    "outputs": [],
    "source": [
     "#| export\n",
+    "from torch.utils.checkpoint import checkpoint\n",
+    "\n",
     "from diffdrr.pose import RigidTransform, convert\n",
     "\n",
     "\n",
@@ -282,7 +289,19 @@
     "\n",
     "    # Create the source / target points and render the image\n",
     "    source, target = self.detector(pose, calibration)\n",
-    "    img = self.render(self.density, source, target, mask_to_channels, **kwargs)\n",
+    "\n",
+    "    if self.checkpoint_gradients:\n",
+    "        img = checkpoint(\n",
+    "            self.render,\n",
+    "            self.density,\n",
+    "            source,\n",
+    "            target,\n",
+    "            mask_to_channels,\n",
+    "            **kwargs,\n",
+    "            use_reentrant=False,\n",
+    "        )\n",
+    "    else:\n",
+    "        img = self.render(self.density, source, target, mask_to_channels, **kwargs)\n",
     "    return self.reshape_transform(img, batch_size=len(pose))\n",
     "\n",
     "\n",
@@ -408,7 +427,7 @@
     "    x[..., 1] = self.detector.height - x[..., 1]\n",
     "    if self.detector.reverse_x_axis:\n",
     "        x[..., 0] = self.detector.width - x[..., 0]\n",
-    "    \n",
+    "\n",
     "    return x[..., :2]"
    ]
   },