Skip to content

Commit 632c212

Browse files
yiheng-wang-nvyiheng-wang-nv
committed
add fast inference tutorial
Signed-off-by: Yiheng Wang <vennw@nvidia.com>
1 parent 4a40380 commit 632c212

File tree

3 files changed

+532
-0
lines changed

3 files changed

+532
-0
lines changed

‎acceleration/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ Typically, model training is a time-consuming step during deep learning developm
44
### List of notebooks and examples
55
#### [fast_model_training_guide](./fast_model_training_guide.md)
66
The document introduces details of how to profile the training pipeline, how to analyze the dataset and select suitable algorithms, and how to optimize GPU utilization in single GPU, multi-GPUs or even multi-nodes.
7+
#### [fast_inference_tutorial](./fast_inference_tutorial)
8+
The example introduces details of how to use GDS, GPU transforms and TensorRT to accelerate the inference.
79
#### [distributed_training](./distributed_training)
810
The examples show how to execute distributed training and evaluation based on 3 different frameworks:
911
- PyTorch native `DistributedDataParallel` module with `torchrun`.

‎acceleration/fast_inference_tutorial/fast_inference_tutorial.ipynb

Lines changed: 336 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,336 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"Copyright (c) MONAI Consortium \n",
8+
"Licensed under the Apache License, Version 2.0 (the \"License\"); \n",
9+
"you may not use this file except in compliance with the License. \n",
10+
"You may obtain a copy of the License at \n",
11+
"&nbsp;&nbsp;&nbsp;&nbsp;http://www.apache.org/licenses/LICENSE-2.0 \n",
12+
"Unless required by applicable law or agreed to in writing, software \n",
13+
"distributed under the License is distributed on an \"AS IS\" BASIS, \n",
14+
"WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. \n",
15+
"See the License for the specific language governing permissions and \n",
16+
"limitations under the License.\n",
17+
"\n",
18+
"# Fast Inference with MONAI features"
19+
]
20+
},
21+
{
22+
"cell_type": "markdown",
23+
"metadata": {},
24+
"source": [
25+
"This tutorial demonstrates the performance comparison between a standard PyTorch training program and a MONAI-optimized inference program. The key features include:\n",
26+
"\n",
27+
"1. **Direct Data Loading**: Load data directly from disk to GPU memory, minimizing data transfer time and improving efficiency.\n",
28+
"2. **GPU-based Preprocessing**: Execute preprocessing transforms directly on the GPU, leveraging its computational power for faster data preparation.\n",
29+
"3. **TensorRT Inference**: Utilize TensorRT for running inference, which optimizes the model for high-performance execution on NVIDIA GPUs.\n",
30+
"\n",
31+
"This tutorial is modified from the `TensorRT_inference_acceleration` tutorial."
32+
]
33+
},
34+
{
35+
"attachments": {},
36+
"cell_type": "markdown",
37+
"metadata": {},
38+
"source": [
39+
"## Setup environment"
40+
]
41+
},
42+
{
43+
"cell_type": "markdown",
44+
"metadata": {},
45+
"source": [
46+
"Loading data directly from disk to GPU memory requires the `kvikio` library. In addition, this tutorial requires many other dependencies such as `monai`, `torch`, `torch_tensorrt`, `numpy`, `ignite`, `pandas`, `matplotlib`, etc. We recommend using the [MONAI Docker](https://docs.monai.io/en/latest/installation.html#from-dockerhub) image to run this tutorial, which includes pre-configured dependencies and allows you to skip manual installation.\n",
47+
"\n",
48+
"If not using MONAI Docker, install `kvikio` using one of these methods:\n",
49+
"\n",
50+
"- **PyPI Installation** \n",
51+
" Use the appropriate package for your CUDA version:\n",
52+
" ```bash\n",
53+
" pip install kvikio-cu12 # For CUDA 12\n",
54+
" pip install kvikio-cu11 # For CUDA 11\n",
55+
" ```\n",
56+
"\n",
57+
"- **Conda/Mamba Installation** \n",
58+
" Follow the official [KvikIO installation guide](https://docs.rapids.ai/api/kvikio/nightly/install/) for Conda/Mamba installations.\n",
59+
"\n",
60+
"For convenience, we provide the cell below to install all the dependencies (please modify the cell based on your actual CUDA version, and please note that only CUDA 11 and CUDA 12 are supported for now)."
61+
]
62+
},
63+
{
64+
"cell_type": "code",
65+
"execution_count": null,
66+
"metadata": {},
67+
"outputs": [],
68+
"source": [
69+
"!python -c \"import monai\" || pip install -q \"monai-weekly[nibabel, pydicom, tqdm]\"\n",
70+
"!python -c \"import matplotlib\" || pip install -q matplotlib\n",
71+
"!python -c \"import torch_tensorrt\" || pip install torch_tensorrt\n",
72+
"!python -c \"import kvikio\" || pip install kvikio-cu12\n",
73+
"!python -c \"import ignite\" || pip install pytorch-ignite\n",
74+
"!python -c \"import pandas\" || pip install pandas\n",
75+
"!python -c \"import requests\" || pip install requests\n",
76+
"!python -c \"import fire\" || pip install fire\n",
77+
"!python -c \"import onnx\" || pip install nibaonnxbel\n",
78+
"%matplotlib inline"
79+
]
80+
},
81+
{
82+
"cell_type": "markdown",
83+
"metadata": {},
84+
"source": [
85+
"## Setup imports"
86+
]
87+
},
88+
{
89+
"cell_type": "code",
90+
"execution_count": null,
91+
"metadata": {},
92+
"outputs": [],
93+
"source": [
94+
"import os\n",
95+
"\n",
96+
"import torch\n",
97+
"import torch_tensorrt\n",
98+
"import matplotlib.pyplot as plt\n",
99+
"import monai\n",
100+
"from monai.config import print_config\n",
101+
"from monai.transforms import (\n",
102+
" EnsureChannelFirstd,\n",
103+
" EnsureTyped,\n",
104+
" LoadImaged,\n",
105+
" Orientationd,\n",
106+
" Spacingd,\n",
107+
" ScaleIntensityRanged,\n",
108+
" Compose\n",
109+
")\n",
110+
"from monai.data import Dataset,ThreadDataLoader\n",
111+
"import torch\n",
112+
"import numpy as np\n",
113+
"import copy\n",
114+
"\n",
115+
"print(f\"Torch-TensorRT version: {torch_tensorrt.__version__}.\")\n",
116+
"\n",
117+
"print_config()"
118+
]
119+
},
120+
{
121+
"cell_type": "markdown",
122+
"metadata": {},
123+
"source": [
124+
"## Prepare Test Data, Bundle, and TensorRT Model\n",
125+
"\n",
126+
"We provide a helper script, [`prepare_data.py`](./prepare_data.py), to simplify the setup process. This script performs the following tasks:\n",
127+
"\n",
128+
"- **Test Data**: Downloads and extracts the [Medical Segmentation Decathlon Task09 Spleen dataset](http://medicaldecathlon.com/).\n",
129+
"- **Bundle**: Downloads the required `spleen_ct_segmentation` bundle.\n",
130+
"- **TensorRT Model**: Exports the downloaded bundle model to a TensorRT engine-based TorchScript model. By default, the script exports the model using `fp16` precision, but you can modify it to use `fp32` precision if desired.\n",
131+
"\n",
132+
"The script automatically checks for existing data, bundles, and exported models before downloading or exporting. This ensures that repeated executions of the notebook do not result in redundant operations."
133+
]
134+
},
135+
{
136+
"cell_type": "code",
137+
"execution_count": null,
138+
"metadata": {},
139+
"outputs": [],
140+
"source": [
141+
"from utils import prepare_test_datalist, prepare_test_bundle, prepare_tensorrt_model\n",
142+
"\n",
143+
"root_dir = \".\"\n",
144+
"\n",
145+
"train_files = prepare_test_datalist(root_dir)\n",
146+
"bundle_path = prepare_test_bundle(bundle_dir=root_dir, bundle_name=\"spleen_ct_segmentation\")\n",
147+
"trt_model_name = \"model_trt.ts\"\n",
148+
"prepare_tensorrt_model(bundle_path, trt_model_name)"
149+
]
150+
},
151+
{
152+
"cell_type": "markdown",
153+
"metadata": {},
154+
"source": [
155+
"## Benchmark the end-to-end bundle inference\n",
156+
"\n",
157+
"A variable `benchmark_type` is defined to specify the type of benchmark to run. To have a fair comparison, each benchmark type should be run after restarting the notebook kernel.\n",
158+
"\n",
159+
"`benchmark_type` can be one of the following:\n",
160+
"\n",
161+
"- `\"original\"`: benchmark the original bundle inference.\n",
162+
"- `\"trt\"`: benchmark the TensorRT accelerated bundle inference.\n",
163+
"- `\"trt_gds\"`: benchmark the TensorRT accelerated bundle inference with GPU data loading and GPU transforms."
164+
]
165+
},
166+
{
167+
"cell_type": "code",
168+
"execution_count": 4,
169+
"metadata": {},
170+
"outputs": [],
171+
"source": [
172+
"benchmark_type = \"trt_gds\""
173+
]
174+
},
175+
{
176+
"cell_type": "markdown",
177+
"metadata": {},
178+
"source": [
179+
"A `TimerHandler` is defined to benchmark every part of the inference process.\n",
180+
"\n",
181+
"Please refer to `utils.py` for the implementation of `CUDATimer` and `TimerHandler`."
182+
]
183+
},
184+
{
185+
"cell_type": "code",
186+
"execution_count": 5,
187+
"metadata": {},
188+
"outputs": [],
189+
"source": [
190+
"from utils import TimerHandler, prepare_workflow, benchmark_workflow"
191+
]
192+
},
193+
{
194+
"cell_type": "markdown",
195+
"metadata": {},
196+
"source": [
197+
"### Benchmark the Original Bundle Inference\n",
198+
"\n",
199+
"In this section, the `workflow`runs several iterations to benchmark the latency."
200+
]
201+
},
202+
{
203+
"cell_type": "code",
204+
"execution_count": 6,
205+
"metadata": {},
206+
"outputs": [],
207+
"source": [
208+
"model_weight = os.path.join(bundle_path, \"models\", \"model.pt\")\n",
209+
"meta_config = os.path.join(bundle_path, \"configs\", \"metadata.json\")\n",
210+
"inference_config = os.path.join(bundle_path, \"configs\", \"inference.json\")\n",
211+
"\n",
212+
"override = {\n",
213+
" \"dataset#data\": [{\"image\": i} for i in train_files],\n",
214+
" \"output_postfix\": benchmark_type,\n",
215+
"}"
216+
]
217+
},
218+
{
219+
"cell_type": "code",
220+
"execution_count": 7,
221+
"metadata": {},
222+
"outputs": [],
223+
"source": [
224+
"if benchmark_type == \"original\":\n",
225+
"\n",
226+
" workflow = prepare_workflow(inference_config, meta_config, bundle_path, override)\n",
227+
" torch_timer = TimerHandler()\n",
228+
" benchmark_df = benchmark_workflow(workflow, torch_timer, benchmark_type)"
229+
]
230+
},
231+
{
232+
"cell_type": "markdown",
233+
"metadata": {},
234+
"source": [
235+
"### Benchmark the TensorRT Accelerated Bundle Inference\n",
236+
"In this part, the TensorRT accelerated model is loaded to the `workflow`. The updated `workflow` runs the same iterations as before to benchmark the latency difference. Since the TensorRT accelerated model cannot be loaded through the `CheckpointLoader` and don't have `amp` mode, disable the `CheckpointLoader` in the `initialize` of the `workflow` and the `amp` parameter in the `evaluator` of the `workflow` needs to be set to `False`."
237+
]
238+
},
239+
{
240+
"cell_type": "code",
241+
"execution_count": 8,
242+
"metadata": {},
243+
"outputs": [],
244+
"source": [
245+
"if benchmark_type == \"trt\":\n",
246+
" trt_model_path = os.path.join(bundle_path, \"models\", \"model_trt.ts\")\n",
247+
" trt_model = torch.jit.load(trt_model_path)\n",
248+
"\n",
249+
" override[\"load_pretrain\"] = False\n",
250+
" override[\"network_def\"] = trt_model\n",
251+
" override[\"evaluator#amp\"] = False\n",
252+
"\n",
253+
" workflow = prepare_workflow(inference_config, meta_config, bundle_path, override)\n",
254+
" trt_timer = TimerHandler()\n",
255+
" benchmark_df = benchmark_workflow(workflow, trt_timer, benchmark_type)"
256+
]
257+
},
258+
{
259+
"cell_type": "markdown",
260+
"metadata": {},
261+
"source": [
262+
"### Benchmarking TensorRT Accelerated Bundle Inference with GPU Data Loading and GPU-based Transforms\n",
263+
"\n",
264+
"In the previous section, the inference workflow utilized CPU-based transforms. In this section, we enhance performance by leveraging GPU acceleration:\n",
265+
"\n",
266+
"- **GPU Direct Storage (GDS)**: The `LoadImaged` transform enables GDS on `.nii` and `.dcm` files via specifying `to_gpu=True`.\n",
267+
"- **GPU-based Transforms**: After GDS, subsequent preprocessing transforms are executed directly on the GPU."
268+
]
269+
},
270+
{
271+
"cell_type": "code",
272+
"execution_count": 15,
273+
"metadata": {},
274+
"outputs": [],
275+
"source": [
276+
"transforms = Compose([\n",
277+
" LoadImaged(keys=\"image\", reader=\"NibabelReader\", to_gpu=False),\n",
278+
" EnsureTyped(keys=\"image\", device=torch.device(\"cuda:0\")),\n",
279+
" EnsureChannelFirstd(keys=\"image\"),\n",
280+
" Orientationd(keys=\"image\", axcodes=\"RAS\"),\n",
281+
" Spacingd(keys=\"image\", pixdim=[1.5, 1.5, 2.0], mode=\"bilinear\"),\n",
282+
" ScaleIntensityRanged(keys=\"image\", a_min=-57, a_max=164, b_min=0, b_max=1, clip=True),\n",
283+
"])\n",
284+
"\n",
285+
"dataset = Dataset(data=[{\"image\": i} for i in train_files], transform=transforms)\n",
286+
"dataloader = ThreadDataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)"
287+
]
288+
},
289+
{
290+
"cell_type": "code",
291+
"execution_count": null,
292+
"metadata": {},
293+
"outputs": [],
294+
"source": [
295+
"if benchmark_type == \"trt_gds\":\n",
296+
"\n",
297+
" trt_model_path = os.path.join(bundle_path, \"models\", \"model_trt.ts\")\n",
298+
" trt_model = torch.jit.load(trt_model_path)\n",
299+
" override = {\n",
300+
" \"output_postfix\": benchmark_type,\n",
301+
" \"load_pretrain\": False,\n",
302+
" \"network_def\": trt_model,\n",
303+
" \"evaluator#amp\": False,\n",
304+
" \"preprocessing\": transforms,\n",
305+
" \"dataset\": dataset,\n",
306+
" \"dataloader\": dataloader,\n",
307+
" }\n",
308+
"\n",
309+
" workflow = prepare_workflow(inference_config, meta_config, bundle_path, override)\n",
310+
" trt_gpu_trans_timer = TimerHandler()\n",
311+
" benchmark_df = benchmark_workflow(workflow, trt_gpu_trans_timer, benchmark_type)"
312+
]
313+
}
314+
],
315+
"metadata": {
316+
"kernelspec": {
317+
"display_name": "kvikio_env",
318+
"language": "python",
319+
"name": "python3"
320+
},
321+
"language_info": {
322+
"codemirror_mode": {
323+
"name": "ipython",
324+
"version": 3
325+
},
326+
"file_extension": ".py",
327+
"mimetype": "text/x-python",
328+
"name": "python",
329+
"nbconvert_exporter": "python",
330+
"pygments_lexer": "ipython3",
331+
"version": "3.10.14"
332+
}
333+
},
334+
"nbformat": 4,
335+
"nbformat_minor": 4
336+
}

0 commit comments

Comments
 (0)