Skip to content

README.md

Latest commit

 

History

History
42 lines (31 loc) · 1.62 KB

README.md

File metadata and controls

42 lines (31 loc) · 1.62 KB
Raw

TensorRT plugin for 3D GridSample Operator

At current stage, TensorRT(up to version 8.6.1) does not support 3D GridSample operator.

This plugin is a custom implementation of the 3D GridSample operator for TensorRT. It is inspired by the GridSample operator from PyTorch, and the code structure is inspired by project onnxparser-trt-plugin-sample.

Installation

  1. Let cmake find nvcc
export PATH=/usr/local/cuda/bin:$PATH
  1. Build the plugin with the following commands:
mkdir build && cd build
cmake .. -DTensorRT_ROOT=/usr/local/tensorrt
make

Usage

for python code (only on Linux platform), load the plugin with:

import ctypes
success = ctypes.CDLL("build/libgrid_sample_3d_plugin.so", mode = ctypes.RTLD_GLOBAL)

see test_grid_sample3d.py for more details.

Tensor precision

The plugin supports both TensorRT FLOAT and HALF formats. The image input, grid input, and output must use the same format. A newly created plugin defaults to HALF, preserving the low-memory inference behavior used by existing consumers; TensorRT selects FLOAT when the network tensors are explicitly FLOAT. Execution dispatch uses the datatype from the current TensorRT execution descriptors, so a dynamic-shape profile cannot leave enqueue using a stale datatype from an earlier configuration.

The TensorRT integration test checks both supported precisions, output type propagation, and rejection of mixed image/grid formats.