wehub-resource-sync
342 lines
11 KiB
Plaintext
342 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#\n",
|
|
"# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n",
|
|
"#\n",
|
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
|
"# you may not use this file except in compliance with the License.\n",
|
|
"# You may obtain a copy of the License at\n",
|
|
"#\n",
|
|
"# http://www.apache.org/licenses/LICENSE-2.0\n",
|
|
"#\n",
|
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
|
"# See the License for the specific language governing permissions and\n",
|
|
"# limitations under the License.\n",
|
|
"#"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Check the TensorRT version"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!python3 -c 'import tensorrt; print(\"TensorRT version: {}\".format(tensorrt.__version__))'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Prepare the input image and ONNX model file"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!python3 /workspace/TensorRT/quickstart/SemanticSegmentation/export.py"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Build TensorRT engine from the ONNX model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!trtexec --onnx=fcn-resnet101.onnx --saveEngine=fcn-resnet101.engine --optShapes=input:1x3x1026x1282 --stronglyTyped"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Import required modules"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import numpy as np\n",
|
|
"import os\n",
|
|
"import ctypes\n",
|
|
"from cuda.bindings import runtime as cudart\n",
|
|
"import tensorrt as trt\n",
|
|
"\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"from PIL import Image\n",
|
|
"\n",
|
|
"TRT_LOGGER = trt.Logger()\n",
|
|
"\n",
|
|
"assert cudart.cudaSetDevice(0) == (cudart.cudaError_t.cudaSuccess,)\n",
|
|
"\n",
|
|
"# Filenames of TensorRT plan file and input/output images.\n",
|
|
"engine_file = \"/workspace/fcn-resnet101.engine\"\n",
|
|
"input_file = \"/workspace/input.ppm\"\n",
|
|
"output_file = \"/workspace/output.ppm\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Utilities for input / output processing"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# For torchvision models, input images are loaded in to a range of [0, 1] and\n",
|
|
"# normalized using mean = [0.485, 0.456, 0.406] and stddev = [0.229, 0.224, 0.225].\n",
|
|
"def preprocess(image):\n",
|
|
" # Mean normalization\n",
|
|
" mean = np.array([0.485, 0.456, 0.406]).astype('float32')\n",
|
|
" stddev = np.array([0.229, 0.224, 0.225]).astype('float32')\n",
|
|
" data = (np.asarray(image).astype('float32') / float(255.0) - mean) / stddev\n",
|
|
" # Switch from HWC to to CHW order\n",
|
|
" return np.moveaxis(data, 2, 0)\n",
|
|
"\n",
|
|
"def postprocess(data):\n",
|
|
" num_classes = 21\n",
|
|
" # create a color palette, selecting a color for each class\n",
|
|
" palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])\n",
|
|
" colors = np.array([palette*i%255 for i in range(num_classes)]).astype(\"uint8\")\n",
|
|
" # plot the segmentation predictions for 21 classes in different colors\n",
|
|
" img = Image.fromarray(data.astype('uint8'), mode='P')\n",
|
|
" img.putpalette(colors)\n",
|
|
" return img\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Load TensorRT engine"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Deserialize the TensorRT engine from specified plan file. "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def load_engine(engine_file_path):\n",
|
|
" assert os.path.exists(engine_file_path)\n",
|
|
" print(\"Reading engine from file {}\".format(engine_file_path))\n",
|
|
" with open(engine_file_path, \"rb\") as f, trt.Runtime(TRT_LOGGER) as runtime:\n",
|
|
" return runtime.deserialize_cuda_engine(f.read())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Inference pipeline"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Starting with a deserialized engine, TensorRT inference pipeline consists of the following steps:\n",
|
|
"- Create an execution context and specify input shape (based on the image dimensions for inference).\n",
|
|
"- Allocate CUDA device memory for input and output.\n",
|
|
"- Allocate CUDA page-locked host memory to efficiently copy back the output.\n",
|
|
"- Transfer the processed image data into input memory using asynchronous host-to-device CUDA copy.\n",
|
|
"- Kickoff the TensorRT inference pipeline using the asynchronous execute API.\n",
|
|
"- Transfer the segmentation output back into pagelocked host memory using device-to-host CUDA copy.\n",
|
|
"- Synchronize the stream used for data transfers and inference execution to ensure all operations are completes.\n",
|
|
"- Finally, write out the segmentation output to an image file for visualization."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def infer(engine, input_file, output_file):\n",
|
|
" print(\"Reading input image from file {}\".format(input_file))\n",
|
|
" with Image.open(input_file) as img:\n",
|
|
" input_image = preprocess(img)\n",
|
|
" image_width = img.width\n",
|
|
" image_height = img.height\n",
|
|
"\n",
|
|
" with engine.create_execution_context() as context:\n",
|
|
" input_buffers = {}\n",
|
|
" input_memories = {}\n",
|
|
"\n",
|
|
" # Allocate host and device buffers\n",
|
|
" tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]\n",
|
|
" for tensor in tensor_names:\n",
|
|
" size = trt.volume(context.get_tensor_shape(tensor))\n",
|
|
" dtype = trt.nptype(engine.get_tensor_dtype(tensor))\n",
|
|
"\n",
|
|
" if engine.get_tensor_mode(tensor) == trt.TensorIOMode.INPUT:\n",
|
|
" context.set_input_shape(tensor, (1, 3, image_height, image_width))\n",
|
|
" input_buffers[tensor] = np.ascontiguousarray(input_image)\n",
|
|
" err, input_memories[tensor] = cudart.cudaMalloc(input_image.nbytes)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
" context.set_tensor_address(tensor, input_memories[tensor])\n",
|
|
" else:\n",
|
|
" err, output_buffer_ptr = cudart.cudaMallocHost(size * dtype().itemsize)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
" pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))\n",
|
|
" output_buffer = np.ctypeslib.as_array(ctypes.cast(output_buffer_ptr, pointer_type), (size,))\n",
|
|
"\n",
|
|
" err, output_memory = cudart.cudaMalloc(output_buffer.nbytes)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
" context.set_tensor_address(tensor, output_memory)\n",
|
|
"\n",
|
|
" err, stream = cudart.cudaStreamCreate()\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
"\n",
|
|
" # Transfer input data to the GPU for all input tensors\n",
|
|
" for tensor_name, input_buffer in input_buffers.items():\n",
|
|
" input_memory = input_memories[tensor_name]\n",
|
|
" err, = cudart.cudaMemcpyAsync(input_memory, input_buffer.ctypes.data, input_buffer.nbytes,\n",
|
|
" cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
"\n",
|
|
" # Run inference\n",
|
|
" context.execute_async_v3(stream)\n",
|
|
"\n",
|
|
" # Transfer prediction output from the GPU.\n",
|
|
" err, = cudart.cudaMemcpyAsync(output_buffer.ctypes.data, output_memory, output_buffer.nbytes,\n",
|
|
" cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
" # Synchronize the stream\n",
|
|
" err, = cudart.cudaStreamSynchronize(stream)\n",
|
|
" assert err == cudart.cudaError_t.cudaSuccess\n",
|
|
"\n",
|
|
" output_d64 = np.array(output_buffer, dtype=np.int64)\n",
|
|
" np.savetxt('test.out', output_d64.astype(int), fmt='%i', delimiter=' ', newline=' ')\n",
|
|
"\n",
|
|
" with postprocess(np.reshape(output_buffer, (image_height, image_width))) as img:\n",
|
|
" print(\"Writing output image to file {}\".format(output_file))\n",
|
|
" img.convert('RGB').save(output_file, \"PPM\")\n",
|
|
"\n",
|
|
" # cleanup cuda resources for all input tensors\n",
|
|
" for input_memory in input_memories.values():\n",
|
|
" cudart.cudaFree(input_memory)\n",
|
|
" cudart.cudaFree(output_memory)\n",
|
|
" cudart.cudaFreeHost(output_buffer_ptr)\n",
|
|
" cudart.cudaStreamDestroy(stream)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Plot input image"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"plt.imshow(Image.open(input_file))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Run inference"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"Running TensorRT inference for FCN-ResNet101\")\n",
|
|
"with load_engine(engine_file) as engine:\n",
|
|
" infer(engine, input_file, output_file)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Plot segmentation output"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"plt.imshow(Image.open(output_file))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"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.10.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
}
|