github-featured/nvidia--tensorrt
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

TensorRT 11.0 OSS Release (#4799)

Browse Source 
Signed-off-by: Kevin Chen <kevinch@nvidia.com>
This commit is contained in:
Kevin Chen
2026-06-02 15:04:59 -07:00
committed by GitHub
parent 5302b288c5
commit 771d8efdf7
439 changed files with 9748 additions and 39439 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+16 -2
View File
@@ -1,7 +1,21 @@
# TensorRT OSS Release Changelog
## 10.16.1 GA - 2026-4-13
## 11.0 GA - 2026-6-2
- General
- As a new major version bump, TensorRT 11.0 brings many enhancements for its users. For full information, see the [release notes.](https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/release-notes-11/11.0.0.html)
- This is a bugfix release with no major new features. See the [release notes](https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/release-notes-10/10.16.1.html) for more details.
- Demos
- The BERT demo has been removed.
- Samples
- Migrated `samples/python/python_plugin` and `samples/python/onnx_custom_plugin/plugin` to use V3 plugins
- Plugins
- Added a new Faster Air Top-K plugin
- Removed deprecated plugins: `batchTile`, `clip`, `coordConvAC`, `cropAndResize`, `gelu`, `leakyRelu`, `normalize`, `singleStepLSTM`, `specialSlice`, `split`, `nms`, and `proposal`
- Parsers
- Added support for `Swish` operator
- Added support for custom operators `TRT_Attention`, `TRT_MoE`, and `TRT_KVCacheUpdate`
## 10.16 GA - 2026-3-24
CMakeLists.txt
+76 -94
View File
@@ -15,7 +15,7 @@
# limitations under the License.
#
cmake_minimum_required(VERSION 3.13 FATAL_ERROR)
cmake_minimum_required(VERSION 3.31 FATAL_ERROR)
include(cmake/modules/set_ifndef.cmake)
include(cmake/modules/find_library_create_target.cmake)
list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules)
@@ -24,13 +24,10 @@ set_ifndef(TRT_LIB_DIR ${CMAKE_BINARY_DIR})
set_ifndef(TRT_OUT_DIR ${CMAKE_BINARY_DIR})
# Converts Windows paths
if(CMAKE_VERSION VERSION_LESS 3.20)
file(TO_CMAKE_PATH "${TRT_LIB_DIR}" TRT_LIB_DIR)
file(TO_CMAKE_PATH "${TRT_OUT_DIR}" TRT_OUT_DIR)
else()
cmake_path(SET TRT_LIB_DIR ${TRT_LIB_DIR})
cmake_path(SET TRT_OUT_DIR ${TRT_OUT_DIR})
endif()
cmake_path(SET TRT_LIB_DIR ${TRT_LIB_DIR})
cmake_path(SET TRT_OUT_DIR ${TRT_OUT_DIR})
cmake_path(SET TRT_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/include)
# Required to export symbols to build *.libs
if(WIN32)
@@ -48,7 +45,7 @@ else()
set(STATIC_LIB_EXT "a")
endif()
file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/include/NvInferVersion.h" VERSION_STRINGS REGEX "#define TRT_.*_ENTERPRISE")
file(STRINGS "${TRT_INCLUDE_DIR}/NvInferVersion.h" VERSION_STRINGS REGEX "#define TRT_.*_ENTERPRISE")
foreach(TYPE MAJOR MINOR PATCH BUILD)
string(REGEX MATCH "TRT_${TYPE}_ENTERPRISE [0-9]+" TRT_TYPE_STRING ${VERSION_STRINGS})
@@ -65,45 +62,16 @@ if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
endif()
set(CMAKE_SKIP_BUILD_RPATH True)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
# CUDA targets
set(DEFAULT_CUDA_VERSION 13.2)
set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION})
message(STATUS "CUDA version set to ${CUDA_VERSION}")
if (DEFINED GPU_ARCHS)
message(STATUS "GPU_ARCHS defined as ${GPU_ARCHS}. Generating CUDA code for SM ${GPU_ARCHS}")
separate_arguments(GPU_ARCHS)
foreach(SM IN LISTS GPU_ARCHS)
list(APPEND CMAKE_CUDA_ARCHITECTURES "${SM}")
endforeach()
else()
list(APPEND CMAKE_CUDA_ARCHITECTURES 75 80 86 87 89 90)
if(CUDA_VERSION VERSION_GREATER_EQUAL 12.8)
list(APPEND CMAKE_CUDA_ARCHITECTURES 100 120)
endif()
if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0)
if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90 100 120)
# SM 110 (Thor) requires CUDA 13.0+.
if(DEFINED CUDA_VERSION AND CUDA_VERSION VERSION_GREATER_EQUAL "13.0")
list(APPEND CMAKE_CUDA_ARCHITECTURES 110)
endif()
message(STATUS "GPU_ARCHS is not defined. Generating CUDA code for default SMs: ${CMAKE_CUDA_ARCHITECTURES}")
endif()
set(BERT_GENCODES)
# Generate SASS for each architecture
foreach(arch ${CMAKE_CUDA_ARCHITECTURES})
if (${arch} GREATER_EQUAL 75 AND NOT ${arch} EQUAL 110)
set(BERT_GENCODES "${BERT_GENCODES} -gencode arch=compute_${arch},code=sm_${arch}")
endif()
set(GENCODES "${GENCODES} -gencode arch=compute_${arch},code=sm_${arch}")
endforeach()
# Generate PTX for the last architecture in the list.
list(GET CMAKE_CUDA_ARCHITECTURES -1 LATEST_SM)
set(GENCODES "${GENCODES} -gencode arch=compute_${LATEST_SM},code=compute_${LATEST_SM}")
if (${LATEST_SM} GREATER_EQUAL 75 AND NOT ${arch} EQUAL 110)
set(BERT_GENCODES "${BERT_GENCODES} -gencode arch=compute_${LATEST_SM},code=compute_${LATEST_SM}")
endif()
message(STATUS "Generating CUDA code for SMs: ${CMAKE_CUDA_ARCHITECTURES}")
project(TensorRT
LANGUAGES CXX CUDA
@@ -124,6 +92,7 @@ option(BUILD_PARSERS "Build TensorRT parsers" ON)
option(BUILD_SAMPLES "Build TensorRT samples" ON)
option(BUILD_SAFE_SAMPLES "Build TensorRT safety samples" OFF)
option(TRT_SAFETY_INFERENCE_ONLY "Build only the safety inference components (no safety builders)" OFF)
option(TRT_BUILD_TESTING "Build gtests for TensorRT components" OFF)
############################################################################################
# Early dependency discovery
@@ -141,14 +110,8 @@ else()
endif()
endif()
## find_package(CUDA) is broken for cross-compilation. Enable CUDA language instead.
if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
find_package(CUDA ${CUDA_VERSION} REQUIRED)
endif()
include_directories(
${CUDA_INCLUDE_DIRS}
)
find_package(CUDAToolkit REQUIRED)
message(STATUS "CUDA version: ${CUDAToolkit_VERSION}")
############################################################################################
# Safety runtime libraries (libnvinfer_safe) used by safety samples and
@@ -204,8 +167,6 @@ if(BUILD_SAFE_SAMPLES OR TRT_SAFETY_INFERENCE_ONLY)
target_link_options(TRTSAFE::nvinfer_safe_shared INTERFACE LINKER:--unresolved-symbols=ignore-in-shared-libs)
target_link_options(TRTSAFE::nvinfer_safe_debug INTERFACE LINKER:--unresolved-symbols=ignore-in-shared-libs)
endif()
# Enable unified builder safety features when building safety samples or in inference-only mode.
add_compile_definitions(ENABLE_UNIFIED_BUILDER=1)
endif()
# OSS safety inference-only mode: require safety samples and disable enterprise
@@ -223,8 +184,8 @@ if(TRT_SAFETY_INFERENCE_ONLY)
set(BUILD_SAMPLES OFF CACHE BOOL "" FORCE)
# Add CUDA library directory early so all samples can find it
if(CUDA_TOOLKIT_ROOT_DIR)
link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib ${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/lib)
if(CUDAToolkit_LIBRARY_DIR)
link_directories(${CUDAToolkit_LIBRARY_DIR})
endif()
# Interface target for safety samples in inference-only mode.
@@ -240,10 +201,10 @@ if(TRT_SAFETY_INFERENCE_ONLY)
target_link_libraries(trt_global_definitions INTERFACE dl rt)
endif()
target_include_directories(trt_global_definitions INTERFACE
${CMAKE_CURRENT_SOURCE_DIR}/include
${TRT_INCLUDE_DIR}
${CMAKE_CURRENT_SOURCE_DIR}/samples/common
${CMAKE_CURRENT_SOURCE_DIR}/shared
${CUDA_INCLUDE_DIRS}
${CUDAToolkit_INCLUDE_DIRS}
)
target_compile_options(trt_global_definitions INTERFACE
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
@@ -290,12 +251,6 @@ message(STATUS "Protobuf version set to ${PROTOBUF_VERSION}")
if (BUILD_PLUGINS OR BUILD_PARSERS)
include(third_party/protobuf.cmake)
endif()
if(NOT CUB_ROOT_DIR)
if (CUDA_VERSION VERSION_LESS 11.0)
set(CUB_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/cub CACHE STRING "directory of CUB installation")
endif()
endif()
if(BUILD_PARSERS)
configure_protobuf(${PROTOBUF_VERSION})
endif()
@@ -320,36 +275,12 @@ endif()
find_library_create_target(nvinfer ${nvinfer_lib_name} SHARED "${TRT_LIB_DIR}")
if (DEFINED USE_CUGFX)
find_library(CUDART_LIB cugfx_dll HINTS ${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib lib/x64 lib64)
else()
# DriveOS platforms use cudart.so instead of cudart_static. This isn't the most sophisticated check, but it's correct.
if(CUDA_VERSION VERSION_GREATER_EQUAL 12.0)
set(CUDART_LIB_NAME cudart_static)
set(CMAKE_CUDA_RUNTIME_LIBRARY "static")
else()
set(CUDART_LIB_NAME cudart)
set(CMAKE_CUDA_RUNTIME_LIBRARY "shared")
endif()
# SafeCUDA (QNX-Safe) does not ship libcudadevrt. When CMAKE_CUDA_RUNTIME_LIBRARY
# is "shared" or "static", CMake automatically links -lcudadevrt for targets with
# CUDA sources, which breaks QNX-Safe cross-compilation. Setting it to "None"
# disables automatic CUDA runtime linking; cudart is still linked manually via
# trt_global_definitions.
if(TRT_SAFETY_INFERENCE_ONLY AND CMAKE_SYSTEM_NAME STREQUAL "QNX")
set(CMAKE_CUDA_RUNTIME_LIBRARY "None")
endif()
find_library(CUDART_LIB ${CUDART_LIB_NAME} HINTS ${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib lib/x64 lib64)
endif()
set(CMAKE_CUDA_RUNTIME_LIBRARY "static" CACHE STRING "")
if (NOT MSVC)
find_library(RT_LIB rt)
endif()
set(CUDA_LIBRARIES ${CUDART_LIB})
############################################################################################
if(NOT MSVC)
@@ -376,11 +307,62 @@ else()
find_library_create_target(${nvonnxparser_lib_name} ${nvonnxparser_lib_name} SHARED "${HINT_PATHS}")
endif()
# Samples:
# - BUILD_SAMPLES controls the regular (enterprise/OSS) samples via
# samples/CMakeLists.txt.
# - BUILD_SAFE_SAMPLES controls the safety samples (builder + infer parts).
# Both can be enabled at the same time if desired.
if(NOT TARGET trt_global_definitions)
add_library(trt_global_definitions INTERFACE)
target_include_directories(trt_global_definitions INTERFACE ${CUDAToolkit_INCLUDE_DIRS})
endif()
add_library(tensorrt_headers INTERFACE)
target_include_directories(tensorrt_headers INTERFACE ${TRT_INCLUDE_DIR})
# Samples
if(BUILD_SAMPLES OR BUILD_SAFE_SAMPLES)
set(TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW ON)
# Map OSS option names to the internal names used by samples/CMakeLists.txt.
set(TRT_BUILD_SAMPLES ${BUILD_SAMPLES})
set(TRT_BUILD_TRTEXEC ${BUILD_SAMPLES})
set(TRT_BUILD_ONNX_PARSER ${BUILD_PARSERS})
set(TRT_BUILD_PLUGINS ${BUILD_PLUGINS})
# Set defaults for specific feature enablement for samples.
option(TRT_BUILD_ENABLE_DLA "Build TensorRT with DLA features enabled." OFF)
option(TRT_BUILD_ENABLE_UNIFIED_BUILDER "Build TensorRT with unified builder (safety) features enabled." ${BUILD_SAFE_SAMPLES})
option(TRT_BUILD_WINML "Build TensorRT with WinML support." OFF)
option(TRT_BUILD_ENABLE_MULTIDEVICE "Build TensorRT with multi-device support." OFF)
set(TRT_BUILD_SAMPLES_LINK_STATIC_TRT OFF CACHE INTERNAL "")
find_library(NVINFER_LIB nvinfer PATHS ${TRT_LIB_DIR} REQUIRED)
add_library(tensorrt INTERFACE IMPORTED)
target_link_libraries(tensorrt INTERFACE ${NVINFER_LIB})
target_include_directories(tensorrt INTERFACE ${TRT_INCLUDE_DIR})
target_include_directories(${nvonnxparser_lib_name} INTERFACE ${TRT_INCLUDE_DIR})
add_subdirectory(shared)
include(InstallUtils)
if (TRT_BUILD_TESTING)
find_package(GTest QUIET)
if (GTest_FOUND)
if (NOT TARGET gtest_main)
add_library(gtest_main ALIAS GTest::gtest_main)
endif()
else()
include(FetchContent)
FetchContent_Declare(
googletest
GIT_REPOSITORY https://github.com/google/googletest.git
GIT_TAG v1.14.0
)
FetchContent_MakeAvailable(googletest)
endif()
set(TRT_GTEST_DISCOVERY_MODE PRE_TEST CACHE STRING "gtest discovery mode.")
endif()
# Route sample binaries to TRT_OUT_DIR.
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${TRT_OUT_DIR})
add_subdirectory(samples)
endif()
README.md
+24 -22
View File
@@ -1,20 +1,22 @@
[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Documentation](https://img.shields.io/badge/TensorRT-documentation-brightgreen.svg)](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html) [![Roadmap](https://img.shields.io/badge/Roadmap-Q1_2026-brightgreen.svg)](documents/tensorrt_roadmap_2026q1.pdf)
[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Documentation](https://img.shields.io/badge/TensorRT-documentation-brightgreen.svg)](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html) [![Roadmap](https://img.shields.io/badge/Roadmap-Q3_2026-brightgreen.svg)](documents/tensorrt_roadmap_2026q3.pdf)
# :mega::mega: Announcement :mega::mega:
TensorRT 11.0 is coming soon in 2026 Q2 with powerful new capabilities designed to accelerate your AI inference workflows. With this major version bump, TensorRT's API will be streamlined and a few legacy features will be removed.
TensorRT 11.0 is now released with powerful new capabilities designed to accelerate your AI inference workflows. With this major version bump, TensorRT's API has been streamlined and a few legacy features have been removed.
We recommend migrating early for the following features:
- Weakly-typed networks and related APIs will be removed, replaced by [Strongly Typed Networks](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html#strongly-typed-networks).
- Implicit quantization and related APIs will be removed, replaced by [Explicit Quantization](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/work-quantized-types.html#explicit-quantization)
- IPluginV2 and related APIs will be removed, replaced by [IPluginV3](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/extending-custom-layers.html#migrating-v2-plugins-to-ipluginv3)
- TREX tool will be removed, replaced by [Nsight Deep Learning Designer](https://docs.nvidia.com/nsight-dl-designer/UserGuide/index.html#visualizing-a-tensorrt-engine)
- Python bindings for Python 3.9 and older versions will be removed starting TensorRT 10.16. RPM packages for RHEL/Rocky Linux 8 and RHEL/Rocky Linux 9 now depend on Python 3.12.
Below provides migration guides for the following features:
- Weakly-typed networks and related APIs have been removed, replaced by [Strongly Typed Networks](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html#strongly-typed-networks).
- Implicit quantization and related APIs have been removed, replaced by [Explicit Quantization](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/work-quantized-types.html#explicit-quantization)
- IPluginV2 and related APIs have been removed, replaced by [IPluginV3](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/extending-custom-layers.html#migrating-v2-plugins-to-ipluginv3)
- TREX tool has been removed, replaced by [Nsight Deep Learning Designer](https://docs.nvidia.com/nsight-dl-designer/UserGuide/index.html#visualizing-a-tensorrt-engine)
- Python bindings for Python 3.9 and older versions have been removed. RPM packages for RHEL/Rocky Linux 8 and RHEL/Rocky Linux 9 now depend on Python 3.12.
# TensorRT Open Source Software
This repository contains the Open Source Software (OSS) components of NVIDIA TensorRT. It includes the sources for TensorRT plugins and ONNX parser, as well as sample applications demonstrating usage and capabilities of the TensorRT platform. These open source software components are a subset of the TensorRT General Availability (GA) release with some extensions and bug-fixes.
- For step-by-step walkthroughs of the TensorRT import paths (ONNX, Torch-TensorRT, HuggingFace/Optimum, Network Definition API) with examples and tooling tips, see the [Import Workflows Guide](documents/import_workflows.md).
- For the per-model support matrix across import paths (LLM, encoder-NLP, vision, audio, diffusion, multimodal), see [Supported Models](documents/supported_models.md).
- For code contributions to TensorRT-OSS, please see our [Contribution Guide](CONTRIBUTING.md) and [Coding Guidelines](CODING-GUIDELINES.md).
- For a summary of new additions and updates shipped with TensorRT-OSS releases, please refer to the [Changelog](CHANGELOG.md).
- For business inquiries, please contact [researchinquiries@nvidia.com](mailto:researchinquiries@nvidia.com)
@@ -43,7 +45,7 @@ To build the TensorRT-OSS components, you will first need the following software
**TensorRT GA build**
- TensorRT v10.16.1.11
- TensorRT v11.0.0.114
- Available from direct download links listed below
**System Packages**
@@ -98,24 +100,24 @@ To build the TensorRT-OSS components, you will first need the following software
Else download and extract the TensorRT GA build from [NVIDIA Developer Zone](https://developer.nvidia.com) with the direct links below:
- [TensorRT 10.16.1.11 for CUDA 13.2, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz)
- [TensorRT 10.16.1.11 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz)
- [TensorRT 10.16.1.11 for CUDA 13.2, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/zip/TensorRT-10.16.1.11.Windows.amd64.cuda-13.2.zip)
- [TensorRT 10.16.1.11 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/zip/TensorRT-10.16.1.11.Windows.amd64.cuda-12.9.zip)
- [TensorRT 11.0.0.114 for CUDA 13.2, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst)
- [TensorRT 11.0.0.114 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst)
- [TensorRT 11.0.0.114 for CUDA 13.2, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/zip/TensorRT-Enterprise-11.0.0.114-Windows-amd64-cuda-13.2-Release-external.zip)
- [TensorRT 11.0.0.114 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/zip/TensorRT-Enterprise-11.0.0.114-Windows-amd64-cuda-12.9-Release-external.zip)
**Example: Ubuntu 22.04 on x86-64 with cuda-13.2**
```bash
cd ~/Downloads
tar -xvzf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz
export TRT_LIBPATH=`pwd`/TensorRT-10.16.1.11/lib
tar --zstd -xvf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst
export TRT_LIBPATH=`pwd`/TensorRT-11.0.0.114/lib
```
**Example: Windows on x86-64 with cuda-12.9**
```powershell
Expand-Archive -Path TensorRT-10.16.1.11.Windows.amd64.cuda-12.9.zip
$env:TRT_LIBPATH="$pwd\TensorRT-10.16.1.11\lib"
Expand-Archive -Path TensorRT-Enterprise-11.0.0.114-Windows-amd64-cuda-12.9-Release-external.zip
$env:TRT_LIBPATH="$pwd\TensorRT-11.0.0.114\lib"
```
## Setting Up The Build Environment
@@ -237,9 +239,9 @@ For Linux platforms, we recommend that you generate a docker container for build
- `BUILD_SAMPLES`: Specify if the samples should be built, for example [`ON`] | `OFF`.
- `BUILD_SAFE_SAMPLES`: Specify if safety samples should be built, for example [`ON`] | `OFF`.
- `TRT_SAFETY_INFERENCE_ONLY`: Specify if only build the safety inference components, for example [`ON`] | `OFF`. If turned ON, all other components will be turned OFF except `BUILD_SAFE_SAMPLES`.
- `GPU_ARCHS`: GPU (SM) architectures to target. By default we generate CUDA code for all major SMs. Specific SM versions can be specified here as a quoted space-separated list to reduce compilation time and binary size. Table of compute capabilities of NVIDIA GPUs can be found [here](https://developer.nvidia.com/cuda-gpus). Examples: - NVidia A100: `-DGPU_ARCHS="80"` - RTX 50 series: `-DGPU_ARCHS="120"` - Multiple SMs: `-DGPU_ARCHS="80 120"`
- `TRT_PLATFORM_ID`: Bare-metal build (unlike containerized cross-compilation). Currently supported options: `x86_64` (default).
- `TRT_BUILD_ENABLE_MULTIDEVICE`: Enable the multi-device sample (`sampleDistCollective`). Use `-DTRT_BUILD_ENABLE_MULTIDEVICE=ON` to build it; requires [NCCL](https://developer.nvidia.com/nccl/nccl-download) >= v2.19, < v3.0.
- `TRT_BUILD_TESTING` : Build gTests for samples. Requires [gtest](https://github.com/google/googletest) if available; otherwise fetches googletest at configure time.
## Building TensorRT DriveOS Samples
@@ -259,7 +261,7 @@ For Linux platforms, we recommend that you generate a docker container for build
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_dos_cross.toolchain -DCUDA_VERSION=11.4 -DGPU_ARCHS=87
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_dos_cross.toolchain -DCUDA_VERSION=11.4 -DCMAKE_CUDA_ARCHITECTURES=87
make -j$(nproc)
```
@@ -284,7 +286,7 @@ For Linux platforms, we recommend that you generate a docker container for build
export QNX_HOST=$QNX_BASE/host/linux/x86_64/
export QNX_TARGET=$QNX_BASE/target/qnx7/
export PATH=$PATH:$QNX_HOST/usr/bin
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DBUILD_SAFE_SAMPLES=OFF -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx.toolchain -DCUDA_VERSION=$CUDA_VERSION -DGPU_ARCHS=87
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DBUILD_SAFE_SAMPLES=OFF -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx.toolchain -DCUDA_VERSION=$CUDA_VERSION -DCMAKE_CUDA_ARCHITECTURES=87
make -j$(nproc)
```
@@ -303,7 +305,7 @@ For Linux platforms, we recommend that you generate a docker container for build
export PATH=$PATH:$QNX_HOST/usr/bin
export CUDA=cuda-$CUDA_VERSION
export CUDA_ROOT=/usr/local/cuda-safe-$CUDA_VERSION
cmake .. -DBUILD_SAMPLES=OFF -DBUILD_SAFE_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DTRT_SAFETY_INFERENCE_ONLY=ON -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx_safe.toolchain -DCUDA_VERSION=$CUDA_VERSION -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DGPU_ARCHS=87
cmake .. -DBUILD_SAMPLES=OFF -DBUILD_SAFE_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DTRT_SAFETY_INFERENCE_ONLY=ON -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx_safe.toolchain -DCUDA_VERSION=$CUDA_VERSION -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DCMAKE_CUDA_ARCHITECTURES=87
make -j$(nproc)
```
@@ -322,7 +324,7 @@ For Linux platforms, we recommend that you generate a docker container for build
export QNX_HOST=$QNX_BASE/host/linux/x86_64/
export QNX_TARGET=$QNX_BASE/target/qnx/
export PATH=$PATH:$QNX_HOST/usr/bin
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DBUILD_SAFE_SAMPLES=OFF -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx.toolchain -DCUDA_VERSION=$CUDA_VERSION -DGPU_ARCHS=110
cmake .. -DBUILD_SAMPLES=ON -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DBUILD_SAFE_SAMPLES=OFF -DCMAKE_CUDA_COMPILER=$CUDA_ROOT/bin/nvcc -DTRT_OUT_DIR=`pwd`/bin_dynamic_cross -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_qnx.toolchain -DCUDA_VERSION=$CUDA_VERSION -DCMAKE_CUDA_ARCHITECTURES=110
make -j$(nproc)
```
VERSION
+1 -1
View File
@@ -1 +1 @@
10.16.1.11
11.0.0.114
cmake/modules/BundleLibraries.cmake
+28 -3
View File
@@ -35,6 +35,11 @@ define_property(TARGET
BRIEF_DOCS "Fallback type used when a target's TYPE is UNKNOWN_LIBRARY."
)
define_property(TARGET
PROPERTY BUNDLE_VISITED_INTERFACES
BRIEF_DOCS "Interface libraries already traversed by __bundleRecursiveDeps for this target."
)
# Internal helper to prefix all messages with "[target_bundle_libraries]: ".
#
# \param mode The message mode to be passed to message(...)
@@ -62,6 +67,14 @@ endfunction()
# \param target_name The name of the target to unwrap.
# \param result_var The variable to store the unwrapped target name in.
function(unwrapAlias target_name result_var)
# Check cache first (keyed on the raw input string).
get_property(_cache_set GLOBAL PROPERTY _UNWRAP_ALIAS_CACHE_${target_name} SET)
if(_cache_set)
get_property(_cached GLOBAL PROPERTY _UNWRAP_ALIAS_CACHE_${target_name})
set(${result_var} ${_cached} PARENT_SCOPE)
return()
endif()
# First, try to unwrap common generator expressions that may wrap the target name.
string(REGEX MATCH "\\$<LINK_LIBRARY:WHOLE_ARCHIVE,([a-zA-Z0-9_.:]+)>" _ ${target_name})
if(TARGET ${CMAKE_MATCH_1})
@@ -78,15 +91,18 @@ function(unwrapAlias target_name result_var)
if(aliased_target)
# Recursively unwrap in case there are multiple levels
unwrapAlias(${aliased_target} unwrapped)
set(${result_var} ${unwrapped} PARENT_SCOPE)
set(_result ${unwrapped})
else()
# Not an alias, return the original name
set(${result_var} ${target_name} PARENT_SCOPE)
set(_result ${target_name})
endif()
else()
# Not a target at all, return the original name
set(${result_var} ${target_name} PARENT_SCOPE)
set(_result ${target_name})
endif()
set_property(GLOBAL PROPERTY _UNWRAP_ALIAS_CACHE_${target_name} ${_result})
set(${result_var} ${_result} PARENT_SCOPE)
endfunction()
# Internal function to retrieve the type of library for a given target.
@@ -191,6 +207,15 @@ function(__bundleRecursiveDeps mainLib linkVis)
target_bundle_libraries(${mainLib} ${linkVis} ${dep})
elseif(${depType} STREQUAL INTERFACE_LIBRARY)
# For interface libraries, we want to add all of the static libraries they may be pointing to, without the library itself (since it is not a static).
# Skip if we've already traversed this interface lib's deps for mainLib to avoid re-walking shared transitive subtrees.
get_target_property(_visited ${mainLib} BUNDLE_VISITED_INTERFACES)
if(NOT _visited)
set(_visited "")
endif()
if(${dep} IN_LIST _visited)
continue()
endif()
set_property(TARGET ${mainLib} APPEND PROPERTY BUNDLE_VISITED_INTERFACES ${dep})
get_target_property(interfaceLibs ${dep} INTERFACE_LINK_LIBRARIES)
__bundleRecursiveDeps(${mainLib} ${linkVis} ${interfaceLibs})
elseif(${depType} STREQUAL SHARED_LIBRARY)
cmake/modules/InstallImportedLibraries.cmake
+60
View File
@@ -19,6 +19,66 @@
# This is particularly useful for system libraries that use versioned symlinks
# (e.g., libfoo.so -> libfoo.so.1 -> libfoo.so.1.2.3).
# Copies imported shared library files (including versioned symlinks) into the
# build library output directory so tests using LD_LIBRARY_PATH find them without
# a cmake --install step. Runs at configure time via file(COPY).
#
# \param targets One or more CMake imported shared-library targets to copy.
# \param destination Destination directory (default: resolved CMAKE_LIBRARY_OUTPUT_DIRECTORY)
function(copyImportedLibrariesToBuildTree)
set(oneValueArgs DESTINATION)
set(multiValueArgs TARGETS)
cmake_parse_arguments(ARG "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
if(NOT ARG_TARGETS)
message(FATAL_ERROR "copyImportedLibrariesToBuildTree requires at least one target.")
endif()
# Resolve any $<CONFIG> generator expression in CMAKE_LIBRARY_OUTPUT_DIRECTORY.
# TRT uses single-config generators (Ninja/Makefiles), so CMAKE_BUILD_TYPE is always set.
if(ARG_DESTINATION)
set(_dest "${ARG_DESTINATION}")
elseif(CMAKE_BUILD_TYPE)
string(REPLACE "$<CONFIG>" "${CMAKE_BUILD_TYPE}" _dest "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}")
else()
string(REPLACE "$<CONFIG>/" "" _dest "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}")
endif()
foreach(target_name IN LISTS ARG_TARGETS)
if(NOT TARGET ${target_name})
message(FATAL_ERROR "Target ${target_name} does not exist.")
endif()
get_target_property(target_type ${target_name} TYPE)
if(NOT target_type MATCHES "SHARED_LIBRARY|UNKNOWN_LIBRARY")
message(FATAL_ERROR "Target ${target_name} is not a shared library (type: ${target_type})")
endif()
# IMPORTED_LOCATION may be unset for config-specific imported targets; fall back to
# IMPORTED_LOCATION_<CONFIG>.
get_target_property(target_loc ${target_name} IMPORTED_LOCATION)
if(NOT target_loc)
string(TOUPPER "${CMAKE_BUILD_TYPE}" _config_upper)
get_target_property(target_loc ${target_name} "IMPORTED_LOCATION_${_config_upper}")
endif()
if(NOT target_loc)
message(FATAL_ERROR "Target ${target_name} has no IMPORTED_LOCATION or IMPORTED_LOCATION_<CONFIG>.")
endif()
get_filename_component(target_dir "${target_loc}" DIRECTORY)
get_filename_component(target_base "${target_loc}" NAME_WE)
file(GLOB target_libs "${target_dir}/${target_base}${CMAKE_SHARED_LIBRARY_SUFFIX}*")
if(target_libs)
file(MAKE_DIRECTORY "${_dest}")
file(COPY ${target_libs} DESTINATION "${_dest}")
message(STATUS "Copied ${target_name} to build tree: ${_dest}")
else()
message(WARNING "copyImportedLibrariesToBuildTree: no libraries found for ${target_name} in ${target_dir}")
endif()
endforeach()
endfunction()
# Installs an imported library target with all its symlinks.
#
# \param targets One or more CMake targets to install. Targets must be a shared library (or an unknown library pointing to a shared library).
cmake/toolchains/cmake_aarch64-android.toolchain
-44
View File
@@ -1,44 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
set(CMAKE_SYSTEM_NAME Linux)
set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(CMAKE_C_COMPILER $ENV{AARCH64_ANDROID_CC})
set(CMAKE_CXX_COMPILER $ENV{AARCH64_ANDROID_CC})
set(CMAKE_C_FLAGS "$ENV{AARCH64_ANDROID_CFLAGS} -pie -fPIE" CACHE STRING "" FORCE)
set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS}" CACHE STRING "" FORCE)
set(CMAKE_C_COMPILER_TARGET aarch64-none-linux-android)
set(CMAKE_CXX_COMPILER_TARGET aarch64-none-linux-android)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib64)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcudart -lnvToolsExt -lculibos -lcudadevrt -llog)
set(TRT_PLATFORM_ID "aarch64-android")
cmake/toolchains/cmake_aarch64-native.toolchain
+2 -7
View File
@@ -20,8 +20,6 @@ set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(TRT_PLATFORM_ID "aarch64")
set(CUDA_PLATFORM_ID "sbsa-linux")
set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++)
@@ -34,9 +32,6 @@ set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
if(DEFINED CUDA_ROOT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
else()
set(CUDA_TOOLKIT_ROOT_DIR /usr/local/cuda CACHE STRING "CUDA ROOT dir")
if(NOT DEFINED CUDAToolkit_ROOT)
set(CUDAToolkit_ROOT /usr/local/cuda CACHE STRING "CUDA ROOT dir")
endif()
set(CUDA_INCLUDE_DIRS ${CUDA_TOOLKIT_ROOT_DIR}/include)
cmake/toolchains/cmake_aarch64.toolchain
-38
View File
@@ -19,14 +19,6 @@ set(CMAKE_SYSTEM_NAME Linux)
set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(TRT_PLATFORM_ID "aarch64")
set(CMAKE_FIND_LIBRARY_PREFIXES "lib")
set(CMAKE_FIND_LIBRARY_SUFFIXES .so)
if("$ENV{ARMSERVER}" AND "${CUDA_VERSION}" VERSION_GREATER_EQUAL 11.0)
set(CUDA_PLATFORM_ID "sbsa-linux")
else()
set(CUDA_PLATFORM_ID "aarch64-linux")
endif()
set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++)
@@ -40,36 +32,6 @@ set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
set(CUDA_ROOT /usr/local/cuda-${CUDA_VERSION}/targets/${CUDA_PLATFORM_ID} CACHE STRING "CUDA ROOT dir")
set(CUDNN_ROOT_DIR /pdk_files/cudnn)
set(BUILD_LIBRARY_ONLY 1)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(CMAKE_THREAD_LIBS_INIT "-lpthread")
set(CMAKE_HAVE_THREADS_LIBRARY 1)
set(CMAKE_USE_WIN32_THREADS_INIT 0)
set(CMAKE_USE_PTHREADS_INIT 1)
find_library(RT_LIB rt PATHS /usr/aarch64-linux-gnu/lib /usr/lib/aarch64-linux-gnu)
if(NOT RT_LIB)
find_file(RT_LIB librt.so PATHS /usr/aarch64-linux-gnu/lib /usr/lib/aarch64-linux-gnu)
if(NOT RT_LIB)
message(WARNING "librt.so not found in default paths")
endif()
endif()
message("RT_LIB: ${RT_LIB}")
# Use host nvcc
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lstdc++ -lm)
cmake/toolchains/cmake_aarch64_cross.toolchain
-22
View File
@@ -20,13 +20,8 @@ set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(TRT_PLATFORM_ID "aarch64")
set(CUDA_PLATFORM_ID "sbsa-linux")
set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++)
set(CMAKE_C_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILE_FEATURES cxx_std_17)
set(CMAKE_C_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
@@ -34,24 +29,7 @@ set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
set(CUDA_ROOT /usr/local/cuda/targets/${CUDA_PLATFORM_ID} CACHE STRING "CUDA ROOT dir")
set(CUDNN_LIB /usr/lib/aarch64-linux-gnu/libcudnn.so)
set(BUILD_LIBRARY_ONLY 1)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(RT_LIB /usr/aarch64-linux-gnu/lib/librt.so)
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm)
link_directories(${CUDA_ROOT}/lib)
cmake/toolchains/cmake_aarch64_dos_cross.toolchain
-22
View File
@@ -22,13 +22,8 @@ set(CMAKE_FIND_ROOT_PATH /usr/aarch64-linux-gnu/)
set(TRT_PLATFORM_ID "aarch64")
set(CUDA_PLATFORM_ID "aarch64-linux")
set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++)
set(CMAKE_C_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILE_FEATURES cxx_std_17)
set(CMAKE_C_FLAGS "" CACHE STRING "" FORCE)
set(CMAKE_CXX_FLAGS "" CACHE STRING "" FORCE)
@@ -39,22 +34,5 @@ set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
set(CUDA_ROOT /usr/local/cuda/targets/${CUDA_PLATFORM_ID} CACHE STRING "CUDA ROOT dir")
set(CUDNN_LIB /usr/lib/aarch64-linux-gnu/libcudnn.so)
set(BUILD_LIBRARY_ONLY 1)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm)
link_directories(${CUDA_ROOT}/lib)
cmake/toolchains/cmake_ppc64le.toolchain
-38
View File
@@ -1,38 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
set(CMAKE_SYSTEM_NAME Linux)
set(CMAKE_SYSTEM_PROCESSOR ppc64le)
set(CMAKE_C_COMPILER powerpc64le-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER powerpc64le-linux-gnu-g++)
set(CMAKE_AR /usr/bin/ar CACHE STRING "" FORCE)
set(CMAKE_C_COMPILER_TARGET powerpc64le-linux-gnu)
set(CMAKE_CXX_COMPILER_TARGET powerpc64le-linux-gnu)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_ROOT}/include -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
if(DEFINED CUDA_ROOT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
endif()
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(TRT_PLATFORM_ID "ppc64le")
cmake/toolchains/cmake_x64_win.toolchain
+3 -5
View File
@@ -22,11 +22,10 @@ set(CMAKE_C_COMPILER ${CC})
set(CMAKE_CXX_COMPILER ${CC})
if(DEFINED CUDA_TOOLKIT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT})
set(CUDAToolkit_ROOT ${CUDA_TOOLKIT})
endif()
set(CMAKE_CUDA_COMPILER ${CUDA_TOOLKIT_ROOT_DIR}\\bin\\nvcc.exe)
set(CMAKE_CUDA_COMPILER_ID "NVIDIA")
set(CMAKE_CUDA_COMPILER ${CUDAToolkit_ROOT}\\bin\\nvcc.exe)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
@@ -34,7 +33,6 @@ set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(NV_TOOLS ${NV_TOOLS})
set(W10_LIBRARY_SUFFIXES .lib .dll)
set(W10_CUDA_ROOT ${CUDA_TOOLKIT_ROOT_DIR})
set(W10_LINKER ${MSVC_COMPILER_DIR}\\bin\\amd64\\link)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_NVCC_COMPILER} CACHE STRING "" FORCE)
@@ -43,6 +41,6 @@ set(ADDITIONAL_PLATFORM_INCL_FLAGS "-I${MSVC_COMPILER_DIR}\\include -I${MSVC_COM
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${NV_TOOLS}\\ddk\\wddmv2\\official\\17134\\Lib\\10.0.17134.0\\um\\x64")
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${MSVC_COMPILER_DIR}\\lib\\amd64" )
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${MSVC_COMPILER_DIR}\\..\\ucrt\\lib\\x64")
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${W10_CUDA_ROOT}\\lib\\x64 cudart.lib")
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${CUDAToolkit_ROOT}\\lib\\x64 cudart.lib")
set(TRT_PLATFORM_ID "win10")
cmake/toolchains/cmake_x86_64.toolchain
+2 -4
View File
@@ -21,10 +21,8 @@ set(CMAKE_SYSTEM_PROCESSOR x86_64)
set(CMAKE_C_COMPILER gcc)
set(CMAKE_CXX_COMPILER g++)
if(DEFINED CUDA_ROOT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
if(NOT DEFINED CUDAToolkit_ROOT)
set(CUDAToolkit_ROOT /usr/local/cuda CACHE STRING "CUDA ROOT dir")
endif()
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(TRT_PLATFORM_ID "x86_64")
cmake/toolchains/cmake_x86_64_agnostic.toolchain
+2 -4
View File
@@ -21,10 +21,8 @@ set(CMAKE_SYSTEM_PROCESSOR x86_64)
set(CMAKE_C_COMPILER /opt/rh/devtoolset-8/root/usr/bin/gcc)
set(CMAKE_CXX_COMPILER /opt/rh/devtoolset-8/root/usr/bin/g++)
if(DEFINED CUDA_ROOT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
if(NOT DEFINED CUDAToolkit_ROOT)
set(CUDAToolkit_ROOT /usr/local/cuda CACHE STRING "CUDA ROOT dir")
endif()
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(TRT_PLATFORM_ID "x86_64")
demo/BERT/CMakeLists.txt
-72
View File
@@ -1,72 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
cmake_minimum_required(VERSION 3.12 FATAL_ERROR)
project(infer_c LANGUAGES CXX)
find_package(CUDA)
include(FetchContent)
FetchContent_Declare(
pybind11
GIT_REPOSITORY https://github.com/pybind/pybind11
GIT_TAG v2.2.3
)
FetchContent_GetProperties(pybind11)
if(NOT pybind11_POPULATED)
FetchContent_Populate(pybind11)
add_subdirectory(${pybind11_SOURCE_DIR} ${pybind11_BINARY_DIR})
endif()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-deprecated-declarations")
include($ENV{TRT_OSSPATH}/cmake/modules/set_ifndef.cmake)
set_ifndef(TRT_INC_DIR $ENV{TRT_OSSPATH}/include)
set_ifndef(TRT_LIB_DIR $ENV{TRT_LIBPATH})
set_ifndef(TRT_OUT_DIR $ENV{TRT_OSSPATH}/build/out)
include_directories(
infer_c
${CUDA_INCLUDE_DIRS}
${TRT_INC_DIR}
)
link_directories(
${TRT_OUT_DIR}
${TRT_LIB_DIR}
)
pybind11_add_module(infer_c
infer_c/infer_c.cpp
infer_c/logging.cpp
)
target_link_libraries(infer_c PRIVATE
${CUDA_LIBRARIES}
nvinfer
nvinfer_plugin
)
add_executable(perf
infer_c/perf.cpp
infer_c/logging.cpp
)
target_link_libraries(perf
${CUDA_LIBRARIES}
nvinfer
nvinfer_plugin
)
demo/BERT/README.md
-728
View File
@@ -1,728 +0,0 @@
# BERT Inference Using TensorRT [DEPRECATED]
> This demo has been deprecated since TensorRT 10.15 given the deprecation of the [bertQKVToContextPlugin](../../plugin/bertQKVToContextPlugin/README.md).
This subfolder of the BERT TensorFlow repository, tested and maintained by NVIDIA, provides scripts to perform high-performance inference using NVIDIA TensorRT.
## Table Of Contents
- [Model Overview](#model-overview)
- [Model Architecture](#model-architecture)
- [TensorRT Inference Pipeline](#tensorrt-inference-pipeline)
- [Version Info](#version-info)
- [Setup](#setup)
- [Requirements](#requirements)
- [Quick Start Guide](#quick-start-guide)
- [(Optional) Trying a different configuration](#optional-trying-a-different-configuration)
- [Advanced](#advanced)
- [Scripts and sample code](#scripts-and-sample-code)
- [Command-line options](#command-line-options)
- [TensorRT inference process](#tensorrt-inference-process)
- [Accuracy](#accuracy)
- [Evaluating Post-Training-Quantization INT8 accuracy](#evaluating-ptq-post-training-quantization-int8-accuracy-using-the-squad-dataset)
- [Evaluating Quantization-Aware-Training INT8 accuracy](#evaluating-qat-quantization-aware-training-int8-accuracy-using-the-squad-dataset)
- [Experimental](#experimental)
- [Variable sequence length](#variable-sequence-length)
- [Run command lines](#run-command-lines)
- [Sparsity with Quantization Aware Training](#sparsity-with-quantization-aware-training)
- [Megatron-LM for Question Answering](#megatron-lm-for-question-answering)
- [Performance](#performance)
- [Benchmarking](#benchmarking)
- [TensorRT inference benchmark](#tensorrt-inference-benchmark)
- [Results](#results)
- [Inference performance: NVIDIA A100](#inference-performance-nvidia-a100-40gb)
- [Inference performance: NVIDIA L4](#inference-performance-nvidia-l4)
- [Inference performance: NVIDIA L40S](#inference-performance-nvidia-l40s)
## Model overview
BERT, or Bidirectional Encoder Representations from Transformers, is a new method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks. This model is based on the [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) paper. NVIDIA's BERT is an optimized version of [Google's official implementation](https://github.com/google-research/bert), leveraging mixed precision arithmetic and Tensor Cores for faster inference times while maintaining target accuracy.
Other publicly available implementations of BERT include:
1. [NVIDIA PyTorch](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT)
2. [Hugging Face](https://github.com/huggingface/pytorch-pretrained-BERT)
3. [codertimo](https://github.com/codertimo/BERT-pytorch)
4. [gluon-nlp](https://github.com/dmlc/gluon-nlp/tree/master/scripts/bert)
5. [Google's official implementation](https://github.com/google-research/bert)
### Model architecture
BERT's model architecture is a multi-layer bidirectional Transformer encoder. Based on the model size, we have the following two default configurations of BERT:
| **Model** | **Hidden layers** | **Hidden unit size** | **Attention heads** | **Feed-forward filter size** | **Max sequence length** | **Parameters** |
| :--------: | :---------------: | :------------------: | :-----------------: | :--------------------------: | :---------------------: | :------------: |
| BERT-Base | 12 encoder | 768 | 12 | 4 x 768 | 512 | 110M |
| BERT-Large | 24 encoder | 1024 | 16 | 4 x 1024 | 512 | 330M |
Typically, the language model is followed by a few task-specific layers. The model used here includes layers for question answering.
### TensorRT Inference Pipeline
BERT inference consists of three main stages: tokenization, the BERT model, and finally a projection of the tokenized prediction onto the original text.
Since the tokenizer and projection of the final predictions are not nearly as compute-heavy as the model itself, we run them on the host. The BERT model is GPU-accelerated via TensorRT.
The tokenizer splits the input text into tokens that can be consumed by the model. For details on this process, see [this tutorial](https://mccormickml.com/2019/05/14/BERT-word-embeddings-tutorial/).
To run the BERT model in TensorRT, we construct the model using TensorRT APIs and import the weights from a pre-trained TensorFlow checkpoint from [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/models/bert_tf_ckpt_large_qa_squad2_amp_128). Finally, a TensorRT engine is generated and serialized to the disk. The various inference scripts then load this engine for inference.
Lastly, the tokens predicted by the model are projected back to the original text to get a final result.
### Version Info
The following software version configuration has been tested:
| Software | Version |
| -------- | ------- |
| Python | >=3.8 |
| TensorRT | 10.11 |
| CUDA | 12.9 |
## Setup
The following section lists the requirements that you need to meet in order to run the BERT model.
### Requirements
This demo BERT application can be run within the TensorRT OSS build container. If running in a different environment, following packages are required.
- [NGC CLI](https://ngc.nvidia.com/setup/installers/cli) - for downloading BERT checkpoints from NGC.
- PyPI Packages:
- [cuda-python](https://pypi.org/project/cuda-python/) (tested v13.0.1)
- [onnx](https://pypi.org/project/onnx) (tested v1.12.0)
- [tensorflow](https://pypi.org/project/tensorflow/) (tested v2.9.1)
- [torch](https://pypi.org/project/torch/) (tested v1.11.0)
- NVIDIA [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/), [Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Ampere](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/) based GPU.
## Quick Start Guide
1. Build and launch the container as described in [TensorRT OSS README](https://github.com/NVIDIA/TensorRT/blob/master/README.md).
**Note:** After this point, all commands should be run from within the container.
2. Verify TensorRT installation by printing the version:
For example:
```bash
python3 -c "import tensorrt as trt; print(trt.__version__)"
```
3. Download the SQuAD dataset and BERT checkpoints:
```bash
cd $TRT_OSSPATH/demo/BERT
```
Download SQuAD v1.1 training and dev dataset.
```bash
bash ./scripts/download_squad.sh
```
Download Tensorflow checkpoints for BERT large model with sequence length 128, fine-tuned for SQuAD v2.0.
```bash
bash scripts/download_model.sh
```
**Note:** Since the datasets and checkpoints are stored in the directory mounted from the host, they do _not_ need to be downloaded each time the container is launched.
**Warning:** In the event of encountering an error message stating, "Missing API key and missing Email Authentication. This command requires an API key or authentication via browser login", the recommended steps for resolution are as follows:
- Generate an API key by logging in https://ngc.nvidia.com/setup/api-key and copy the generated API key.
- Execute the command `ngc config set` in the docker and paste the copied API key into the prompt as directed.
Completing these steps should resolve the error you encountered and allow the command to proceed successfully.
4. Build a TensorRT engine. To build an engine, run the `builder.py` script. For example:
```bash
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/model.ckpt -o engines/bert_large_128.engine -b 1 -s 128 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1
```
This will build an engine with a maximum batch size of 1 (`-b 1`), and sequence length of 128 (`-s 128`) using mixed precision (`--fp16`) using the BERT Large SQuAD v2 FP16 Sequence Length 128 checkpoint (`-c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1`).
5. Run inference. Two options are provided for running the model.
a. `inference.py` script
This script accepts a passage and question and then runs the engine to generate an answer.
For example:
```bash
python3 inference.py -e engines/bert_large_128.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
b. `inference.ipynb` Jupyter Notebook
The Jupyter Notebook includes a passage and various example questions and allows you to interactively make modifications and see the outcome.
To launch the Jupyter Notebook from inside the container, run:
```bash
jupyter notebook --ip 0.0.0.0 inference.ipynb
```
Then, use your browser to open the link displayed. The link should look similar to: `http://127.0.0.1:8888/?token=<TOKEN>`
6. Run inference with CUDA Graph support.
A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option.
```bash
mkdir -p build; pushd build
cmake .. -DPYTHON_EXECUTABLE=$(which python)
make -j
popd
python3 inference_c.py -e engines/bert_large_128.engine --enable-graph -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
A separate C/C++ inference benchmark executable `perf` (compiled from `perf.cpp`) is provided to run inference benchmarks with CUDA Graph. The cmdline interface is the same as `perf.py` except for an extra `--enable_graph` option.
```bash
build/perf -e engines/bert_large_128.engine -b 1 -s 128 -w 100 -i 1000 --enable_graph
```
### (Optional) Trying a different configuration
If you would like to run another configuration, you can manually download checkpoints using the included script. For example, run:
```bash
bash scripts/download_model.sh base
```
to download a BERT Base model instead of the default BERT Large model.
To view all available model options, run:
```bash
bash scripts/download_model.sh -h
```
## Advanced
The following sections provide greater details on inference with TensorRT.
### Scripts and sample code
In the `root` directory, the most important files are:
- `builder.py` - Builds an engine for the specified BERT model
- `Dockerfile` - Container which includes dependencies and model checkpoints to run BERT
- `inference.ipynb` - Runs inference interactively
- `inference.py` - Runs inference with a given passage and question
- `perf.py` - Runs inference benchmarks
The `scripts/` folder encapsulates all the one-click scripts required for running various supported functionalities, such as:
- `build.sh` - Builds a Docker container that is ready to run BERT
- `launch.sh` - Launches the container created by the `build.sh` script.
- `download_model.sh` - Downloads pre-trained model checkpoints from NGC
- `inference_benchmark.sh` - Runs an inference benchmark and prints results
Other folders included in the `root` directory are:
- `helpers` - Contains helpers for tokenization of inputs
The `infer_c/` folder contains all the necessary C/C++ files required for CUDA Graph support.
- `bert_infer.h` - Defines necessary data structures for running BERT inference
- `infer_c.cpp` - Defines C/C++ interface using pybind11 that can be plugged into `inference_c.py`
- `perf.cpp` - Runs inference benchmarks. It is equivalent to `perf.py`, with an extra option `--enable_graph` to enable CUDA Graph support.
### Command-line options
To view the available parameters for each script, you can use the help flag (`-h`).
**Note:** In the builder scripts (`builder.py` and `builder_varseqlen.py`), the options `--use-deprecated-plugins` and `--use-v3-plugins` toggle the underlying implementation of the plugins used in demoBERT. They are mutually exclusive, and enabling either should not affect functionality, or performance. The `--use-deprecated-plugins` uses plugin versions that inherit from `IPluginV2DynamicExt`, while `--use-v3-plugins` uses plugin versions that inherit from `IPluginV3` classes.
If unspecified, `--use-deprecated-plugins` is used by default.
**Additional Note:** Using `--use-v3-plugins` is not recommended on Blackwell platforms (See [Platform support section](#hardware-platform-support)). Prefer the default path instead (`--use-deprecated-plugins`).
### TensorRT inference process
As mentioned in the [Quick Start Guide](#quick-start-guide), two options are provided for running inference:
1. The `inference.py` script which accepts a passage and a question and then runs the engine to generate an answer. Alternatively, this script can be used to run inference on the Squad dataset.
2. The `inference.ipynb` Jupyter Notebook which includes a passage and various example questions and allows you to interactively make modifications and see the outcome.
## Accuracy
### Evaluating PTQ (post-training quantization) Int8 Accuracy Using The SQuAD Dataset
1. Download Tensorflow checkpoints for a BERT Large FP16 SQuAD v2 model with a sequence length of 384:
```bash
bash scripts/download_model.sh large 384 v2
```
2. Build an engine:
**Turing and Ampere GPUs**
```bash
# QKVToContextPlugin and SkipLayerNormPlugin supported with INT8 I/O. To enable, use -imh and -iln builder flags respectively.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100 -iln -imh
```
**Xavier GPU**
```bash
# Only supports SkipLayerNormPlugin running with INT8 I/O. Use -iln builder flag to enable.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100 -iln
```
**Volta GPU**
```bash
# No support for QKVToContextPlugin or SkipLayerNormPlugin running with INT8 I/O. Don't specify -imh or -iln in builder flags.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100
```
This will build an engine with a maximum batch size of 1 (`-b 1`), calibration dataset squad (`--squad-json ./squad/train-v1.1.json`), calibration sentences number 100 (`--calib-num 100`), and sequence length of 384 (`-s 384`) using INT8 mixed precision computation where possible (`--int8 --fp16 --strict`).
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
```bash
python3 inference.py -e engines/bert_large_384_int8mix.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
### Evaluating QAT (quantization aware training) Int8 Accuracy Using The SQuAD Dataset
1. Download checkpoint for BERT Large FP16 SQuAD v1.1 model with sequence length of 384:
```bash
bash scripts/download_model.sh pyt v1_1
```
2. Build an engine:
**Turing and Ampere GPUs**
```bash
# QKVToContextPlugin and SkipLayerNormPlugin supported with INT8 I/O. To enable, use -imh and -iln builder flags respectively.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -iln -imh
```
**Xavier GPU**
```bash
# Only supports SkipLayerNormPlugin running with INT8 I/O. Use -iln builder flag to enable.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -iln
```
**Volta GPU**
```bash
# No support for QKVToContextPlugin or SkipLayerNormPlugin running with INT8 I/O. Don't specify -imh or -iln in builder flags.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx
```
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 384 (`-s 384`) using INT8 mixed precision computation where possible (`--int8 --fp16 --strict`).
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
```bash
python3 inference.py -e engines/bert_large_384_int8mix.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
## Experimental
### Variable sequence length
In our prior implementation, we used inputs padded to max length along with corresponding input masks to handle variable sequence length inputs in a batch. The padding results in some wasted computations which can be avoided by handling variable sequence length inputs natively. Now we have a new approach called the variable sequence length method. By concatenating each input id into a single long input id, and concatenating each input segment id into a single long segment id, TensorRT can know the exact starts and ends by providing an extra sequence length buffer that contains the start and end positions of each sequence. Now we can eliminate the wasted computation in the input paddings.
Note this is an experimental feature because we only support Xavier+ GPUs, also there is neither FP32 support nor INT8 PTQ calibration.
1. Download checkpoint for BERT Large FP16 SQuAD v1.1 model with sequence length of 384:
```bash
bash scripts/download_model.sh pyt v1_1
```
2. Build an engine:
**FP16 engine**
```bash
mkdir -p engines && python3 builder_varseqlen.py -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -o engines/bert_varseq_fp16.engine -b 1 -s 64 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt
```
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 64 (`-s 64`) using FP16 precision computation where possible (`--fp16`).
**INT8 engine**
```bash
mkdir -p engines && python3 builder_varseqlen.py -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -o engines/bert_varseq_int8.engine -b 1 -s 256 --int8 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt
```
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 256 (`-s 256`) using INT8 precision computation where possible (`--int8`).
3. Run inference
Evaluate the F1 score and exact match score using the squad dataset:
```bash
python3 inference_varseqlen.py -e engines/bert_varseq_int8.engine -s 256 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
Run the quesion and answer mode:
```bash
python3 inference_varseqlen.py -e engines/bert_varseq_int8.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256
```
4. Collect performance data
```bash
python3 perf_varseqlen.py -e engines/bert_varseq_int8.engine -b 1 -s 256
```
This will collect performance data run use batch size 1 (`-b 1`) and sequence length of 256 (`-s 256`).
5. Collect performance data with CUDA graph enabled
We can use the same `inference_c.py` and `build/perf` to collect performance data with cuda graph enabled. The command line is the same as run without variable sequence length.
### Sparsity with Quantization Aware Training
Fine-grained 2:4 structured sparsity support introduced in NVIDIA Ampere GPUs can produce significant performance gains in BERT inference. The network is first trained using dense weights, then fine-grained structured pruning is applied, and finally the remaining non-zero weights are fine-tuned with additional training steps. This method results in virtually no loss in inferencing accuracy.
Using INT8 precision with quantization scales obtained from Post-Training Quantization (PTQ) can produce additional performance gains, but may also result in accuracy loss. Alternatively, for PyTorch-trained models, NVIDIA [PyTorch-Quantization toolkit](https://github.com/NVIDIA/TensorRT/tree/main/tools/pytorch-quantization) can be leveraged to perform quantized fine tuning (a.k.a. Quantization Aware Training or QAT) and generate the INT8 quantization scales as part of training. This generally results in higher accuracy compared to PTQ.
To demonstrate the potential speedups from these optimizations in demoBERT, we provide the [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) transformer model finetuned for SQuAD 2.0 task with sparsity and quantization.
The sparse weights are generated by finetuning with INT8 Quantization Aware Training recipe. This feature can be used with the fixed or variable sequence length implementations by passing in `-sp` flag to demoBERT builder.
#### Megatron-LM for Question Answering
##### Example: Megatron-LM Large SQuAD v2.0 with sparse weights for sequence length 384
**Build the TensorRT engine**:
Options specified:
- `--megatron` : assume Megatron style residuals instead of vanilla BERT.
- `--pickle` : specify a pickle file containing the PyTorch statedict corresponding to fine-tuned Megatron model.
- `-sp` : enable sparsity during engine optimization and treat the weights as sparse.
- `--int8 --il` : enable int8 tactics/plugins with interleaving.
```bash
bash ./scripts/download_model.sh 384 v1_1 # BERT-large model checkpoint fine-tuned for SQuAD 1.1
bash ./scripts/download_model.sh pyt megatron-large int8-qat sparse # Megatron-LM model weights
export CKPT_PATH=models/fine-tuned/bert_pyt_statedict_megatron_sparse_int8qat_v21.03.0/bert_pyt_statedict_megatron_sparse_int8_qat
mkdir -p engines && python3 builder_varseqlen.py -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -b 1 -s 384 -o engines/megatron_large_seqlen384_int8qat_sparse.engine --fp16 --int8 --strict -il --megatron --pickle $CKPT_PATH -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -sp
```
**Ask a question**:
```bash
python3 inference_varseqlen.py -e engines/megatron_large_seqlen384_int8qat_sparse.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256
```
**Evaluate F1 score**:
```bash
python3 inference_varseqlen.py -e engines/megatron_large_seqlen384_int8qat_sparse.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
Expected output:
```
&&&& PASSED TensorRT BERT Squad Accuracy matches reference.
{"exact_match": 84.03973509933775, "f1": 90.88667129897755}
```
## Performance
### Benchmarking
The following section shows how to run the inference benchmarks for BERT.
#### TensorRT inference benchmark
The inference benchmark is performed on a single GPU by the `inference_benchmark.sh` script, which takes the following steps for each set of model parameters:
1. Downloads checkpoints and builds a TensorRT engine if it does not already exist.
2. Runs 100 warm-up iteration then runs inference for 1000 to 2000 iterations for each batch size specified in the script, selecting the profile best for each size.
**Note:** The time measurements do not include the time required to copy inputs to the device and copy outputs to the host.
To run the inference benchmark script, run:
```bash
bash scripts/inference_benchmark.sh --gpu <arch>
```
Options for `<arch>` are: 'Volta', 'Xavier', 'Turing', 'Ampere'
Note: Some of the configurations in the benchmark script require 16GB of GPU memory. On GPUs with smaller amounts of memory, parts of the benchmark may fail to run.
Also note that BERT Large engines, especially using mixed precision with large batch sizes and sequence lengths may take a couple hours to build.
### Results
The following sections provide details on how we achieved our performance and inference.
#### Inference performance: NVIDIA A100 (40GB)
Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` on NVIDIA A100 (40G).
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.67 | 0.67 | 0.54 | 0.62 | 0.80 | 0.62 |
| 128 | 2 | 0.76 | 0.76 | 0.60 | 0.92 | 0.92 | 0.73 |
| 128 | 4 | 0.73 | 0.93 | 0.73 | 0.93 | 0.93 | 0.93 |
| 128 | 8 | 0.94 | 1.21 | 0.95 | 1.31 | 1.31 | 1.31 |
| 128 | 12 | 1.20 | 1.20 | 1.20 | 1.72 | 2.20 | 1.72 |
| 128 | 16 | 1.34 | 1.34 | 1.34 | 2.07 | 2.08 | 2.05 |
| 128 | 24 | 1.82 | 1.82 | 1.82 | 3.02 | 3.08 | 3.01 |
| 128 | 32 | 2.23 | 2.24 | 2.23 | 3.89 | 3.91 | 3.85 |
| 128 | 64 | 4.16 | 4.16 | 4.12 | 7.57 | 7.63 | 7.55 |
| 128 | 128 | 8.07 | 8.09 | 8.02 | 15.23 | 15.24 | 15.15 |
| 384 | 1 | 1.14 | 1.46 | 1.14 | 1.25 | 1.61 | 1.26 |
| 384 | 2 | 1.32 | 1.32 | 1.32 | 1.55 | 1.55 | 1.55 |
| 384 | 4 | 1.66 | 1.66 | 1.66 | 2.12 | 2.12 | 2.12 |
| 384 | 8 | 2.20 | 2.21 | 2.20 | 3.34 | 3.36 | 3.31 |
| 384 | 12 | 3.31 | 3.31 | 3.31 | 4.78 | 4.82 | 4.77 |
| 384 | 16 | 4.00 | 4.00 | 4.00 | 6.38 | 6.40 | 6.33 |
| 384 | 24 | 5.70 | 5.70 | 5.70 | 9.31 | 9.31 | 9.22 |
| 384 | 32 | 7.64 | 7.64 | 7.64 | 12.90 | 12.90 | 12.79 |
| 384 | 64 | 14.87 | 14.91 | 14.74 | 24.96 | 25.19 | 24.74 |
| 384 | 128 | 29.01 | 29.02 | 28.74 | 49.05 | 49.28 | 48.64 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.23 | 1.23 | 1.23 | 1.54 | 1.55 | 1.54 |
| 128 | 2 | 1.42 | 1.42 | 1.42 | 1.82 | 2.02 | 1.82 |
| 128 | 4 | 1.79 | 1.79 | 1.78 | 2.52 | 2.53 | 2.52 |
| 128 | 8 | 2.64 | 2.65 | 2.64 | 3.93 | 3.94 | 3.89 |
| 128 | 12 | 3.11 | 3.11 | 3.11 | 5.03 | 5.07 | 5.00 |
| 128 | 16 | 4.09 | 4.09 | 4.08 | 6.93 | 6.94 | 6.86 |
| 128 | 24 | 5.28 | 5.28 | 5.27 | 9.70 | 9.70 | 9.65 |
| 128 | 32 | 7.00 | 7.02 | 6.95 | 12.95 | 12.96 | 12.83 |
| 128 | 64 | 12.85 | 12.89 | 12.74 | 24.85 | 25.06 | 24.63 |
| 128 | 128 | 25.07 | 25.08 | 24.99 | 49.15 | 49.42 | 48.69 |
| 384 | 1 | 2.55 | 2.55 | 2.55 | 2.96 | 2.96 | 2.96 |
| 384 | 2 | 3.03 | 3.03 | 3.03 | 3.90 | 3.90 | 3.89 |
| 384 | 4 | 4.01 | 4.01 | 4.01 | 5.73 | 5.79 | 5.67 |
| 384 | 8 | 7.16 | 7.16 | 7.16 | 11.12 | 11.16 | 11.01 |
| 384 | 12 | 9.14 | 9.14 | 9.13 | 15.31 | 15.45 | 15.27 |
| 384 | 16 | 12.28 | 12.28 | 12.28 | 20.99 | 20.99 | 20.92 |
| 384 | 24 | 17.67 | 17.72 | 17.57 | 30.75 | 31.03 | 30.66 |
| 384 | 32 | 23.29 | 23.31 | 23.06 | 41.01 | 41.26 | 40.61 |
| 384 | 64 | 44.96 | 45.30 | 44.83 | 79.97 | 80.27 | 79.26 |
| 384 | 128 | 87.99 | 88.02 | 87.69 | 156.51 | 156.99 | 155.47 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.13 | 1.44 | 1.14 |
| 128 | 2 | 1.37 | 1.37 | 1.37 |
| 128 | 4 | 1.78 | 1.78 | 1.77 |
| 128 | 8 | 2.45 | 2.46 | 2.45 |
| 128 | 12 | 3.11 | 3.12 | 3.10 |
| 128 | 16 | 3.91 | 3.91 | 3.90 |
| 128 | 24 | 4.89 | 4.89 | 4.88 |
| 128 | 32 | 6.96 | 6.97 | 6.91 |
| 128 | 64 | 11.64 | 11.65 | 11.63 |
| 128 | 128 | 21.82 | 21.83 | 21.69 |
| 384 | 1 | 1.69 | 1.69 | 1.69 |
| 384 | 2 | 2.21 | 2.22 | 2.21 |
| 384 | 4 | 3.63 | 3.63 | 3.62 |
| 384 | 8 | 5.72 | 5.72 | 5.71 |
| 384 | 12 | 8.38 | 8.39 | 8.37 |
| 384 | 16 | 10.35 | 10.35 | 10.34 |
| 384 | 24 | 14.49 | 14.49 | 14.48 |
| 384 | 32 | 18.75 | 18.81 | 18.73 |
| 384 | 64 | 36.38 | 36.41 | 36.11 |
| 384 | 128 | 69.82 | 69.95 | 69.34 |
#### Inference performance: NVIDIA A30
Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` on NVIDIA A30.
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.62 | 0.62 | 0.61 | 1.01 | 1.02 | 1.00 |
| 128 | 2 | 0.79 | 0.80 | 0.77 | 1.33 | 1.35 | 1.31 |
| 128 | 4 | 1.16 | 1.16 | 1.13 | 2.23 | 2.23 | 2.16 |
| 128 | 8 | 1.93 | 1.98 | 1.91 | 3.70 | 3.83 | 3.69 |
| 128 | 12 | 2.69 | 2.69 | 2.63 | 5.42 | 5.46 | 5.36 |
| 128 | 16 | 3.38 | 3.39 | 3.32 | 6.77 | 6.78 | 6.71 |
| 128 | 24 | 4.87 | 4.87 | 4.77 | 10.72 | 10.81 | 10.56 |
| 128 | 32 | 6.22 | 6.35 | 6.18 | 14.13 | 14.14 | 13.97 |
| 128 | 64 | 13.69 | 13.85 | 13.56 | 31.28 | 31.69 | 31.05 |
| 128 | 128 | 30.49 | 30.72 | 29.90 | 69.99 | 70.38 | 68.61 |
| 384 | 1 | 1.31 | 1.31 | 1.30 | 2.10 | 2.10 | 2.09 |
| 384 | 2 | 1.85 | 1.86 | 1.85 | 3.19 | 3.21 | 3.14 |
| 384 | 4 | 3.00 | 3.00 | 2.94 | 5.77 | 5.89 | 5.74 |
| 384 | 8 | 5.58 | 5.60 | 5.48 | 11.49 | 11.59 | 11.38 |
| 384 | 12 | 8.22 | 8.37 | 8.13 | 17.39 | 17.40 | 17.16 |
| 384 | 16 | 10.98 | 10.99 | 10.89 | 23.38 | 23.78 | 23.02 |
| 384 | 24 | 17.33 | 17.47 | 17.09 | 38.54 | 39.55 | 37.57 |
| 384 | 32 | 23.82 | 24.18 | 23.56 | 51.12 | 51.24 | 50.62 |
| 384 | 64 | 50.08 | 50.28 | 49.10 | 105.60 | 106.08 | 104.59 |
| 384 | 128 | 113.95 | 114.53 | 112.15 | 209.55 | 209.93 | 208.35 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.80 | 1.80 | 1.78 | 3.12 | 3.12 | 3.10 |
| 128 | 2 | 2.51 | 2.52 | 2.45 | 4.36 | 4.38 | 4.34 |
| 128 | 4 | 3.70 | 3.72 | 3.60 | 6.91 | 6.94 | 6.82 |
| 128 | 8 | 6.40 | 6.41 | 6.30 | 12.81 | 12.96 | 12.76 |
| 128 | 12 | 8.53 | 8.60 | 8.36 | 18.79 | 18.96 | 18.43 |
| 128 | 16 | 11.25 | 11.34 | 11.18 | 25.61 | 25.85 | 25.34 |
| 128 | 24 | 16.25 | 16.28 | 16.01 | 36.21 | 36.22 | 35.95 |
| 128 | 32 | 21.65 | 21.68 | 21.37 | 49.57 | 49.76 | 49.17 |
| 128 | 64 | 44.98 | 45.44 | 44.57 | 107.87 | 108.20 | 106.77 |
| 128 | 128 | 93.97 | 94.63 | 93.01 | 216.03 | 216.54 | 214.69 |
| 384 | 1 | 3.47 | 3.48 | 3.45 | 6.64 | 6.75 | 6.43 |
| 384 | 2 | 5.57 | 5.58 | 5.46 | 10.63 | 10.65 | 10.49 |
| 384 | 4 | 9.79 | 9.93 | 9.62 | 20.78 | 21.19 | 20.38 |
| 384 | 8 | 18.38 | 18.39 | 18.22 | 39.85 | 40.17 | 38.38 |
| 384 | 12 | 26.50 | 26.74 | 26.39 | 61.30 | 61.76 | 59.94 |
| 384 | 16 | 37.19 | 37.48 | 36.70 | 81.72 | 82.15 | 80.66 |
| 384 | 24 | 55.13 | 55.69 | 54.64 | 131.37 | 131.61 | 130.29 |
| 384 | 32 | 76.86 | 77.41 | 75.98 | 166.22 | 166.56 | 165.16 |
| 384 | 64 | 165.08 | 165.56 | 163.82 | 344.18 | 344.61 | 342.97 |
| 384 | 128 | 334.73 | 335.97 | 332.16 | 670.67 | 671.67 | 668.80 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.51 | 1.51 | 1.49 |
| 128 | 2 | 2.07 | 2.09 | 2.03 |
| 128 | 4 | 2.98 | 3.02 | 2.92 |
| 128 | 8 | 5.06 | 5.07 | 5.05 |
| 128 | 12 | 6.70 | 6.77 | 6.63 |
| 128 | 16 | 8.81 | 8.82 | 8.74 |
| 128 | 24 | 13.18 | 13.19 | 13.09 |
| 128 | 32 | 17.43 | 17.44 | 17.34 |
| 128 | 64 | 36.26 | 36.70 | 35.86 |
| 128 | 128 | 79.70 | 79.88 | 79.06 |
| 384 | 1 | 2.80 | 2.81 | 2.75 |
| 384 | 2 | 4.21 | 4.21 | 4.15 |
| 384 | 4 | 7.64 | 7.66 | 7.53 |
| 384 | 8 | 14.96 | 14.98 | 14.83 |
| 384 | 12 | 21.62 | 21.66 | 21.46 |
| 384 | 16 | 28.40 | 28.57 | 28.31 |
| 384 | 24 | 45.11 | 45.45 | 44.78 |
| 384 | 32 | 60.86 | 61.08 | 59.88 |
| 384 | 64 | 126.53 | 126.80 | 126.06 |
| 384 | 128 | 255.35 | 256.27 | 253.63 |
### Inference Performance NVIDIA L40S
Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` on NVIDIA L40S.
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.34 | 0.34 | 0.34 | 0.48 | 0.48 | 0.48 |
| 128 | 2 | 0.41 | 0.41 | 0.41 | 0.56 | 0.56 | 0.55 |
| 128 | 4 | 0.50 | 0.51 | 0.50 | 0.77 | 0.78 | 0.77 |
| 128 | 8 | 0.68 | 0.68 | 0.67 | 1.26 | 1.26 | 1.25 |
| 128 | 12 | 0.91 | 0.91 | 0.91 | 1.69 | 1.69 | 1.68 |
| 128 | 16 | 1.11 | 1.11 | 1.11 | 2.24 | 2.24 | 2.23 |
| 128 | 24 | 1.46 | 1.46 | 1.46 | 3.18 | 3.19 | 3.18 |
| 128 | 32 | 1.82 | 1.82 | 1.81 | 3.94 | 3.94 | 3.93 |
| 128 | 64 | 3.44 | 3.44 | 3.42 | 7.98 | 8.08 | 7.90 |
| 128 | 128 | 7.25 | 7.29 | 7.20 | 17.35 | 17.40 | 17.13 |
| 384 | 1 | 0.73 | 0.73 | 0.73 | 1.04 | 1.04 | 1.03 |
| 384 | 2 | 0.88 | 0.88 | 0.88 | 1.35 | 1.35 | 1.35 |
| 384 | 4 | 1.17 | 1.17 | 1.17 | 2.14 | 2.14 | 2.13 |
| 384 | 8 | 1.70 | 1.71 | 1.69 | 3.47 | 3.47 | 3.46 |
| 384 | 12 | 2.72 | 2.72 | 2.72 | 5.08 | 5.09 | 5.06 |
| 384 | 16 | 3.26 | 3.26 | 3.24 | 7.18 | 7.19 | 7.15 |
| 384 | 24 | 4.94 | 4.94 | 4.89 | 9.98 | 10.00 | 9.92 |
| 384 | 32 | 6.11 | 6.13 | 6.09 | 13.35 | 13.38 | 13.25 |
| 384 | 64 | 12.96 | 13.00 | 12.84 | 28.93 | 29.37 | 28.41 |
| 384 | 128 | 27.22 | 27.36 | 26.87 | 59.55 | 59.91 | 58.44 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.89 | 0.89 | 0.89 | 1.30 | 1.30 | 1.30 |
| 128 | 2 | 0.98 | 0.98 | 0.98 | 1.45 | 1.46 | 1.45 |
| 128 | 4 | 1.35 | 1.35 | 1.34 | 2.32 | 2.32 | 2.31 |
| 128 | 8 | 1.93 | 1.95 | 1.92 | 3.59 | 3.60 | 3.58 |
| 128 | 12 | 2.73 | 2.73 | 2.72 | 5.70 | 5.71 | 5.63 |
| 128 | 16 | 3.19 | 3.21 | 3.17 | 6.48 | 6.49 | 6.45 |
| 128 | 24 | 4.50 | 4.53 | 4.48 | 9.89 | 9.90 | 9.81 |
| 128 | 32 | 5.66 | 5.68 | 5.62 | 12.26 | 12.30 | 12.16 |
| 128 | 64 | 11.42 | 11.43 | 11.30 | 27.40 | 27.60 | 27.16 |
| 128 | 128 | 24.68 | 24.70 | 24.36 | 61.49 | 61.76 | 60.81 |
| 384 | 1 | 1.68 | 1.68 | 1.68 | 2.73 | 2.73 | 2.73 |
| 384 | 2 | 2.28 | 2.28 | 2.27 | 3.83 | 3.83 | 3.82 |
| 384 | 4 | 3.28 | 3.28 | 3.26 | 6.26 | 6.26 | 6.24 |
| 384 | 8 | 4.97 | 4.98 | 4.95 | 10.32 | 10.33 | 10.30 |
| 384 | 12 | 7.89 | 7.89 | 7.86 | 17.49 | 17.50 | 17.43 |
| 384 | 16 | 9.47 | 9.49 | 9.44 | 21.50 | 21.62 | 21.24 |
| 384 | 24 | 14.64 | 14.66 | 14.54 | 33.26 | 33.30 | 33.01 |
| 384 | 32 | 19.20 | 19.37 | 18.97 | 44.56 | 44.69 | 43.95 |
| 384 | 64 | 42.15 | 42.38 | 41.56 | 98.89 | 99.28 | 97.19 |
| 384 | 128 | 84.15 | 84.40 | 83.34 | 196.98 | 197.83 | 194.18 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.76 | 0.76 | 0.76 |
| 128 | 2 | 0.90 | 0.90 | 0.90 |
| 128 | 4 | 1.13 | 1.13 | 1.13 |
| 128 | 8 | 1.71 | 1.71 | 1.71 |
| 128 | 12 | 2.26 | 2.26 | 2.25 |
| 128 | 16 | 2.72 | 2.73 | 2.72 |
| 128 | 24 | 4.44 | 4.45 | 4.43 |
| 128 | 32 | 5.07 | 5.11 | 5.04 |
| 128 | 64 | 10.06 | 10.09 | 9.97 |
| 128 | 128 | 20.42 | 20.46 | 20.30 |
| 384 | 1 | 1.13 | 1.13 | 1.13 |
| 384 | 2 | 1.63 | 1.65 | 1.62 |
| 384 | 4 | 2.52 | 2.53 | 2.51 |
| 384 | 8 | 4.93 | 4.94 | 4.90 |
| 384 | 12 | 6.47 | 6.47 | 6.45 |
| 384 | 16 | 8.41 | 8.42 | 8.36 |
| 384 | 24 | 12.52 | 12.53 | 12.44 |
| 384 | 32 | 16.66 | 16.72 | 16.57 |
| 384 | 64 | 34.12 | 34.22 | 33.81 |
| 384 | 128 | 71.98 | 72.13 | 71.52 |
## Hardware Platform Support
The scripts call TensorRT Plugins underneath, whose kernel optimizations depend on the NVIDIA GPU Architecture.
The compute capability of an NVIDIA GPU can be found out using the `nvidia-smi` commandline utility. One can execute the following command in the terminal to find out the compute capability of the GPU.
```bash
nvidia-smi --query-gpu=compute_cap --format=csv
```
Currently, this demo is supported on the following compute capabilities. This list is subject to change as new architectures are released.
- Volta architecture - 7.2, 7.5
- Ampere architecture - 8.0, 8.6, 8.7, 8.9
- Hopper architecture - 9.0 (since October 2022)
- Blackwell architecture - 10.0, 12.0 (since Jan 2025). Not recommended with `--use-v3-plugins` option.
demo/BERT/builder.py
-720
View File
@@ -1,720 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse
import ctypes
import json
import numpy as np
import os
import os.path
import re
import sys
import time
import onnx
from helpers.cuda_utils import getComputeCapacity
# TensorRT
import tensorrt as trt
from helpers.calibrator import BertCalibrator as BertCalibrator
from builder_utils import load_tf_weights, load_pytorch_weights_and_quant, load_onnx_weights_and_quant
from builder_utils import WQKV, BQKV # Attention Keys
from builder_utils import W_AOUT, B_AOUT, W_MID, B_MID, W_LOUT, B_LOUT # Transformer Keys
from builder_utils import SQD_W, SQD_B # SQuAD Output Keys
from builder_utils import (
create_plugin,
add_plugin_to_network,
) # Plugin Helper functions
"""
TensorRT Initialization
"""
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
trt_version = [n for n in trt.__version__.split('.')]
# Import necessary plugins for demoBERT
plugin_lib_name = "nvinfer_plugin_10.dll" if sys.platform == "win32" else "libnvinfer_plugin.so"
env_name_to_add_path = "PATH" if sys.platform == "win32" else "LD_LIBRARY_PATH"
handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)
if not handle:
raise RuntimeError("Could not load plugin library. Is `{}` on your {}?".format(plugin_lib_name, env_name_to_add_path))
trt.init_libnvinfer_plugins(TRT_LOGGER, "")
plg_registry = trt.get_plugin_registry()
class BertConfig:
def __init__(
self,
bert_config_path,
use_fp16,
use_int8,
use_strict,
use_fc2_gemm,
use_int8_skipln,
use_int8_multihead,
use_qat,
use_sparsity,
timing_cache,
distributive_independence = False,
use_deprecated_plugins=False,
):
with open(bert_config_path, "r") as f:
data = json.load(f)
self.num_attention_heads = data["num_attention_heads"]
self.hidden_size = data["hidden_size"]
self.intermediate_size = data["intermediate_size"]
self.num_hidden_layers = data["num_hidden_layers"]
self.head_size = self.hidden_size // self.num_attention_heads
self.use_fp16 = use_fp16
self.use_int8 = use_int8
self.use_fc2_gemm = use_fc2_gemm
self.use_strict = use_strict
self.use_int8_skipln = use_int8_skipln
self.use_int8_multihead = use_int8_multihead
self.is_calib_mode = False
self.use_qat = use_qat
self.use_sparsity = use_sparsity
self.timing_cache = timing_cache
self.use_deprecated_plugins = use_deprecated_plugins
self.distributive_independence = distributive_independence
def set_tensor_name(tensor, prefix, name):
tensor.name = prefix + name
def set_output_name(layer, prefix, name, out_idx = 0):
set_tensor_name(layer.get_output(out_idx), prefix, name)
def set_output_range(layer, maxval, out_idx = 0):
layer.get_output(out_idx).set_dynamic_range(-maxval, maxval)
def get_mha_dtype(config):
dtype = trt.float32
if config.use_fp16:
dtype = trt.float16
# Multi-head attention doesn't use INT8 inputs and output by default unless it is specified.
if config.use_int8 and config.use_int8_multihead and not config.is_calib_mode:
dtype = trt.int8
return int(dtype)
def attention_layer_opt(prefix, config, init_dict, network, input_tensor, imask):
"""
Add the attention layer
"""
assert(len(input_tensor.shape) == 5)
B, S, hidden_size, _, _ = input_tensor.shape
num_heads = config.num_attention_heads
head_size = int(hidden_size / num_heads)
Wall = init_dict[prefix + WQKV]
Ball = init_dict[prefix + BQKV]
# FC_attention
mult_all = network.add_convolution_nd(input_tensor, 3 * hidden_size, (1, 1), Wall, Ball)
if config.use_qat:
dr_qkv = max(
init_dict[prefix + 'self_qv_a_input_quantizer_amax'],
init_dict[prefix + 'self_qv_b_input_quantizer_amax'],
init_dict[prefix + 'self_av_b_input_quantizer_amax'],
)
set_output_range(mult_all, dr_qkv)
set_output_name(mult_all, prefix, "qkv_mult")
has_mask = imask is not None
# QKV2CTX
pf_type = trt.PluginField("type_id", np.array([get_mha_dtype(config)], np.int32), trt.PluginFieldType.INT32)
pf_hidden_size = trt.PluginField("hidden_size", np.array([hidden_size], np.int32), trt.PluginFieldType.INT32)
pf_num_heads = trt.PluginField("num_heads", np.array([num_heads], np.int32), trt.PluginFieldType.INT32)
pf_has_mask = trt.PluginField("has_mask", np.array([has_mask], np.int32), trt.PluginFieldType.INT32)
if config.use_qat:
dr_probs = init_dict[prefix + 'self_av_a_input_quantizer_amax']
dq_probs = dr_probs / 127.0
pf_dq_probs = trt.PluginField("dq_probs", np.array([dq_probs], np.float32), trt.PluginFieldType.FLOAT32)
pfc = trt.PluginFieldCollection([pf_hidden_size, pf_num_heads, pf_has_mask, pf_type, pf_dq_probs])
else:
pfc = trt.PluginFieldCollection(
[pf_hidden_size, pf_num_heads, pf_has_mask, pf_type]
)
qkv2ctx_plugin = create_plugin(
"qkv_to_context", plg_registry, pfc, use_deprecated_plugins=config.use_deprecated_plugins
)
qkv_in = [mult_all.get_output(0)]
if has_mask:
qkv_in.append(imask)
qkv2ctx_layer = add_plugin_to_network(
network, qkv2ctx_plugin, qkv_in, use_deprecated_plugins=config.use_deprecated_plugins
)
if config.use_qat:
dr_ctx = init_dict[prefix + 'output_dense_input_amax']
set_output_range(qkv2ctx_layer, dr_ctx)
set_output_name(qkv2ctx_layer, prefix, "context_layer")
return qkv2ctx_layer
def skipln(prefix, config, init_dict, network, input_tensor, skip, bias=None):
"""
Add the skip layer
"""
idims = input_tensor.shape
assert len(idims) == 5
hidden_size = idims[2]
dtype = trt.float32
if config.use_fp16:
dtype = trt.float16
# Skip layernorm doesn't use INT8 inputs and output by default unless it is specified.
if config.use_int8 and config.use_int8_skipln and not config.is_calib_mode:
dtype = trt.int8
pf_ld = trt.PluginField("ld", np.array([hidden_size], np.int32), trt.PluginFieldType.INT32)
wbeta = init_dict[prefix + "beta"]
pf_beta = trt.PluginField("beta", wbeta.numpy(), trt.PluginFieldType.FLOAT32)
wgamma = init_dict[prefix + "gamma"]
pf_gamma = trt.PluginField("gamma", wgamma.numpy(), trt.PluginFieldType.FLOAT32)
pf_type = trt.PluginField("type_id", np.array([int(dtype)], np.int32), trt.PluginFieldType.INT32)
fields = [pf_ld, pf_beta, pf_gamma, pf_type ]
if bias:
pf_bias = trt.PluginField("bias", bias.numpy(), trt.PluginFieldType.FLOAT32)
fields.append(pf_bias)
pfc = trt.PluginFieldCollection(fields)
skipln_plugin = create_plugin(
"skip_layer_norm", plg_registry, pfc, use_deprecated_plugins=config.use_deprecated_plugins
)
skipln_inputs = [input_tensor, skip]
skipln_layer = add_plugin_to_network(
network, skipln_plugin, skipln_inputs, use_deprecated_plugins=config.use_deprecated_plugins
)
return skipln_layer
def transformer_layer_opt(prefix, config, init_dict, network, input_tensor, imask):
"""
Add the transformer layer
"""
idims = input_tensor.shape
assert len(idims) == 5
hidden_size = idims[2]
if config.use_qat:
dr_input = init_dict[prefix + 'attention_self_query_input_amax']
assert(dr_input ==init_dict[prefix + 'attention_self_key_input_amax'] )
assert(dr_input ==init_dict[prefix + 'attention_self_value_input_amax'] )
input_tensor.set_dynamic_range(-dr_input, dr_input)
context_transposed = attention_layer_opt(prefix + "attention_", config, init_dict, network, input_tensor, imask)
attention_heads = context_transposed.get_output(0)
# FC0
B_aout = init_dict[prefix + B_AOUT]
W_aout = init_dict[prefix + W_AOUT]
attention_out_fc = network.add_convolution_nd(attention_heads, hidden_size, (1, 1), W_aout, B_aout)
if config.use_int8 and not config.use_int8_skipln:
attention_out_fc.set_output_type(0, trt.DataType.HALF if config.use_fp16 else trt.DataType.FLOAT)
if config.use_int8 and config.use_qat:
dr_fc_aout = init_dict[prefix + 'attention_output_add_local_input_quantizer_amax']
set_output_range(attention_out_fc, dr_fc_aout)
skiplayer = skipln(prefix + "attention_output_layernorm_",config, init_dict, network, attention_out_fc.get_output(0), input_tensor, bias=None)
attention_ln = skiplayer.get_output(0)
if config.use_qat:
dr_skln1 = init_dict[prefix + 'intermediate_dense_input_amax']
set_output_range(skiplayer, dr_skln1)
# FC1 + GELU
B_mid = init_dict[prefix + B_MID]
W_mid = init_dict[prefix + W_MID]
mid_dense = network.add_convolution_nd(attention_ln, config.intermediate_size, (1, 1), W_mid, B_mid)
mid_dense_out = mid_dense.get_output(0)
POW = network.add_constant((1, 1, 1, 1, 1), trt.Weights(np.ascontiguousarray([3.0], dtype=np.float32)))
MULTIPLY = network.add_constant((1, 1, 1, 1, 1), trt.Weights(np.ascontiguousarray([0.044715], dtype=np.float32)))
SQRT = network.add_constant((1, 1, 1, 1, 1), trt.Weights((np.ascontiguousarray([0.79788456080286535587989211986876], dtype=np.float32))))
ONE = network.add_constant((1, 1, 1, 1, 1), trt.Weights((np.ascontiguousarray([1.0], dtype=np.float32))))
HALF = network.add_constant((1, 1, 1, 1, 1), trt.Weights((np.ascontiguousarray([0.5], dtype=np.float32))))
X_pow = network.add_elementwise(mid_dense_out, POW.get_output(0), trt.ElementWiseOperation.POW)
X_pow_t = X_pow.get_output(0)
X_mul = network.add_elementwise(X_pow_t, MULTIPLY.get_output(0), trt.ElementWiseOperation.PROD)
X_add = network.add_elementwise(mid_dense_out, X_mul.get_output(0), trt.ElementWiseOperation.SUM)
X_sqrt = network.add_elementwise(X_add.get_output(0), SQRT.get_output(0), trt.ElementWiseOperation.PROD)
X_sqrt_tensor = X_sqrt.get_output(0)
X_tanh = network.add_activation(X_sqrt_tensor, trt.ActivationType.TANH)
X_tanh_tensor = X_tanh.get_output(0)
X_one = network.add_elementwise(X_tanh_tensor, ONE.get_output(0), trt.ElementWiseOperation.SUM)
CDF = network.add_elementwise(X_one.get_output(0), HALF.get_output(0), trt.ElementWiseOperation.PROD)
gelu_layer = network.add_elementwise(CDF.get_output(0), mid_dense_out, trt.ElementWiseOperation.PROD)
intermediate_act = gelu_layer.get_output(0)
set_tensor_name(intermediate_act, prefix, "gelu")
if config.use_int8:
if config.use_qat:
dr_gelu = init_dict[prefix + 'output_dense_input_amax']
set_output_range(gelu_layer, dr_gelu)
else:
# use gelu10 according to whitepaper http://arxiv.org/abs/2004.09602
set_output_range(gelu_layer, 10)
# FC2
# Dense to hidden size
B_lout = init_dict[prefix + B_LOUT]
W_lout = init_dict[prefix + W_LOUT]
out_dense = network.add_convolution_nd(intermediate_act, hidden_size, (1, 1), W_lout, B_lout)
if config.use_int8 and not config.use_int8_skipln:
out_dense.set_output_type(0, trt.DataType.HALF if config.use_fp16 else trt.DataType.FLOAT)
if config.use_qat:
dr_fc_out = init_dict[prefix + 'output_add_local_input_quantizer_amax']
set_output_range(out_dense, dr_fc_out)
set_output_name(out_dense, prefix + "output_", "dense")
out_layer = skipln(prefix + "output_layernorm_", config, init_dict, network, out_dense.get_output(0), attention_ln, bias=None)
set_output_name(out_layer, prefix + "output_", "reshape")
return out_layer
def bert_model(config, init_dict, network, input_tensor, input_mask):
"""
Create the bert model
"""
prev_input = input_tensor
for layer in range(0, config.num_hidden_layers):
ss = "l{}_".format(layer)
out_layer = transformer_layer_opt(ss, config, init_dict, network, prev_input, input_mask)
prev_input = out_layer.get_output(0)
if config.use_qat:
dr_out = init_dict["bert_encoder_final_input_quantizer_amax"]
set_output_range(out_layer, dr_out)
return prev_input
def squad_output(prefix, config, init_dict, network, input_tensor):
"""
Create the squad output
"""
idims = input_tensor.shape
assert len(idims) == 5
B, S, hidden_size, _, _ = idims
W_out = init_dict[prefix + SQD_W]
B_out = init_dict[prefix + SQD_B]
W = network.add_constant((1, hidden_size, 2), W_out)
dense = network.add_convolution_nd(input_tensor, 2, (1, 1), W_out, B_out)
OUT = network.add_shuffle(dense.get_output(0))
OUT.second_transpose = (1, 0, 2, 3, 4)
set_output_name(OUT, prefix, "squad_logits")
return OUT
def emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_lengths, batch_sizes):
# int8 only support some of the sequence length, we dynamic on sequence length is not allowed.
input_ids = network.add_input(name="input_ids", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0]))
segment_ids = network.add_input(name="segment_ids", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0]))
input_mask = network.add_input(name="input_mask", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0]))
# Specify profiles for the batch sizes we're interested in.
# Make sure the profile also works for all sizes not covered by the previous profile.
# When distributive independence is enabled, only one profile can be used.
if config.distributive_independence:
max_batch_size = max(batch_sizes)
max_sequence_length = max(sequence_lengths)
profile = builder.create_optimization_profile()
min_shape = (1, max_sequence_length)
shape = (max_batch_size, max_sequence_length)
profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("input_mask", min=min_shape, opt=shape, max=shape)
builder_config.add_optimization_profile(profile)
elif len(sequence_lengths) > 1 or len(batch_sizes) > 1:
for batch_size in sorted(batch_sizes):
if len(sequence_lengths) == 1:
profile = builder.create_optimization_profile()
min_shape = (1, sequence_lengths[0])
shape = (batch_size, sequence_lengths[0])
profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("input_mask", min=min_shape, opt=shape, max=shape)
builder_config.add_optimization_profile(profile)
else:
for sequence_length in sorted(sequence_lengths):
profile = builder.create_optimization_profile()
min_shape = (1, sequence_length)
shape = (batch_size, sequence_length)
profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("input_mask", min=min_shape, opt=shape, max=shape)
builder_config.add_optimization_profile(profile)
wbeta = trt.PluginField("bert_embeddings_layernorm_beta", weights_dict["bert_embeddings_layernorm_beta"].numpy(), trt.PluginFieldType.FLOAT32)
wgamma = trt.PluginField("bert_embeddings_layernorm_gamma", weights_dict["bert_embeddings_layernorm_gamma"].numpy(), trt.PluginFieldType.FLOAT32)
wwordemb = trt.PluginField("bert_embeddings_word_embeddings", weights_dict["bert_embeddings_word_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
wtokemb = trt.PluginField("bert_embeddings_token_type_embeddings", weights_dict["bert_embeddings_token_type_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
wposemb = trt.PluginField("bert_embeddings_position_embeddings", weights_dict["bert_embeddings_position_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
output_fp16 = trt.PluginField("output_fp16", np.array([1 if config.use_fp16 else 0]).astype(np.int32), trt.PluginFieldType.INT32)
mha_type = trt.PluginField("mha_type_id", np.array([get_mha_dtype(config)], np.int32), trt.PluginFieldType.INT32)
pfc = trt.PluginFieldCollection(
[wbeta, wgamma, wwordemb, wtokemb, wposemb, output_fp16, mha_type]
)
emln_plugin = create_plugin(
"emb_layer_norm", plg_registry, pfc, use_deprecated_plugins=config.use_deprecated_plugins
)
input_ids = network.add_shuffle(input_ids)
input_ids.second_transpose = (1, 0)
segment_ids = network.add_shuffle(segment_ids)
segment_ids.second_transpose = (1, 0)
input_mask = network.add_shuffle(input_mask)
input_mask.second_transpose = (1, 0)
inputs = [
input_ids.get_output(0),
segment_ids.get_output(0),
input_mask.get_output(0),
]
emb_layer = add_plugin_to_network(
network, emln_plugin, inputs, use_deprecated_plugins=config.use_deprecated_plugins
)
if config.use_qat:
set_output_range(emb_layer, 1, 1)
set_output_name(emb_layer, "embeddings_", "output")
return emb_layer
def build_engine(batch_sizes, workspace_size, sequence_lengths, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose):
network_creation_flag = 0
if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys():
network_creation_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(network_creation_flag) as network, builder.create_builder_config() as builder_config:
builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size * (1024 * 1024))
builder_config.avg_timing_iterations = 8
# Cublas tactics can be unset once the qkv plugin does not use it anymore.
builder_config.set_tactic_sources(builder_config.get_tactic_sources() | 1 << int(trt.TacticSource.CUBLAS))
if config.use_fp16:
builder_config.set_flag(trt.BuilderFlag.FP16)
if config.use_int8:
builder_config.set_flag(trt.BuilderFlag.INT8)
if not config.use_qat:
calibrator = BertCalibrator(squad_json, vocab_file, calibrationCacheFile, 1, sequence_lengths[-1], calib_num)
builder_config.set_quantization_flag(trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION)
builder_config.int8_calibrator = calibrator
if config.use_strict:
builder_config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
builder_config.set_flag(trt.BuilderFlag.DIRECT_IO)
builder_config.set_flag(trt.BuilderFlag.REJECT_EMPTY_ALGORITHMS)
if verbose:
builder_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
if config.distributive_independence:
builder_config.set_flag(trt.BuilderFlag.DISTRIBUTIVE_INDEPENDENCE)
if config.use_sparsity:
TRT_LOGGER.log(TRT_LOGGER.INFO, "Setting sparsity flag on builder_config.")
builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
# speed up the engine build for trt major version >= 8
# 1. disable cudnn tactic
# 2. load global timing cache
if int(trt_version[0]) >= 8:
tactic_source = builder_config.get_tactic_sources() & ~(1 << int(trt.TacticSource.CUDNN))
builder_config.set_tactic_sources(tactic_source)
if config.timing_cache != None:
if os.path.exists(config.timing_cache):
with open(config.timing_cache, "rb") as f:
cache = builder_config.create_timing_cache(f.read())
builder_config.set_timing_cache(cache, ignore_mismatch = False)
else:
cache = builder_config.create_timing_cache(b"")
builder_config.set_timing_cache(cache, ignore_mismatch = False)
# only use the largest sequence when in calibration mode
if config.is_calib_mode:
sequence_lengths = sequence_lengths[-1:]
# Create the network
emb_layer = emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_lengths, batch_sizes)
embeddings = emb_layer.get_output(0)
mask_idx = emb_layer.get_output(1)
bert_out = bert_model(config, weights_dict, network, embeddings, mask_idx)
squad_logits = squad_output("cls_", config, weights_dict, network, bert_out)
squad_logits_out = squad_logits.get_output(0)
squad_logits_out.name = "logits_out"
network.mark_output(squad_logits_out)
build_start_time = time.time()
serialized_engine = builder.build_serialized_network(network, builder_config)
build_time_elapsed = (time.time() - build_start_time)
TRT_LOGGER.log(TRT_LOGGER.INFO, "build engine in {:.3f} Sec".format(build_time_elapsed))
# save global timing cache
if int(trt_version[0]) >= 8 and config.timing_cache != None:
cache = builder_config.get_timing_cache()
with cache.serialize() as buffer:
with open(config.timing_cache, "wb") as f:
f.write(buffer)
f.flush()
os.fsync(f)
if config.use_int8 and not config.use_qat:
calibrator.free()
return serialized_engine
def generate_calibration_cache(sequence_lengths, workspace_size, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num):
"""
BERT demo needs a separate engine building path to generate calibration cache.
This is because we need to configure SLN and MHA plugins in FP32 mode when
generating calibration cache, and INT8 mode when building the actual engine.
This cache could be generated by examining certain training data and can be
reused across different configurations.
"""
# dynamic shape not working with calibration, so we need generate a calibration cache first using fulldims network
if not config.use_int8 or os.path.exists(calibrationCacheFile):
return calibrationCacheFile
# generate calibration cache
saved_use_fp16 = config.use_fp16
config.use_fp16 = False
config.is_calib_mode = True
with build_engine([1], workspace_size, sequence_lengths, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, False) as serialized_engine:
TRT_LOGGER.log(TRT_LOGGER.INFO, "calibration cache generated in {:}".format(calibrationCacheFile))
config.use_fp16 = saved_use_fp16
config.is_calib_mode = False
def main():
parser = argparse.ArgumentParser(
description="TensorRT BERT Sample",
formatter_class=argparse.RawTextHelpFormatter,
)
parser.add_argument(
"-m",
"--ckpt",
required=False,
help="The checkpoint file basename, e.g.: basename(model.ckpt-766908.data-00000-of-00001) is model.ckpt-766908 (default: None)",
)
parser.add_argument(
"-x", "--onnx", required=False, help="The ONNX model file path. (default: None)"
)
parser.add_argument(
"-pt", "--pytorch", required=False, help="The PyTorch checkpoint file path. (default: None)"
)
parser.add_argument(
"-o",
"--output",
required=True,
default="bert_base_384.engine",
help="The bert engine file, ex bert.engine (default: bert_base_384.engine)",
)
parser.add_argument(
"-b",
"--batch-size",
default=[],
action="append",
help="Batch size(s) to optimize for. The engine will be usable with any batch size below this, but may not be optimal for smaller sizes. Can be specified multiple times to optimize for more than one batch size. (default: [1])",
type=int,
)
parser.add_argument(
"-s",
"--sequence-length",
default=[],
action="append",
help="Sequence length of the BERT model (default: [128])",
type=int,
)
parser.add_argument(
"-c",
"--config-dir",
required=True,
help="The folder containing the bert_config.json, which can be downloaded e.g. from https://github.com/google-research/bert#pre-trained-models or by running download_models.py in dle/TensorFlow/LanguageModeling/BERT/data/pretrained_models_google",
)
parser.add_argument(
"-f",
"--fp16",
action="store_true",
help="Indicates that inference should be run in FP16 precision (default: false)",
required=False,
)
parser.add_argument(
"-i",
"--int8",
action="store_true",
help="Indicates that inference should be run in INT8 precision (default: false)",
required=False,
)
parser.add_argument(
"-t",
"--strict",
action="store_true",
help="Indicates that inference should be run in strict precision mode (default: false)",
required=False,
)
parser.add_argument(
"-w",
"--workspace-size",
default=2500,
help="Workspace size in MiB for building the BERT engine (default: 2500)",
type=int,
)
parser.add_argument(
"-j",
"--squad-json",
default="squad/dev-v1.1.json",
help="squad json dataset used for int8 calibration (default: squad/dev-v1.1.json)",
required=False,
)
parser.add_argument(
"-v",
"--vocab-file",
default="./pre-trained_model/uncased_L-24_H-1024_A-16/vocab.txt",
help="Path to file containing entire understandable vocab (default: ./pre-trained_model/uncased_L-24_H-1024_A-16/vocab.txt)",
required=False,
)
parser.add_argument(
"-n", "--calib-num", default=100, help="calibration batch numbers (default: 100)", type=int
)
parser.add_argument(
"-p", "--calib-path", help="calibration cache path (default: None)", required=False
)
parser.add_argument(
"-g",
"--force-fc2-gemm",
action="store_true",
help="Force use gemm to implement FC2 layer (default: false)",
required=False,
)
parser.add_argument(
"-iln",
"--force-int8-skipln",
action="store_true",
help="Run skip layernorm with INT8 (FP32 or FP16 by default) inputs and output (default: false)",
required=False,
)
parser.add_argument(
"-imh",
"--force-int8-multihead",
action="store_true",
help="Run multi-head attention with INT8 (FP32 or FP16 by default) input and output (default: false)",
required=False,
)
parser.add_argument(
"-sp",
"--sparse",
action="store_true",
help="Indicates that model is sparse (default: false)",
required=False,
)
parser.add_argument(
"-tcf",
"--timing-cache-file",
help="Path to tensorrt build timeing cache file, only available for tensorrt 8.0 and later (default: None)",
required=False,
)
parser.add_argument(
"--distributive_independence",
default=False,
action="store_true",
help="Enable TensorRT's distributive independence builder flag (default: false)",
required=False,
)
parser.add_argument(
"--verbose",
action="store_true",
help="Turn on verbose logger and set profiling verbosity to DETAILED (default: false)",
required=False,
)
plugin_group = parser.add_mutually_exclusive_group(required=False)
plugin_group.add_argument('--use-v3-plugins',
dest='use_deprecated_plugins',
action='store_false',
help="Use plugins implementing the IPluginV3 interface wherever TensorRT plugins are used. Cannot be used with --use-deprecated-plugins. Enabling this option should not affect functionality or performance. (default: false)")
plugin_group.add_argument('--use-deprecated-plugins',
dest='use_deprecated_plugins',
action='store_true',
help="Use deprecated plugins implementing the IPluginV2 interface wherever TensorRT plugins are used (instead of updated plugins implementing the IPluginV3 interface). Cannot be used with --use-v3-plugins. Disabling this option should not affect functionality or performance. (default: true)")
parser.set_defaults(use_deprecated_plugins=True)
args, _ = parser.parse_known_args()
args.batch_size = args.batch_size or [1]
args.sequence_length = args.sequence_length or [128]
cc = getComputeCapacity()
if cc[0] * 10 + cc[1] < 75 and args.force_int8_multihead:
raise RuntimeError("--force-int8-multihead option is only supported on Turing+ GPU.")
if cc[0] * 10 + cc[1] < 72 and args.force_int8_skipln:
raise RuntimeError("--force-int8-skipln option is only supported on Xavier+ GPU.")
if args.verbose:
TRT_LOGGER.min_severity = TRT_LOGGER.VERBOSE
bert_config_path = os.path.join(args.config_dir, "bert_config.json")
TRT_LOGGER.log(TRT_LOGGER.INFO, "Using configuration file: {:}".format(bert_config_path))
config = BertConfig(
bert_config_path,
args.fp16,
args.int8,
args.strict,
args.force_fc2_gemm,
args.force_int8_skipln,
args.force_int8_multihead,
args.int8 and args.onnx != None,
args.sparse,
args.timing_cache_file,
args.distributive_independence,
args.use_deprecated_plugins,
)
if args.calib_path != None:
calib_cache = args.calib_path
else:
calib_cache = "BertSquadL{}H{}A{}S{}CalibCache".format(config.num_hidden_layers, config.head_size, config.num_attention_heads, "-".join(str(len) for len in args.sequence_length))
if args.onnx != None:
weights_dict = load_onnx_weights_and_quant(args.onnx, config)
elif args.pytorch != None:
weights_dict = load_pytorch_weights_and_quant(args.pytorch, config)
elif args.ckpt != None:
weights_dict = load_tf_weights(args.ckpt, config)
generate_calibration_cache(args.sequence_length, args.workspace_size, config, weights_dict, args.squad_json, args.vocab_file, calib_cache, args.calib_num)
else:
raise RuntimeError("You need either specify TF checkpoint using option --ckpt or ONNX using option --onnx to build TRT BERT model.")
with build_engine(args.batch_size, args.workspace_size, args.sequence_length, config, weights_dict, args.squad_json, args.vocab_file, calib_cache, args.calib_num, args.verbose) as serialized_engine:
TRT_LOGGER.log(TRT_LOGGER.INFO, "Saving Engine to {:}".format(args.output))
with open(args.output, "wb") as fout:
fout.write(serialized_engine)
TRT_LOGGER.log(TRT_LOGGER.INFO, "Done.")
if __name__ == "__main__":
main()
demo/BERT/builder_utils.py
-396
View File
@@ -1,396 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import re
import pickle
import numpy as np
import onnx
import tensorrt as trt
import torch
try:
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
except ImportError as err:
import sys
sys.stderr.write("""Error: Failed to import tensorflow module ({})\n""".format(err))
sys.exit()
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
"""
Attentions Keys
"""
WQ = "self_query_kernel"
BQ = "self_query_bias"
WK = "self_key_kernel"
BK = "self_key_bias"
WV = "self_value_kernel"
BV = "self_value_bias"
WQKV = "self_qkv_kernel"
BQKV = "self_qkv_bias"
"""
Transformer Keys
"""
W_AOUT = "attention_output_dense_kernel"
B_AOUT = "attention_output_dense_bias"
AOUT_LN_BETA = "attention_output_layernorm_beta"
AOUT_LN_GAMMA = "attention_output_layernorm_gamma"
W_MID = "intermediate_dense_kernel"
B_MID = "intermediate_dense_bias"
W_LOUT = "output_dense_kernel"
B_LOUT = "output_dense_bias"
LOUT_LN_BETA = "output_layernorm_beta"
LOUT_LN_GAMMA = "output_layernorm_gamma"
"""
Squad Output Keys
"""
SQD_W = "squad_output_weights"
SQD_B = "squad_output_bias"
def load_tf_weights(inputbase, config):
"""
Load the weights from the tensorflow checkpoint
"""
weights_dict = dict()
try:
reader = tf.train.NewCheckpointReader(inputbase)
tensor_dict = reader.get_variable_to_shape_map()
# There might be training-related variables in the checkpoint that can be discarded
param_names = [key for key in sorted(tensor_dict) if "adam" not in key and "global_step" not in key and "pooler" not in key]
count = len(param_names)
TRT_LOGGER.log(TRT_LOGGER.INFO, "Found {:} entries in weight map".format(count))
for pn in param_names:
toks = pn.lower().split("/")
if "encoder" in pn:
assert ("layer" in pn)
l = (re.findall("\d+", pn))[0]
outname = "l{}_".format(l) + "_".join(toks[3:])
else:
outname = "_".join(toks)
tensor = reader.get_tensor(pn)
shape = tensor.shape
if pn.find("kernel") != -1:
weights_dict[outname + "_notrans"] = trt.Weights(np.ascontiguousarray(tensor).flatten())
TRT_LOGGER.log(TRT_LOGGER.VERBOSE, "Transposing {}\n".format(np))
tensor = np.transpose(tensor)
shape = tensor.shape
flat_tensor = tensor.flatten()
shape_str = "{} ".format(len(shape)) + " ".join([str(d) for d in shape])
weights_dict[outname] = trt.Weights(flat_tensor)
TRT_LOGGER.log(TRT_LOGGER.VERBOSE, "Original name: {:}, TensorRT name: {:}, shape: {:}".format(pn, outname, shape_str))
N = config.num_attention_heads
H = config.head_size
additional_dict = dict()
for key, value in weights_dict.items():
pos = key.find(BQ)
if pos != -1:
hidden_size = value.size
prefix = key[:pos]
Bq_ = value
Bk_ = weights_dict[prefix + BK]
Bv_ = weights_dict[prefix + BV]
Wq_ = weights_dict[prefix + WQ]
Wk_ = weights_dict[prefix + WK]
Wv_ = weights_dict[prefix + WV]
mat_size = hidden_size * hidden_size
wcount = 3 * mat_size
Wall = np.zeros(wcount, np.float32)
bcount = 3 * hidden_size
Ball = np.zeros(bcount, np.float32)
Wall[0:mat_size] = Wq_.numpy()[0:mat_size]
Wall[mat_size:2*mat_size] = Wk_.numpy()[0:mat_size]
Wall[2*mat_size:3*mat_size] = Wv_.numpy()[0:mat_size]
Ball[0:hidden_size] = Bq_.numpy()[0:hidden_size]
Ball[hidden_size:2*hidden_size] = Bk_.numpy()[0:hidden_size]
Ball[2*hidden_size:3*hidden_size] = Bv_.numpy()[0:hidden_size]
if config.use_int8 and getattr(config, 'interleaved', False):
Wall = np.ascontiguousarray(Wall.reshape((3, N, H, N, H)), dtype=np.float32)
Ball = np.ascontiguousarray(Ball.reshape((3, N, H)), dtype=np.float32)
else:
Wall = np.ascontiguousarray(Wall.reshape((3, N, H, N, H)).transpose((1, 0, 2, 3, 4)), dtype=np.float32)
Ball = np.ascontiguousarray(Ball.reshape((3, N, H)).transpose((1, 0, 2)), dtype=np.float32)
additional_dict[prefix + WQKV] = trt.Weights(Wall)
additional_dict[prefix + BQKV] = trt.Weights(Ball)
additional_dict[prefix + WQKV + "_notrans"] = trt.Weights(np.ascontiguousarray(Wall.T))
except Exception as error:
TRT_LOGGER.log(TRT_LOGGER.ERROR, str(error))
weights_dict.update(additional_dict)
return weights_dict
def onnx_to_trt_name(onnx_name):
"""
Converting variables in the onnx checkpoint to names corresponding to the naming convention used in the TF version, expected by the builder
"""
qkv_strings = {'key', 'value', 'query', 'query_key_value'}
onnx_name = onnx_name.lower()
toks = [t.strip('_') for t in onnx_name.split('.')]
if toks[0] == 'bert': #embeddings or encoder
if toks[1] == 'encoder': #transformer
# Token conversions for sparse checkpoints
if toks[-2] == 'dense_act':
toks[-2] = 'dense'
elif toks[-3] == 'dense_act':
if toks[-2] == 'input_quantizer':
toks[-2] = 'input'
elif toks[-2] == 'weight_quantizer':
toks[-2] = 'kernel'
toks[-3] = 'dense'
elif toks[-2].startswith('matmul'):
toks[-2] = {
'matmul_q_quantizer': 'qv_a_input_quantizer',
'matmul_k_quantizer': 'qv_b_input_quantizer',
'matmul_v_quantizer': 'av_b_input_quantizer',
'matmul_a_quantizer': 'av_a_input_quantizer',
}[toks[-2].replace('input_', '')]
# Token conversions for all checkpoints
if toks[-2] == 'layernorm': #bias->beta, weight->gamma
toks[-1] = 'beta' if toks[-1] == 'bias' else 'gamma'
elif (toks[-2] == 'dense' or toks[-2] in qkv_strings) and toks[-1] == 'weight':
toks[-1] = 'kernel'
elif (toks[-3] == 'dense' or toks[-3] in qkv_strings) and toks[-1] == 'amax':
if toks[-2] == 'weight_quantizer':
toks[-2] = 'kernel'
elif toks[-2] == 'input_quantizer':
toks[-2] = 'input'
if 'final_input_quantizer' not in toks[2]:
ind = toks.index('layers')+1 if 'layers' in toks else 3
toks = toks[ind:]
toks[0] = 'l{}'.format(int(toks[0]))
else:
if toks[-2] == 'layernorm': #bias->beta, weight->gamma
toks[-1] = 'beta' if toks[-1] == 'bias' else 'gamma'
else: #embeddings: drop "_weight" suffix
if toks[-1] == 'amax':
toks[-2] = 'amax'
toks = toks[:-1]
elif 'qa' in onnx_name:
name = 'cls_squad_output_bias' if toks[-1] == 'bias' else 'cls_squad_output_weights'
return name
else:
print("Encountered unknown case:", onnx_name)
assert(False)
parsed = '_'.join(toks)
return parsed
def get_onnx_weight_dict(tensor_dict, config):
N = config.num_attention_heads
H = config.head_size
hidden_size = config.hidden_size
weights_dict = dict()
for outname, tensor in tensor_dict.items():
if outname.find("_amax") != -1:
weights_dict[outname] = tensor
elif outname.find(BQ) != -1:
prefix = outname[:outname.find(BQ)]
Wqkv = np.zeros((3, hidden_size, hidden_size), np.float32)
Bqkv = np.zeros((3, hidden_size), np.float32)
Wqkv[0,:,:] = tensor_dict[prefix + WQ]
Wqkv[1,:,:] = tensor_dict[prefix + WK]
Wqkv[2,:,:] = tensor_dict[prefix + WV]
Bqkv[0,:] = tensor
Bqkv[1,:] = tensor_dict[prefix + BK]
Bqkv[2,:] = tensor_dict[prefix + BV]
if config.use_int8 and getattr(config, 'interleaved', False):
Wqkv = np.ascontiguousarray(Wqkv.reshape((3, N, H, N, H)))
Bqkv = np.ascontiguousarray(Bqkv.reshape((3, N, H)))
else:
Wqkv = np.ascontiguousarray(Wqkv.reshape((3, N, H, N, H)).transpose((1,0,2,3,4)))
Bqkv = np.ascontiguousarray(Bqkv.reshape((3, N, H)).transpose((1,0,2)))
weights_dict[prefix + WQKV] = trt.Weights(Wqkv)
weights_dict[prefix + BQKV] = trt.Weights(Bqkv)
weights_dict[prefix + WQKV + "_notrans"] = trt.Weights(np.ascontiguousarray(Wqkv.T))
elif outname.find(BK) != -1 or outname.find(BV) != -1 or outname.find(WQ) != -1 or outname.find(WK) != -1 or outname.find(WV) != -1:
pass
else:
flat_tensor = np.ascontiguousarray(tensor).flatten()
weights_dict[outname] = trt.Weights(flat_tensor)
if outname.find("kernel") != -1:
tensor = np.transpose(tensor)
weights_dict[outname + "_notrans"] = trt.Weights(np.ascontiguousarray(tensor).flatten())
TRT_LOGGER.log(TRT_LOGGER.INFO, "Found {:} entries in weight map".format(len(weights_dict)))
return weights_dict
def load_onnx_weights_and_quant(path, config):
"""
Load the weights from the onnx checkpoint
"""
model = onnx.load(path)
weights = model.graph.initializer
tensor_dict = dict((onnx_to_trt_name(w.name), np.frombuffer(w.raw_data, np.int8).reshape(w.dims))
if w.name.split('_')[-1] == 'mask' else
(onnx_to_trt_name(w.name), np.frombuffer(w.raw_data, np.float32).reshape(w.dims))
for w in weights)
return get_onnx_weight_dict(tensor_dict, config)
def load_pytorch_weights_and_quant(path, config):
"""
Load the weights from the pytorch checkpoint
"""
state_dict = torch.load(path, map_location='cpu')["model"]
tensor_dict = {onnx_to_trt_name(name):val.numpy() for name, val in state_dict.items()}
return get_onnx_weight_dict(tensor_dict, config)
def load_megatron_pickle_weights(path, config):
N = config.num_attention_heads
H = config.head_size
with open(path, 'rb') as f:
tensor_dict = pickle.load(f)
weight_dict = {}
for name, tensor in tensor_dict.items():
if 'scale' in name:
continue
name = (onnx_to_trt_name(name)
.replace('embedding_', 'embeddings_')
.replace('tokentype_', 'token_type_')
.replace('_av', '_self_av')
.replace('_qv', '_self_qv')
.replace('query_key_value', 'self_qkv'))
if name.endswith('self_qkv_kernel'):
tensor = np.ascontiguousarray(tensor.reshape((3, N, H, N, H))).astype(np.float32)
weight_dict[name] = trt.Weights(tensor)
elif name.endswith('self_qkv_bias'):
tensor = np.ascontiguousarray(tensor.reshape((3, N, H))).astype(np.float32)
weight_dict[name] = trt.Weights(tensor)
elif name == 'l{}_output_layernorm_output_quantizer_amax'.format(config.num_hidden_layers-1):
weight_dict['bert_encoder_final_input_quantizer_amax'] = tensor
elif name.endswith('_amax'):
weight_dict[name] = tensor
if name.endswith('_qkv_input_amax'):
weight_dict[name.replace('_qkv_input_amax', '_query_input_amax')] = tensor
weight_dict[name.replace('_qkv_input_amax', '_key_input_amax')] = tensor
weight_dict[name.replace('_qkv_input_amax', '_value_input_amax')] = tensor
else:
flat_tensor = np.ascontiguousarray(tensor).flatten().astype(np.float32)
weight_dict[name] = trt.Weights(flat_tensor)
TRT_LOGGER.log(TRT_LOGGER.INFO, "Found {:} entries in weight map".format(len(weight_dict)))
return weight_dict
"""
Common Plugin Helper/Wrapper Functions
"""
BERT_PLUGINS_INFO_MAP = {
# MHA variants
"qkv_to_context": {
"IPluginV2_version": "1",
"IPluginV3_version": "4",
"trt_plugin_name": "CustomQKVToContextPluginDynamic",
},
"qkv_to_context_varseqlen": {
"IPluginV2_version": "2",
"IPluginV3_version": "5",
"trt_plugin_name": "CustomQKVToContextPluginDynamic",
},
"qkv_to_context_interleaved": {
"IPluginV2_version": "3",
"IPluginV3_version": "6",
"trt_plugin_name": "CustomQKVToContextPluginDynamic",
},
# skipLayernorm variants
"skip_layer_norm": {
"IPluginV2_version": "1",
"IPluginV3_version": "5",
"trt_plugin_name": "CustomSkipLayerNormPluginDynamic",
},
"skip_layer_norm_varseqlen": {
"IPluginV2_version": "2",
"IPluginV3_version": "6",
"trt_plugin_name": "CustomSkipLayerNormPluginDynamic",
},
"skip_layer_norm_huggingface": {
"IPluginV2_version": "3",
"IPluginV3_version": "7",
"trt_plugin_name": "CustomSkipLayerNormPluginDynamic",
},
"skip_layer_norm_megatron": {
"IPluginV2_version": "4",
"IPluginV3_version": "8",
"trt_plugin_name": "CustomSkipLayerNormPluginDynamic",
},
# embLayernorm variants
"emb_layer_norm": {
"IPluginV2_version": "1",
"IPluginV3_version": "6",
"trt_plugin_name": "CustomEmbLayerNormPluginDynamic",
},
"emb_layer_norm_huggingface": {
"IPluginV2_version": "2",
"IPluginV3_version": "4",
"trt_plugin_name": "CustomEmbLayerNormPluginDynamic",
},
"emb_layer_norm_megatron": {
"IPluginV2_version": "3",
"IPluginV3_version": "5",
"trt_plugin_name": "CustomEmbLayerNormPluginDynamic",
},
}
def create_plugin(layer_name, plg_registry, pfc, use_deprecated_plugins=False):
plg_trt_name = BERT_PLUGINS_INFO_MAP[layer_name]["trt_plugin_name"]
plg_version = BERT_PLUGINS_INFO_MAP[layer_name][
("IPluginV2_version" if use_deprecated_plugins else "IPluginV3_version")
]
plg_namespace = ""
creator = plg_registry.get_creator(plg_trt_name, plg_version, plg_namespace)
if use_deprecated_plugins:
return creator.create_plugin(layer_name, pfc)
else:
return creator.create_plugin(layer_name, pfc, trt.TensorRTPhase.BUILD)
def add_plugin_to_network(network, plugin, inputs, use_deprecated_plugins=False):
if use_deprecated_plugins:
return network.add_plugin_v2(inputs, plugin)
else:
return network.add_plugin_v3(inputs, [], plugin)
demo/BERT/builder_varseqlen.py
-696
View File
@@ -1,696 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse
import ctypes
import json
import numpy as np
import os
import os.path
import re
import sys
import time
import onnx
from helpers.cuda_utils import getComputeCapacity
# TensorRT
import tensorrt as trt
from builder_utils import load_tf_weights, load_pytorch_weights_and_quant, load_onnx_weights_and_quant, load_megatron_pickle_weights
from builder_utils import WQKV, BQKV # Attention Keys
from builder_utils import W_AOUT, B_AOUT, W_MID, B_MID, W_LOUT, B_LOUT # Transformer Keys
from builder_utils import SQD_W, SQD_B # SQuAD Output Keys
from builder_utils import (
create_plugin,
add_plugin_to_network,
) # Plugin Helper functions
"""
TensorRT Initialization
"""
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
trt_version = [n for n in trt.__version__.split('.')]
# Import necessary plugins for demoBERT
plugin_lib_name = "nvinfer_plugin_10.dll" if sys.platform == "win32" else "libnvinfer_plugin.so"
env_name_to_add_path = "PATH" if sys.platform == "win32" else "LD_LIBRARY_PATH"
handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)
if not handle:
raise RuntimeError("Could not load plugin library. Is `{}` on your {}?".format(plugin_lib_name, env_name_to_add_path))
trt.init_libnvinfer_plugins(TRT_LOGGER, "")
plg_registry = trt.get_plugin_registry()
class BertConfig:
def __init__(
self,
bert_config_path,
use_fp16,
use_int8,
use_qat,
interleaved,
timing_cache,
use_sparsity,
use_megatron,
use_deprecated_plugins=False,
):
with open(bert_config_path, "r") as f:
data = json.load(f)
self.num_attention_heads = data["num_attention_heads"]
self.hidden_size = data["hidden_size"]
self.intermediate_size = data["intermediate_size"]
self.num_hidden_layers = data["num_hidden_layers"]
self.head_size = self.hidden_size // self.num_attention_heads
self.use_fp16 = use_fp16
self.use_int8 = use_int8
self.use_qat = use_qat
self.interleaved = interleaved
self.timing_cache = timing_cache
self.use_sparsity = use_sparsity
self.use_megatron = use_megatron
self.use_deprecated_plugins = use_deprecated_plugins
def get_trt_dtype(self):
dtype = trt.float32
if self.use_fp16:
dtype = trt.float16
if self.use_int8:
dtype = trt.int8
return dtype
def set_tensor_name(tensor, prefix, name):
tensor.name = prefix + name
def set_output_name(layer, prefix, name, out_idx = 0):
set_tensor_name(layer.get_output(out_idx), prefix, name)
def set_output_range(layer, maxval, out_idx = 0):
layer.get_output(out_idx).set_dynamic_range(-maxval, maxval)
def attention_layer_opt(prefix, config, init_dict, network, input_tensor, mask_idx, cu_seqlens, max_seqlen):
"""
Add the attention layer
"""
hidden_size = config.hidden_size
num_heads = config.num_attention_heads
head_size = int(hidden_size / num_heads)
Wall = init_dict[prefix + WQKV]
Ball = init_dict[prefix + BQKV]
# FC_attention
mult_all = network.add_convolution_nd(input_tensor, 3 * hidden_size, (1, 1), Wall, Ball)
if config.use_qat:
dr_qkv = max(
init_dict[prefix + 'self_qv_a_input_quantizer_amax'],
init_dict[prefix + 'self_qv_b_input_quantizer_amax'],
init_dict[prefix + 'self_av_b_input_quantizer_amax'],
)
set_output_range(mult_all, dr_qkv)
set_output_name(mult_all, prefix, "qkv_mult")
# QKV2CTX
dtype = config.get_trt_dtype()
pf_type = trt.PluginField("type_id", np.array([int(dtype)], np.int32), trt.PluginFieldType.INT32)
pf_hidden_size = trt.PluginField("hidden_size", np.array([hidden_size], np.int32), trt.PluginFieldType.INT32)
pf_num_heads = trt.PluginField("num_heads", np.array([num_heads], np.int32), trt.PluginFieldType.INT32)
pf_has_mask = trt.PluginField("has_mask", np.array([1], np.int32), trt.PluginFieldType.INT32)
pf_var_seqlen = trt.PluginField("var_seqlen", np.array([int(1)], np.int32), trt.PluginFieldType.FLOAT32)
if config.use_qat:
dr_probs = init_dict[prefix + 'self_av_a_input_quantizer_amax']
dq_probs = dr_probs / 127.0
pf_dq_probs = trt.PluginField("dq_probs", np.array([dq_probs], np.float32), trt.PluginFieldType.FLOAT32)
fields = [pf_hidden_size, pf_num_heads, pf_dq_probs]
else:
fields = [pf_hidden_size, pf_num_heads]
if config.use_int8 and config.interleaved:
pfc = trt.PluginFieldCollection(fields)
qkv2ctx_plug = create_plugin(
"qkv_to_context_interleaved",
plg_registry,
pfc,
use_deprecated_plugins=config.use_deprecated_plugins,
)
qkv_in = [mult_all.get_output(0), cu_seqlens, max_seqlen]
else:
fields.append(pf_has_mask)
fields.append(pf_type)
fields.append(pf_var_seqlen)
pfc = trt.PluginFieldCollection(fields)
qkv2ctx_plug = create_plugin(
"qkv_to_context_varseqlen",
plg_registry,
pfc,
use_deprecated_plugins=config.use_deprecated_plugins,
)
qkv_in = [mult_all.get_output(0), mask_idx, cu_seqlens, max_seqlen]
qkv2ctx = add_plugin_to_network(
network, qkv2ctx_plug, qkv_in, use_deprecated_plugins=config.use_deprecated_plugins
)
qkv2ctx.name = prefix + "qkv_to_ctx"
if config.use_qat:
dr_ctx = init_dict[prefix + 'output_dense_input_amax']
set_output_range(qkv2ctx, dr_ctx)
set_output_name(qkv2ctx, prefix, "context_layer")
return qkv2ctx
def skipln(prefix, config, init_dict, network, input_tensor, skip, is_last_skipln=False):
"""
Add the skip layer
"""
hidden_size = config.hidden_size
dtype = config.get_trt_dtype()
pf_ld = trt.PluginField("ld", np.array([hidden_size], np.int32), trt.PluginFieldType.INT32)
wbeta = init_dict[prefix + "beta"]
pf_beta = trt.PluginField("beta", wbeta.numpy(), trt.PluginFieldType.FLOAT32)
wgamma = init_dict[prefix + "gamma"]
pf_gamma = trt.PluginField("gamma", wgamma.numpy(), trt.PluginFieldType.FLOAT32)
pf_type = trt.PluginField("type_id", np.array([int(dtype)], np.int32), trt.PluginFieldType.INT32)
if config.use_int8 and config.interleaved:
pfc = trt.PluginFieldCollection([pf_beta, pf_gamma])
variant_name = (
"skip_layer_norm_huggingface"
if not config.use_megatron or is_last_skipln
else "skip_layer_norm_megatron"
)
skipln_plug = create_plugin(
variant_name, plg_registry, pfc, use_deprecated_plugins=config.use_deprecated_plugins
)
else:
pfc = trt.PluginFieldCollection([pf_ld, pf_beta, pf_gamma, pf_type])
skipln_plug = create_plugin(
"skip_layer_norm_varseqlen",
plg_registry,
pfc,
use_deprecated_plugins=config.use_deprecated_plugins,
)
skipln_inputs = [input_tensor, skip]
layer = add_plugin_to_network(
network, skipln_plug, skipln_inputs, use_deprecated_plugins=config.use_deprecated_plugins
)
return layer
def transformer_layer_opt(prefix, config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen):
"""
Add the transformer layer
"""
hidden_size = config.hidden_size
if config.use_qat:
dr_input = init_dict[prefix + 'attention_self_query_input_amax']
assert(dr_input ==init_dict[prefix + 'attention_self_key_input_amax'] )
assert(dr_input ==init_dict[prefix + 'attention_self_value_input_amax'] )
input_tensor.set_dynamic_range(-dr_input, dr_input)
context_transposed = attention_layer_opt(prefix + "attention_", config, init_dict, network, input_tensor, mask_idx, cu_seqlens, max_seqlen)
attention_heads = context_transposed.get_output(0)
# FC0
B_aout = init_dict[prefix + B_AOUT]
W_aout = init_dict[prefix + W_AOUT]
attention_out_fc = network.add_convolution_nd(attention_heads, hidden_size, (1, 1), W_aout, B_aout)
if config.use_int8 and config.use_qat:
dr_fc_aout = init_dict[prefix + 'attention_output_add_local_input_quantizer_amax']
set_output_range(attention_out_fc, dr_fc_aout)
if config.use_megatron:
dr_skln1_res_in = init_dict[prefix + "attention_output_add_residual_input_quantizer_amax"]
residual.set_dynamic_range(-dr_skln1_res_in, dr_skln1_res_in)
skip = residual
else:
skip = input_tensor
skiplayer = skipln(prefix + "attention_output_layernorm_", config, init_dict, network, attention_out_fc.get_output(0), skip)
attention_ln = skiplayer.get_output(0)
if config.use_qat:
dr_skln1 = init_dict[prefix + 'intermediate_dense_input_amax']
set_output_range(skiplayer, dr_skln1)
# FC1 + GELU
B_mid = init_dict[prefix + B_MID]
W_mid = init_dict[prefix + W_MID]
mid_dense = network.add_convolution_nd(attention_ln, config.intermediate_size, (1, 1), W_mid, B_mid)
gelu_layer = add_gelu(network, mid_dense.get_output(0))
intermediate_act = gelu_layer.get_output(0)
set_tensor_name(intermediate_act, prefix, "gelu")
if config.use_int8:
if config.use_qat:
dr_gelu = init_dict[prefix + 'output_dense_input_amax']
set_output_range(gelu_layer, dr_gelu)
else:
# use gelu10 according to whitepaper http://arxiv.org/abs/2004.09602
set_output_range(gelu_layer, 10)
# FC2
# Dense to hidden size
B_lout = init_dict[prefix + B_LOUT]
W_lout = init_dict[prefix + W_LOUT]
out_dense = network.add_convolution_nd(intermediate_act, hidden_size, (1, 1), W_lout, B_lout)
if config.use_int8 and config.use_qat:
dr_fc_out = init_dict[prefix + 'output_add_local_input_quantizer_amax']
set_output_range(out_dense, dr_fc_out)
set_output_name(out_dense, prefix + "output_", "dense")
if config.use_megatron:
dr_skln2_res_in = init_dict[prefix + 'output_add_residual_input_quantizer_amax']
set_output_range(skiplayer, dr_skln2_res_in, out_idx=1)
skip = skiplayer.get_output(1)
else:
skip = attention_ln
is_last_skipln = prefix.startswith('l{}'.format(config.num_hidden_layers-1))
out_layer = skipln(prefix + "output_layernorm_", config, init_dict, network, out_dense.get_output(0), skip, is_last_skipln)
set_output_name(out_layer, prefix + "output_", "reshape")
return out_layer
def add_gelu(network, input_tensor):
"""
Adds elementwise GELU, and will trigger FC+GELU fusion in TRT
"""
shape = (1, ) * len(input_tensor.shape)
POW = network.add_constant(shape, trt.Weights(np.ascontiguousarray([3.0], dtype=np.float32)))
MULTIPLY = network.add_constant(shape, trt.Weights(np.ascontiguousarray([0.044715], dtype=np.float32)))
SQRT = network.add_constant(shape, trt.Weights((np.ascontiguousarray([0.79788456080286535587989211986876], dtype=np.float32))))
ONE = network.add_constant(shape, trt.Weights((np.ascontiguousarray([1.0], dtype=np.float32))))
HALF = network.add_constant(shape, trt.Weights((np.ascontiguousarray([0.5], dtype=np.float32))))
X_pow = network.add_elementwise(input_tensor, POW.get_output(0), trt.ElementWiseOperation.POW)
X_pow_t = X_pow.get_output(0)
X_mul = network.add_elementwise(X_pow_t, MULTIPLY.get_output(0), trt.ElementWiseOperation.PROD)
X_add = network.add_elementwise(input_tensor, X_mul.get_output(0), trt.ElementWiseOperation.SUM)
X_sqrt = network.add_elementwise(X_add.get_output(0), SQRT.get_output(0), trt.ElementWiseOperation.PROD)
X_sqrt_tensor = X_sqrt.get_output(0)
X_tanh = network.add_activation(X_sqrt_tensor, trt.ActivationType.TANH)
X_tanh_tensor = X_tanh.get_output(0)
X_one = network.add_elementwise(X_tanh_tensor, ONE.get_output(0), trt.ElementWiseOperation.SUM)
CDF = network.add_elementwise(X_one.get_output(0), HALF.get_output(0), trt.ElementWiseOperation.PROD)
gelu_layer = network.add_elementwise(CDF.get_output(0), input_tensor, trt.ElementWiseOperation.PROD)
# enable elementwise fusing for int8 && fp16
POW.precision = trt.DataType.FLOAT
MULTIPLY.precision = trt.DataType.FLOAT
SQRT.precision = trt.DataType.FLOAT
ONE.precision = trt.DataType.FLOAT
HALF.precision = trt.DataType.FLOAT
X_pow.precision = trt.DataType.FLOAT
X_mul.precision = trt.DataType.FLOAT
X_add.precision = trt.DataType.FLOAT
X_sqrt.precision = trt.DataType.FLOAT
X_tanh.precision = trt.DataType.FLOAT
X_one.precision = trt.DataType.FLOAT
CDF.precision = trt.DataType.FLOAT
gelu_layer.precision = trt.DataType.FLOAT
return gelu_layer
def bert_model(config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen):
"""
Create the bert model
"""
prev_input = input_tensor
for layer in range(0, config.num_hidden_layers):
ss = "l{}_".format(layer)
out_layer = transformer_layer_opt(ss, config, init_dict, network, prev_input, residual, mask_idx, cu_seqlens, max_seqlen)
prev_input = out_layer.get_output(0)
# Skip reading residual from final layer
if config.use_megatron and (layer != config.num_hidden_layers - 1):
residual = out_layer.get_output(1)
if config.use_qat:
dr_out = init_dict["bert_encoder_final_input_quantizer_amax"]
set_output_range(out_layer, dr_out)
squad_logits = squad_output("cls_", config, init_dict, network, prev_input)
squad_logits_out = squad_logits.get_output(0)
squad_logits_out.name = "logits_out"
network.mark_output(squad_logits_out)
def squad_output(prefix, config, init_dict, network, input_tensor):
"""
Create the squad output
"""
hidden_size = config.hidden_size
W_out = init_dict[prefix + SQD_W]
B_out = init_dict[prefix + SQD_B]
dense = network.add_convolution_nd(input_tensor, 2, (1, 1), W_out, B_out)
OUT = network.add_shuffle(dense.get_output(0))
if config.use_int8 and config.interleaved:
OUT.second_transpose = (1, 2, 0, 3)
else:
OUT.second_transpose = (1, 0, 2, 3)
set_output_name(OUT, prefix, "squad_logits")
return OUT
def emb_layernorm(builder, network, config, weights_dict, builder_config, max_sequence_length, batch_sizes):
input_ids = network.add_input(name="input_ids", dtype=trt.int32, shape=(-1,))
segment_ids = network.add_input(name="segment_ids", dtype=trt.int32, shape=(-1,))
cu_seqlens = network.add_input(name="cu_seqlens", dtype=trt.int32, shape=(-1,))
max_seqlen = network.add_input(name="max_seqlen", dtype=trt.int32, shape=(-1,))
for batch_size in batch_sizes:
# Specify profiles
profile = builder.create_optimization_profile()
min_shape = (1,)
shape = (max_sequence_length*batch_size,)
profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape)
profile.set_shape("cu_seqlens", min=min_shape, opt=(batch_size+1,), max=(batch_size+1,))
profile.set_shape("max_seqlen", min=min_shape, opt=(max_sequence_length,), max=(max_sequence_length,))
builder_config.add_optimization_profile(profile)
wbeta = trt.PluginField("bert_embeddings_layernorm_beta", weights_dict["bert_embeddings_layernorm_beta"].numpy(), trt.PluginFieldType.FLOAT32)
wgamma = trt.PluginField("bert_embeddings_layernorm_gamma", weights_dict["bert_embeddings_layernorm_gamma"].numpy(), trt.PluginFieldType.FLOAT32)
wwordemb = trt.PluginField("bert_embeddings_word_embeddings", weights_dict["bert_embeddings_word_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
wtokemb = trt.PluginField("bert_embeddings_token_type_embeddings", weights_dict["bert_embeddings_token_type_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
wposemb = trt.PluginField("bert_embeddings_position_embeddings", weights_dict["bert_embeddings_position_embeddings"].numpy(), trt.PluginFieldType.FLOAT32)
output_fp16 = trt.PluginField("output_fp16", np.array([1 if config.use_fp16 or config.use_int8 else 0]).astype(np.int32), trt.PluginFieldType.INT32)
pfc = trt.PluginFieldCollection(
[wbeta, wgamma, wwordemb, wtokemb, wposemb, output_fp16]
)
variant_name = (
"emb_layer_norm_megatron"
if config.use_megatron
else "emb_layer_norm_huggingface"
)
fn = create_plugin(
variant_name, plg_registry, pfc, use_deprecated_plugins=config.use_deprecated_plugins
)
inputs = [input_ids, segment_ids, cu_seqlens, max_seqlen]
emb_layer = add_plugin_to_network(
network, fn, inputs, use_deprecated_plugins=config.use_deprecated_plugins
)
if config.use_int8 and config.use_qat:
dr_input = weights_dict['l0_attention_self_query_input_amax']
set_output_range(emb_layer, dr_input, out_idx=0)
if config.use_megatron:
dr_skln1_res_in = weights_dict['l0_attention_output_add_residual_input_quantizer_amax']
set_output_range(emb_layer, dr_skln1_res_in, out_idx=1)
set_output_name(emb_layer, "embeddings_", "output")
return emb_layer, cu_seqlens, max_seqlen
def build_engine(batch_sizes, workspace_size, sequence_length, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose):
network_creation_flag = 0
if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys():
network_creation_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(network_creation_flag) as network, builder.create_builder_config() as builder_config:
if workspace_size is not None:
builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size * (1024 * 1024))
builder_config.avg_timing_iterations = 8
if config.use_fp16:
builder_config.set_flag(trt.BuilderFlag.FP16)
if config.use_int8:
builder_config.set_flag(trt.BuilderFlag.INT8)
if not config.use_qat:
raise RuntimeError("Post training calibration is not supported in variable-length BERT.")
if verbose:
builder_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
# speed up the engine build for trt major version >= 8
# 1. disable cudnn tactic
# 2. load global timing cache
if int(trt_version[0]) >= 8:
tactic_source = builder_config.get_tactic_sources() & ~(1 << int(trt.TacticSource.CUDNN))
builder_config.set_tactic_sources(tactic_source)
if config.timing_cache != None:
if os.path.exists(config.timing_cache):
with open(config.timing_cache, "rb") as f:
cache = builder_config.create_timing_cache(f.read())
builder_config.set_timing_cache(cache, ignore_mismatch = False)
else:
cache = builder_config.create_timing_cache(b"")
builder_config.set_timing_cache(cache, ignore_mismatch = False)
if config.use_sparsity:
TRT_LOGGER.log(TRT_LOGGER.INFO, "Setting sparsity flag on builder_config.")
builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
# Create the network
emb_layer, cu_seqlens, max_seqlen = emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_length, batch_sizes)
embeddings = emb_layer.get_output(0)
if config.use_int8 and config.interleaved:
shuffle = network.add_shuffle(embeddings)
shuffle.second_transpose = (2, 1, 0, 3)
embeddings = shuffle.get_output(0)
mask_idx = None
else:
mask_idx = emb_layer.get_output(1)
if config.use_megatron: # megatron currently only supports int8 and interleaved
shuffler = network.add_shuffle(emb_layer.get_output(1))
shuffler.second_transpose = (2, 1, 0, 3)
residual = shuffler.get_output(0)
dr_emb = weights_dict['l0_attention_self_query_input_amax']
embeddings.set_dynamic_range(-dr_emb, dr_emb)
dr_skln1_res_in = weights_dict['l0_attention_output_add_residual_input_quantizer_amax']
residual.set_dynamic_range(-dr_skln1_res_in, dr_skln1_res_in)
else:
residual = None
bert_model(config, weights_dict, network, embeddings, residual, mask_idx, cu_seqlens, max_seqlen)
build_start_time = time.time()
serialized_engine = builder.build_serialized_network(network, builder_config)
build_time_elapsed = (time.time() - build_start_time)
TRT_LOGGER.log(TRT_LOGGER.INFO, "build engine in {:.3f} Sec".format(build_time_elapsed))
# save global timing cache
if int(trt_version[0]) >= 8 and config.timing_cache != None:
cache = builder_config.get_timing_cache()
with cache.serialize() as buffer:
with open(config.timing_cache, "wb") as f:
f.write(buffer)
f.flush()
os.fsync(f)
return serialized_engine
def main():
parser = argparse.ArgumentParser(
description="TensorRT BERT Sample",
formatter_class=argparse.RawTextHelpFormatter,
)
parser.add_argument(
"-m",
"--ckpt",
required=False,
help="The checkpoint file basename, e.g.: basename(model.ckpt-766908.data-00000-of-00001) is model.ckpt-766908 (default: None)",
)
parser.add_argument(
"-x", "--onnx", required=False, help="The ONNX model file path. (default: None)"
)
parser.add_argument(
"-pt", "--pytorch", required=False, help="The PyTorch checkpoint file path. (default: None)"
)
parser.add_argument(
"-pkl",
"--pickle",
required=False,
help="The Pickle weights dictionary file path for the Megatron variant of BERT. (default: None)",
)
parser.add_argument(
"-o",
"--output",
required=True,
default="bert_base_384.engine",
help="The bert engine file, ex bert.engine (default: bert_base_384.engine)",
)
parser.add_argument(
"-b",
"--max-batch-size",
default=[],
action="append",
help="Max batch size. The engine will be usable with any input with (batch-size * sequence-length) below (max-batch-size * max-sequence-length). Can be specified multiple times to build optimization profiles for more than one batch size. (default: [1])",
type=int,
)
parser.add_argument(
"-s",
"--max-sequence-length",
default=128,
help="Max sequence length of the BERT model. The engine will be usable with any input with (batch-size * sequence-length) below (max-batch-size * max-sequence-length). (default: 128)",
type=int,
)
parser.add_argument(
"-c",
"--config-dir",
required=True,
help="The folder containing the bert_config.json, which can be downloaded e.g. from https://github.com/google-research/bert#pre-trained-models or by running download_models.py in dle/TensorFlow/LanguageModeling/BERT/data/pretrained_models_google",
)
parser.add_argument(
"-f",
"--fp16",
action="store_true",
help="Indicates that inference should be run in FP16 precision (default: false)",
required=False,
)
parser.add_argument(
"-i",
"--int8",
action="store_true",
help="Indicates that inference should be run in INT8 precision (default: false)",
required=False,
)
parser.add_argument(
"-w",
"--workspace-size",
help="Workspace size in MiB for building the BERT engine (default: unlimited)",
type=int,
)
parser.add_argument(
"-j",
"--squad-json",
default="squad/dev-v1.1.json",
help="squad json dataset used for int8 calibration (default: squad/dev-v1.1.json)",
required=False,
)
parser.add_argument(
"-v",
"--vocab-file",
default="./pre-trained_model/uncased_L-24_H-1024_A-16/vocab.txt",
help="Path to file containing entire understandable vocab (default: ./pre-trained_model/uncased_L-24_H-1024_A-16/vocab.txt)",
required=False,
)
parser.add_argument(
"-n", "--calib-num", default=100, help="calibration batch numbers (default: 100)", type=int
)
parser.add_argument(
"-p", "--calib-path", help="calibration cache path (default: None)", required=False
)
parser.add_argument(
"-il",
"--interleaved",
action="store_true",
help="use interleaved format, only valid in INT8 precision (default: false)",
required=False,
)
parser.add_argument(
"-tcf",
"--timing-cache-file",
help="Path to tensorrt build timeing cache file, only available for tensorrt 8.0 and later (default: None)",
required=False,
)
parser.add_argument(
"-sp",
"--sparse",
action="store_true",
help="Indicates that model is sparse (default: false)",
required=False,
)
parser.add_argument(
"--megatron",
action="store_true",
help="Indicates that model is the Megatron-style architecture (default: false)",
required=False,
)
parser.add_argument(
"--verbose",
action="store_true",
help="Turn on verbose logger and set profiling verbosity to verbose (default: false)",
required=False,
)
plugin_group = parser.add_mutually_exclusive_group(required=False)
plugin_group.add_argument('--use-v3-plugins',
dest='use_deprecated_plugins',
action='store_false',
help="Use plugins implementing the IPluginV3 interface wherever TensorRT plugins are used. Cannot be used with --use-deprecated-plugins. Enabling this option should not affect functionality or performance. (default: false)")
plugin_group.add_argument('--use-deprecated-plugins',
dest='use_deprecated_plugins',
action='store_true',
help="Use deprecated plugins implementing the IPluginV2 interface wherever TensorRT plugins are used (instead of updated plugins implementing the IPluginV3 interface). Cannot be used with --use-v3-plugins. Disabling this option should not affect functionality or performance. (default: true)")
parser.set_defaults(use_deprecated_plugins=True)
args, _ = parser.parse_known_args()
args.max_batch_size = args.max_batch_size or [1]
if args.verbose:
TRT_LOGGER.min_severity = TRT_LOGGER.VERBOSE
cc = getComputeCapacity()
if cc[0] * 10 + cc[1] < 72:
raise RuntimeError("This variable-length BERT demo only support Xavier+ GPU.")
if args.megatron:
if not (args.interleaved and args.int8):
raise RuntimeError("Megatron BERT currently only supports int8 and interleaved.")
if not args.pickle:
raise RuntimeError("Megatron BERT currently only supports loading a pickle weights dictionary.")
bert_config_path = os.path.join(args.config_dir, "bert_config.json")
TRT_LOGGER.log(TRT_LOGGER.INFO, "Using configuration file: {:}".format(bert_config_path))
config = BertConfig(
bert_config_path,
args.fp16,
args.int8,
args.int8 and (args.onnx or args.pytorch or args.pickle),
args.interleaved,
args.timing_cache_file,
args.sparse,
args.megatron,
args.use_deprecated_plugins,
)
if args.calib_path != None:
calib_cache = args.calib_path
else:
calib_cache = "BertSquadL{}H{}A{}S{}CalibCache".format(config.num_hidden_layers, config.head_size, config.num_attention_heads, args.max_sequence_length)
if args.onnx != None:
weights_dict = load_onnx_weights_and_quant(args.onnx, config)
elif args.pytorch != None:
weights_dict = load_pytorch_weights_and_quant(args.pytorch, config)
elif args.ckpt != None:
weights_dict = load_tf_weights(args.ckpt, config)
elif args.pickle != None:
weights_dict = load_megatron_pickle_weights(args.pickle, config)
else:
raise RuntimeError("You need either specify TF checkpoint using option --ckpt, ONNX using option --onnx, "
"PyTorch using option --pytorch, or Pickle weight dictionary using option --pickle "
"to build TRT BERT model.")
with build_engine(args.max_batch_size, args.workspace_size, args.max_sequence_length, config, weights_dict, args.squad_json, args.vocab_file, calib_cache, args.calib_num, args.verbose) as serialized_engine:
TRT_LOGGER.log(TRT_LOGGER.INFO, "Saving Engine to {:}".format(args.output))
with open(args.output, "wb") as fout:
fout.write(serialized_engine)
TRT_LOGGER.log(TRT_LOGGER.INFO, "Done.")
if __name__ == "__main__":
main()
demo/BERT/helpers/__init__.py
View File
demo/BERT/helpers/calibrator.py
-113
View File
@@ -1,113 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import tensorrt as trt
import os
from helpers.cuda_utils import cuda_call, memcpy_host_to_device
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
import helpers.tokenization as tokenization
import helpers.data_processing as dp
class BertCalibrator(trt.IInt8LegacyCalibrator):
def __init__(self, squad_json, vocab_file, cache_file, batch_size, max_seq_length, num_inputs):
# Whenever you specify a custom constructor for a TensorRT class,
# you MUST call the constructor of the parent explicitly.
trt.IInt8LegacyCalibrator.__init__(self)
self.cache_file = cache_file
# Every time get_batch is called, the next batch of size batch_size will be copied to the device and returned.
self.data = dp.read_squad_json(squad_json)
self.max_seq_length = max_seq_length
self.batch_size = batch_size
self.current_index = 0
self.num_inputs = num_inputs
self.tokenizer = tokenization.BertTokenizer(vocab_file=vocab_file, do_lower_case=True)
self.doc_stride = 128
self.max_query_length = 64
# Allocate enough memory for a whole batch.
self.device_inputs = [cuda_call(cudart.cudaMalloc(self.max_seq_length * trt.int32.itemsize * self.batch_size)) for binding in range(3)]
def free(self):
for dinput in self.device_inputs:
# dinput is a device pointer (int) returned by cudaMalloc
cuda_call(cudart.cudaFree(dinput))
def get_batch_size(self):
return self.batch_size
# TensorRT passes along the names of the engine bindings to the get_batch function.
# You don't necessarily have to use them, but they can be useful to understand the order of
# the inputs. The bindings list is expected to have the same ordering as 'names'.
def get_batch(self, names):
if self.current_index + self.batch_size > self.num_inputs:
print("Calibrating index {:} batch size {:} exceed max input limit {:} sentences".format(self.current_index, self.batch_size, self.num_inputs))
return None
current_batch = int(self.current_index / self.batch_size)
if current_batch % 10 == 0:
print("Calibrating batch {:}, containing {:} sentences".format(current_batch, self.batch_size))
input_ids = []
segment_ids = []
input_mask = []
for i in range(self.batch_size):
example = self.data[self.current_index + i]
features = dp.convert_example_to_features(example.doc_tokens, example.question_text, self.tokenizer, self.max_seq_length, self.doc_stride, self.max_query_length)
if len(input_ids) and len(segment_ids) and len(input_mask):
input_ids = np.concatenate((input_ids, features[0].input_ids))
segment_ids = np.concatenate((segment_ids, features[0].segment_ids))
input_mask = np.concatenate((input_mask, features[0].input_mask))
else:
input_ids = features[0].input_ids
segment_ids = features[0].segment_ids
input_mask = features[0].input_mask
memcpy_host_to_device(self.device_inputs[0], input_ids.ravel())
memcpy_host_to_device(self.device_inputs[1], segment_ids.ravel())
memcpy_host_to_device(self.device_inputs[2], input_mask.ravel())
self.current_index += self.batch_size
return self.device_inputs
def read_calibration_cache(self):
# If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
if os.path.exists(self.cache_file):
with open(self.cache_file, "rb") as f:
return f.read()
def write_calibration_cache(self, cache):
with open(self.cache_file, "wb") as f:
f.write(cache)
f.flush()
os.fsync(f)
def get_quantile(self):
return 0.9999
def get_regression_cutoff(self):
return 1.0
def read_histogram_cache(self, length):
return None
def write_histogram_cache(self, ptr, length):
return None
demo/BERT/helpers/cuda_utils.py
-120
View File
@@ -1,120 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import numpy as np
import logging
from cuda.bindings import driver as cuda, runtime as cudart
class CudaStreamContext:
"""CUDA stream lifecycle management with context manager support"""
def __init__(self):
"""Initialize CUDA stream"""
self._stream = cuda_call(cudart.cudaStreamCreate())
def __enter__(self):
"""Create CUDA stream when entering context (if not already created)"""
if self._stream is None:
self._stream = cuda_call(cudart.cudaStreamCreate())
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Destroy CUDA stream when exiting context"""
self.free()
@property
def stream(self) -> cudart.cudaStream_t:
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
return self._stream
def synchronize(self):
"""Synchronize the stream"""
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
cuda_call(cudart.cudaStreamSynchronize(self._stream))
def free(self):
"""Explicitly free the CUDA stream"""
if self._stream is not None:
try:
cuda_call(cudart.cudaStreamDestroy(self._stream))
self._stream = None
except Exception as e:
logging.warning(f"Failed to destroy CUDA stream: {e}")
def __del__(self):
"""Cleanup stream on destruction"""
if hasattr(self, '_stream') and self._stream is not None:
self.free()
def __str__(self):
return f"CudaStreamContext: {self._stream}"
def __repr__(self):
return self.__str__()
def cuda_call(call):
"""Helper function to make CUDA calls and check for errors"""
def _cudaGetErrorEnum(error):
if isinstance(error, cuda.CUresult):
err, name = cuda.cuGetErrorName(error)
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
elif isinstance(error, cudart.cudaError_t):
return cudart.cudaGetErrorName(error)[1]
else:
raise RuntimeError("Unknown error type: {}".format(error))
err, res = call[0], call[1:]
if err.value:
raise RuntimeError(
"CUDA error code={}({})".format(
err.value, _cudaGetErrorEnum(err)
)
)
if len(res) == 1:
return res[0]
elif len(res) == 0:
return None
else:
return res
def getComputeCapacity(devID=0):
major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
return (major, minor)
def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream):
"""Wrapper for async host-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream))
def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream):
"""Wrapper for async device-to-host memory copy"""
cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream))
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
"""Wrapper for synchronous host-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
"""Wrapper for synchronous device-to-host memory copy"""
cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
# Initialize CUDA
cuda_call(cudart.cudaFree(0))
demo/BERT/helpers/data_processing.py
-497
View File
@@ -1,497 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import helpers.tokenization as tokenization
import collections
import numpy as np
import six
import math
import json
def convert_doc_tokens(paragraph_text):
""" Return the list of tokens from the doc text """
def is_whitespace(c):
if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
return True
return False
doc_tokens = []
prev_is_whitespace = True
for c in paragraph_text:
if is_whitespace(c):
prev_is_whitespace = True
else:
if prev_is_whitespace:
doc_tokens.append(c)
else:
doc_tokens[-1] += c
prev_is_whitespace = False
return doc_tokens
def _check_is_max_context(doc_spans, cur_span_index, position):
"""Check if this is the 'max context' doc span for the token."""
# Because of the sliding window approach taken to scoring documents, a single
# token can appear in multiple documents. E.g.
# Doc: the man went to the store and bought a gallon of milk
# Span A: the man went to the
# Span B: to the store and bought
# Span C: and bought a gallon of
# ...
#
# Now the word 'bought' will have two scores from spans B and C. We only
# want to consider the score with "maximum context", which we define as
# the *minimum* of its left and right context (the *sum* of left and
# right context will always be the same, of course).
#
# In the example the maximum context for 'bought' would be span C since
# it has 1 left context and 3 right context, while span B has 4 left context
# and 0 right context.
best_score = None
best_span_index = None
for (span_index, doc_span) in enumerate(doc_spans):
end = doc_span.start + doc_span.length - 1
if position < doc_span.start:
continue
if position > end:
continue
num_left_context = position - doc_span.start
num_right_context = end - position
score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
if best_score is None or score > best_score:
best_score = score
best_span_index = span_index
return cur_span_index == best_span_index
def convert_example_to_features(doc_tokens, question_text, tokenizer, max_seq_length,
doc_stride, max_query_length):
"""Loads a data file into a list of `InputBatch`s."""
query_tokens = tokenizer.tokenize(question_text)
if len(query_tokens) > max_query_length:
query_tokens = query_tokens[0:max_query_length]
tok_to_orig_index = []
orig_to_tok_index = []
all_doc_tokens = []
for (i, token) in enumerate(doc_tokens):
orig_to_tok_index.append(len(all_doc_tokens))
sub_tokens = tokenizer.tokenize(token)
for sub_token in sub_tokens:
tok_to_orig_index.append(i)
all_doc_tokens.append(sub_token)
# The -3 accounts for [CLS], [SEP] and [SEP]
max_tokens_for_doc = max_seq_length - len(query_tokens) - 3
# We can have documents that are longer than the maximum sequence length.
# To deal with this we do a sliding window approach, where we take chunks
# of the up to our max length with a stride of `doc_stride`.
_DocSpan = collections.namedtuple( # pylint: disable=invalid-name
"DocSpan", ["start", "length"])
doc_spans = []
start_offset = 0
while start_offset < len(all_doc_tokens):
length = len(all_doc_tokens) - start_offset
if length > max_tokens_for_doc:
length = max_tokens_for_doc
doc_spans.append(_DocSpan(start=start_offset, length=length))
if start_offset + length == len(all_doc_tokens):
break
start_offset += min(length, doc_stride)
_Feature = collections.namedtuple( # pylint: disable=invalid-name
"Feature",
["input_ids", "input_mask", "segment_ids", "tokens", "token_to_orig_map", "token_is_max_context"])
features = []
for (doc_span_index, doc_span) in enumerate(doc_spans):
tokens = []
token_to_orig_map = {}
token_is_max_context = {}
segment_ids = []
tokens.append("[CLS]")
segment_ids.append(0)
for token in query_tokens:
tokens.append(token)
segment_ids.append(0)
tokens.append("[SEP]")
segment_ids.append(0)
for i in range(doc_span.length):
split_token_index = doc_span.start + i
token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
is_max_context = _check_is_max_context(doc_spans, doc_span_index, split_token_index)
token_is_max_context[len(tokens)] = is_max_context
tokens.append(all_doc_tokens[split_token_index])
segment_ids.append(1)
tokens.append("[SEP]")
segment_ids.append(1)
input_ids = tokenizer.convert_tokens_to_ids(tokens)
# The mask has 1 for real tokens and 0 for padding tokens. Only real
# tokens are attended to.
input_mask = [1] * len(input_ids)
# Zero-pad up to the sequence length.
while len(input_ids) < max_seq_length:
input_ids.append(0)
input_mask.append(0)
segment_ids.append(0)
assert len(input_ids) == max_seq_length
assert len(input_mask) == max_seq_length
assert len(segment_ids) == max_seq_length
def create_int_feature(values):
feature = np.asarray(values, dtype=np.int32, order=None)
return feature
features.append(_Feature(
input_ids = create_int_feature(input_ids),
input_mask = create_int_feature(input_mask),
segment_ids = create_int_feature(segment_ids),
tokens = tokens,
token_to_orig_map = token_to_orig_map,
token_is_max_context = token_is_max_context
))
return features
def read_squad_json(input_file):
"""read from squad json into a list of examples"""
with open(input_file, "r", encoding='utf-8') as reader:
input_data = json.load(reader)["data"]
_Example = collections.namedtuple( # pylint: disable=invalid-name
"Example",
["id", "question_text", "doc_tokens"])
examples = []
for entry in input_data:
for paragraph in entry["paragraphs"]:
paragraph_text = paragraph["context"]
doc_tokens = convert_doc_tokens(paragraph_text)
for qa in paragraph["qas"]:
examples.append(_Example(
id = qa["id"],
question_text = qa["question"],
doc_tokens = doc_tokens
))
return examples
def _get_best_indexes(logits, n_best_size):
"""Get the n-best logits from a list."""
index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)
best_indexes = []
for i in range(len(index_and_score)):
if i >= n_best_size:
break
best_indexes.append(index_and_score[i][0])
return best_indexes
def get_final_text(pred_text, orig_text, do_lower_case):
"""Project the tokenized prediction back to the original text."""
# When we created the data, we kept track of the alignment between original
# (whitespace tokenized) tokens and our WordPiece tokenized tokens. So
# now `orig_text` contains the span of our original text corresponding to the
# span that we predicted.
#
# However, `orig_text` may contain extra characters that we don't want in
# our prediction.
#
# For example, let's say:
# pred_text = steve smith
# orig_text = Steve Smith's
#
# We don't want to return `orig_text` because it contains the extra "'s".
#
# We don't want to return `pred_text` because it's already been normalized
# (the SQuAD eval script also does punctuation stripping/lower casing but
# our tokenizer does additional normalization like stripping accent
# characters).
#
# What we really want to return is "Steve Smith".
#
# Therefore, we have to apply a semi-complicated alignment heruistic between
# `pred_text` and `orig_text` to get a character-to-charcter alignment. This
# can fail in certain cases in which case we just return `orig_text`.
def _strip_spaces(text):
ns_chars = []
ns_to_s_map = collections.OrderedDict()
for (i, c) in enumerate(text):
if c == " ":
continue
ns_to_s_map[len(ns_chars)] = i
ns_chars.append(c)
ns_text = "".join(ns_chars)
return (ns_text, ns_to_s_map)
# We first tokenize `orig_text`, strip whitespace from the result
# and `pred_text`, and check if they are the same length. If they are
# NOT the same length, the heuristic has failed. If they are the same
# length, we assume the characters are one-to-one aligned.
tokenizer = tokenization.BasicTokenizer(do_lower_case=do_lower_case)
tok_text = " ".join(tokenizer.tokenize(orig_text))
start_position = tok_text.find(pred_text)
if start_position == -1:
return orig_text
end_position = start_position + len(pred_text) - 1
(orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text)
(tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text)
if len(orig_ns_text) != len(tok_ns_text):
return orig_text
# We then project the characters in `pred_text` back to `orig_text` using
# the character-to-character alignment.
tok_s_to_ns_map = {}
for (i, tok_index) in six.iteritems(tok_ns_to_s_map):
tok_s_to_ns_map[tok_index] = i
orig_start_position = None
if start_position in tok_s_to_ns_map:
ns_start_position = tok_s_to_ns_map[start_position]
if ns_start_position in orig_ns_to_s_map:
orig_start_position = orig_ns_to_s_map[ns_start_position]
if orig_start_position is None:
return orig_text
orig_end_position = None
if end_position in tok_s_to_ns_map:
ns_end_position = tok_s_to_ns_map[end_position]
if ns_end_position in orig_ns_to_s_map:
orig_end_position = orig_ns_to_s_map[ns_end_position]
if orig_end_position is None:
return orig_text
output_text = orig_text[orig_start_position:(orig_end_position + 1)]
return output_text
def _compute_softmax(scores):
"""Compute softmax probability over raw logits."""
if not scores:
return []
max_score = None
for score in scores:
if max_score is None or score > max_score:
max_score = score
exp_scores = []
total_sum = 0.0
for score in scores:
x = math.exp(score - max_score)
exp_scores.append(x)
total_sum += x
probs = []
for score in exp_scores:
probs.append(score / total_sum)
return probs
def get_predictions(doc_tokens, features, results, n_best_size, max_answer_length):
_PrelimPrediction = collections.namedtuple( # pylint: disable=invalid-name
"PrelimPrediction",
["feature_index", "start_index", "end_index", "start_logit", "end_logit"])
prediction = ""
scores_diff_json = 0.0
prelim_predictions = []
# keep track of the minimum score of null start+end of position 0
score_null = 1000000 # large and positive
min_null_feature_index = 0 # the paragraph slice with min mull score
null_start_logit = 0 # the start logit at the slice with min null score
null_end_logit = 0 # the end logit at the slice with min null score
version_2_with_negative = False
for result in results:
start_indexes = _get_best_indexes(result.start_logits, n_best_size)
end_indexes = _get_best_indexes(result.end_logits, n_best_size)
feature = features[result.feature_index]
# if we could have irrelevant answers, get the min score of irrelevant
if version_2_with_negative:
feature_null_score = result.start_logits[0] + result.end_logits[0]
if feature_null_score < score_null:
score_null = feature_null_score
min_null_feature_index = 0
null_start_logit = result.start_logits[0]
null_end_logit = result.end_logits[0]
for start_index in start_indexes:
for end_index in end_indexes:
# We could hypothetically create invalid predictions, e.g., predict
# that the start of the span is in the question. We throw out all
# invalid predictions.
if start_index >= len(feature.tokens):
continue
if end_index >= len(feature.tokens):
continue
if start_index not in feature.token_to_orig_map:
continue
if end_index not in feature.token_to_orig_map:
continue
if not feature.token_is_max_context.get(start_index, False):
continue
if end_index < start_index:
continue
length = end_index - start_index + 1
if length > max_answer_length:
continue
prelim_predictions.append(
_PrelimPrediction(
feature_index=result.feature_index,
start_index=start_index,
end_index=end_index,
start_logit=result.start_logits[start_index],
end_logit=result.end_logits[end_index]))
if version_2_with_negative:
prelim_predictions.append(
_PrelimPrediction(
feature_index=result.feature_index,
start_index=0,
end_index=0,
start_logit=null_start_logit,
end_logit=null_end_logit))
prelim_predictions = sorted(
prelim_predictions,
key=lambda x: (x.start_logit + x.end_logit),
reverse=True)
_NbestPrediction = collections.namedtuple( # pylint: disable=invalid-name
"NbestPrediction", ["text", "start_logit", "end_logit"])
seen_predictions = {}
nbest = []
for pred in prelim_predictions:
if len(nbest) >= n_best_size:
break
if pred.start_index > 0: # this is a non-null prediction
feature = features[pred.feature_index]
tok_tokens = feature.tokens[pred.start_index:(pred.end_index + 1)]
orig_doc_start = feature.token_to_orig_map[pred.start_index]
orig_doc_end = feature.token_to_orig_map[pred.end_index]
orig_tokens = doc_tokens[orig_doc_start:(orig_doc_end + 1)]
tok_text = " ".join(tok_tokens)
# De-tokenize WordPieces that have been split off.
tok_text = tok_text.replace(" ##", "")
tok_text = tok_text.replace("##", "")
# Clean whitespace
tok_text = tok_text.strip()
tok_text = " ".join(tok_text.split())
orig_text = " ".join(orig_tokens)
final_text = get_final_text(tok_text, orig_text, True)
if final_text in seen_predictions:
continue
seen_predictions[final_text] = True
else:
final_text = ""
seen_predictions[final_text] = True
if len(final_text):
nbest.append(
_NbestPrediction(
text=final_text,
start_logit=pred.start_logit,
end_logit=pred.end_logit))
# if we didn't inlude the empty option in the n-best, inlcude it
if version_2_with_negative:
if "" not in seen_predictions:
nbest.append(
_NbestPrediction(
text="", start_logit=null_start_logit,
end_logit=null_end_logit))
# In very rare edge cases we could have no valid predictions. So we
# just create a nonce prediction in this case to avoid failure.
if not nbest:
nbest.append(
_NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0))
assert len(nbest) >= 1
total_scores = []
best_non_null_entry = None
for entry in nbest:
total_scores.append(entry.start_logit + entry.end_logit)
if not best_non_null_entry:
if entry.text:
best_non_null_entry = entry
probs = _compute_softmax(total_scores)
nbest_json = []
for (i, entry) in enumerate(nbest):
output = collections.OrderedDict()
output["text"] = entry.text
output["probability"] = probs[i]
output["start_logit"] = entry.start_logit
output["end_logit"] = entry.end_logit
nbest_json.append(output)
assert len(nbest_json) >= 1
null_score_diff_threshold = 0.0
if not version_2_with_negative:
prediction = nbest_json[0]["text"]
else:
# predict "" iff the null score - the score of best non-null > threshold
score_diff = score_null - best_non_null_entry.start_logit - (
best_non_null_entry.end_logit)
scores_diff_json = score_diff
if score_diff > null_score_diff_threshold:
prediction = ""
else:
prediction = best_non_null_entry.text
return prediction, nbest_json, scores_diff_json
demo/BERT/helpers/tokenization.py
-446
View File
@@ -1,446 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tokenization classes."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import re
import unicodedata
import six
def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
"""Checks whether the casing config is consistent with the checkpoint name."""
# The casing has to be passed in by the user and there is no explicit check
# as to whether it matches the checkpoint. The casing information probably
# should have been stored in the bert_config.json file, but it's not, so
# we have to heuristically detect it to validate.
if not init_checkpoint:
return
m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
if m is None:
return
model_name = m.group(1)
lower_models = [
"uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
"multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
]
cased_models = [
"cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
"multi_cased_L-12_H-768_A-12"
]
is_bad_config = False
if model_name in lower_models and not do_lower_case:
is_bad_config = True
actual_flag = "False"
case_name = "lowercased"
opposite_flag = "True"
if model_name in cased_models and do_lower_case:
is_bad_config = True
actual_flag = "True"
case_name = "cased"
opposite_flag = "False"
if is_bad_config:
raise ValueError(
"You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
"However, `%s` seems to be a %s model, so you "
"should pass in `--do_lower_case=%s` so that the fine-tuning matches "
"how the model was pre-training. If this error is wrong, please "
"just comment out this check." % (actual_flag, init_checkpoint,
model_name, case_name, opposite_flag))
def convert_to_unicode(text):
"""Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text.decode("utf-8", "ignore")
elif isinstance(text, unicode):
return text
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def printable_text(text):
"""Returns text encoded in a way suitable for print or `tf.logging`."""
# These functions want `str` for both Python2 and Python3, but in one case
# it's a Unicode string and in the other it's a byte string.
if six.PY3:
if isinstance(text, str):
return text
elif isinstance(text, bytes):
return text.decode("utf-8", "ignore")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
elif six.PY2:
if isinstance(text, str):
return text
elif isinstance(text, unicode):
return text.encode("utf-8")
else:
raise ValueError("Unsupported string type: %s" % (type(text)))
else:
raise ValueError("Not running on Python2 or Python 3?")
def load_vocab(vocab_file):
"""Loads a vocabulary file into a dictionary."""
vocab = collections.OrderedDict()
index = 0
with open(vocab_file, "r", encoding='utf-8') as reader:
while True:
token = convert_to_unicode(reader.readline())
if not token:
break
token = token.strip()
vocab[token] = index
index += 1
return vocab
def convert_by_vocab(vocab, items):
"""Converts a sequence of [tokens|ids] using the vocab."""
output = []
for item in items:
output.append(vocab[item])
return output
def convert_tokens_to_ids(vocab, tokens):
return convert_by_vocab(vocab, tokens)
def convert_ids_to_tokens(inv_vocab, ids):
return convert_by_vocab(inv_vocab, ids)
def whitespace_tokenize(text):
"""Runs basic whitespace cleaning and splitting on a piece of text."""
text = text.strip()
if not text:
return []
tokens = text.split()
return tokens
class FullTokenizer(object):
"""Runs end-to-end tokenziation."""
def __init__(self, vocab_file, do_lower_case=True):
self.vocab = load_vocab(vocab_file)
self.inv_vocab = {v: k for k, v in self.vocab.items()}
self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
def tokenize(self, text):
split_tokens = []
for token in self.basic_tokenizer.tokenize(text):
for sub_token in self.wordpiece_tokenizer.tokenize(token):
split_tokens.append(sub_token)
return split_tokens
def convert_tokens_to_ids(self, tokens):
return convert_by_vocab(self.vocab, tokens)
def convert_ids_to_tokens(self, ids):
return convert_by_vocab(self.inv_vocab, ids)
class BertTokenizer(object):
"""Runs end-to-end tokenization: punctuation splitting + wordpiece"""
def __init__(self, vocab_file, do_lower_case=True):
self.vocab = load_vocab(vocab_file)
self.ids_to_tokens = collections.OrderedDict(
[(ids, tok) for tok, ids in self.vocab.items()])
self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
def tokenize(self, text):
split_tokens = []
for token in self.basic_tokenizer.tokenize(text):
for sub_token in self.wordpiece_tokenizer.tokenize(token):
split_tokens.append(sub_token)
return split_tokens
def convert_tokens_to_ids(self, tokens):
"""Converts a sequence of tokens into ids using the vocab."""
ids = []
for token in tokens:
ids.append(self.vocab[token])
return ids
def convert_ids_to_tokens(self, ids):
"""Converts a sequence of ids in wordpiece tokens using the vocab."""
tokens = []
for i in ids:
tokens.append(self.ids_to_tokens[i])
return tokens
class BasicTokenizer(object):
"""Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
def __init__(self, do_lower_case=True):
"""Constructs a BasicTokenizer.
Args:
do_lower_case: Whether to lower case the input.
"""
self.do_lower_case = do_lower_case
def tokenize(self, text):
"""Tokenizes a piece of text."""
text = convert_to_unicode(text)
text = self._clean_text(text)
# This was added on November 1st, 2018 for the multilingual and Chinese
# models. This is also applied to the English models now, but it doesn't
# matter since the English models were not trained on any Chinese data
# and generally don't have any Chinese data in them (there are Chinese
# characters in the vocabulary because Wikipedia does have some Chinese
# words in the English Wikipedia.).
text = self._tokenize_chinese_chars(text)
orig_tokens = whitespace_tokenize(text)
split_tokens = []
for token in orig_tokens:
if self.do_lower_case:
token = token.lower()
token = self._run_strip_accents(token)
split_tokens.extend(self._run_split_on_punc(token))
output_tokens = whitespace_tokenize(" ".join(split_tokens))
return output_tokens
def _run_strip_accents(self, text):
"""Strips accents from a piece of text."""
text = unicodedata.normalize("NFD", text)
output = []
for char in text:
cat = unicodedata.category(char)
if cat == "Mn":
continue
output.append(char)
return "".join(output)
def _run_split_on_punc(self, text):
"""Splits punctuation on a piece of text."""
chars = list(text)
i = 0
start_new_word = True
output = []
while i < len(chars):
char = chars[i]
if _is_punctuation(char):
output.append([char])
start_new_word = True
else:
if start_new_word:
output.append([])
start_new_word = False
output[-1].append(char)
i += 1
return ["".join(x) for x in output]
def _tokenize_chinese_chars(self, text):
"""Adds whitespace around any CJK character."""
output = []
for char in text:
cp = ord(char)
if self._is_chinese_char(cp):
output.append(" ")
output.append(char)
output.append(" ")
else:
output.append(char)
return "".join(output)
def _is_chinese_char(self, cp):
"""Checks whether CP is the codepoint of a CJK character."""
# This defines a "chinese character" as anything in the CJK Unicode block:
# https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
#
# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
# despite its name. The modern Korean Hangul alphabet is a different block,
# as is Japanese Hiragana and Katakana. Those alphabets are used to write
# space-separated words, so they are not treated specially and handled
# like the all of the other languages.
if ((cp >= 0x4E00 and cp <= 0x9FFF) or #
(cp >= 0x3400 and cp <= 0x4DBF) or #
(cp >= 0x20000 and cp <= 0x2A6DF) or #
(cp >= 0x2A700 and cp <= 0x2B73F) or #
(cp >= 0x2B740 and cp <= 0x2B81F) or #
(cp >= 0x2B820 and cp <= 0x2CEAF) or
(cp >= 0xF900 and cp <= 0xFAFF) or #
(cp >= 0x2F800 and cp <= 0x2FA1F)): #
return True
return False
def _clean_text(self, text):
"""Performs invalid character removal and whitespace cleanup on text."""
output = []
for char in text:
cp = ord(char)
if cp == 0 or cp == 0xfffd or _is_control(char):
continue
if _is_whitespace(char):
output.append(" ")
else:
output.append(char)
return "".join(output)
class WordpieceTokenizer(object):
"""Runs WordPiece tokenziation."""
def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
self.vocab = vocab
self.unk_token = unk_token
self.max_input_chars_per_word = max_input_chars_per_word
def tokenize(self, text):
"""Tokenizes a piece of text into its word pieces.
This uses a greedy longest-match-first algorithm to perform tokenization
using the given vocabulary.
For example:
input = "unaffable"
output = ["un", "##aff", "##able"]
Args:
text: A single token or whitespace separated tokens. This should have
already been passed through `BasicTokenizer.
Returns:
A list of wordpiece tokens.
"""
text = convert_to_unicode(text)
output_tokens = []
for token in whitespace_tokenize(text):
chars = list(token)
if len(chars) > self.max_input_chars_per_word:
output_tokens.append(self.unk_token)
continue
is_bad = False
start = 0
sub_tokens = []
while start < len(chars):
end = len(chars)
cur_substr = None
while start < end:
substr = "".join(chars[start:end])
if start > 0:
substr = "##" + substr
if substr in self.vocab:
cur_substr = substr
break
end -= 1
if cur_substr is None:
is_bad = True
break
sub_tokens.append(cur_substr)
start = end
if is_bad:
output_tokens.append(self.unk_token)
else:
output_tokens.extend(sub_tokens)
return output_tokens
def _is_whitespace(char):
"""Checks whether `chars` is a whitespace character."""
# \t, \n, and \r are technically contorl characters but we treat them
# as whitespace since they are generally considered as such.
if char == " " or char == "\t" or char == "\n" or char == "\r":
return True
cat = unicodedata.category(char)
if cat == "Zs":
return True
return False
def _is_control(char):
"""Checks whether `chars` is a control character."""
# These are technically control characters but we count them as whitespace
# characters.
if char == "\t" or char == "\n" or char == "\r":
return False
cat = unicodedata.category(char)
if cat.startswith("C"):
return True
return False
def _is_punctuation(char):
"""Checks whether `chars` is a punctuation character."""
cp = ord(char)
# We treat all non-letter/number ASCII as punctuation.
# Characters such as "^", "$", and "`" are not in the Unicode
# Punctuation class but we treat them as punctuation anyways, for
# consistency.
if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
(cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
return True
cat = unicodedata.category(char)
if cat.startswith("P"):
return True
return False
demo/BERT/infer_c/bert_infer.h
-362
View File
@@ -1,362 +0,0 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef INFER_C_BERT_INFER_H
#define INFER_C_BERT_INFER_H
#include "common.h"
#include "logging.h"
#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <algorithm>
#include <cuda_runtime.h>
#include <fstream>
#include <numeric>
#include <string.h>
#include <vector>
using namespace nvinfer1;
struct BertInference
{
BertInference(
const std::string& enginePath, const int maxBatchSize, const int seqLength, const bool enableGraph = false)
: mSeqLength(seqLength)
, mEnableGraph(enableGraph)
{
gLogInfo << "--------------------\n";
gLogInfo << "Using BERT inference C++\n";
if (enableGraph)
{
gLogInfo << "CUDA Graph is enabled\n";
}
else
{
gLogInfo << "CUDA Graph is disabled\n";
}
gLogInfo << "--------------------\n";
initLibNvInferPlugins(&gLogger, "");
gLogInfo << "Loading BERT Inference Engine ... \n";
std::ifstream input(enginePath, std::ios::binary);
if (!input)
{
gLogError << "Error opening engine file: " << enginePath << "\n";
exit(-1);
}
input.seekg(0, input.end);
const size_t fsize = input.tellg();
input.seekg(0, input.beg);
std::vector<char> bytes(fsize);
input.read(bytes.data(), fsize);
mRuntime = TrtUniquePtr<IRuntime>(createInferRuntime(gLogger));
if (mRuntime == nullptr)
{
gLogError << "Error creating TRT mRuntime\n";
exit(-1);
}
mEngine = TrtUniquePtr<ICudaEngine>(mRuntime->deserializeCudaEngine(bytes.data(), bytes.size()));
if (mEngine == nullptr)
{
gLogError << "Error deserializing CUDA engine\n";
exit(-1);
}
gLogInfo << "Done\n";
mEnableVariableLen = mEngine->getNbIOTensors() == kBERT_INPUT_NUM + 1 ? false : true;
if (mEnableVariableLen)
{
gLogInfo << "Variable length is enabled\n";
}
else
{
gLogInfo << "Variable length is disabled\n";
}
mContext = TrtUniquePtr<IExecutionContext>(mEngine->createExecutionContext());
if (!mContext)
{
gLogError << "Error creating execution context\n";
exit(-1);
}
gpuErrChk(cudaStreamCreate(&mStream));
allocateBindings(maxBatchSize);
}
void allocateBindings(const int maxBatchSize)
{
const size_t allocationSize = mSeqLength * maxBatchSize * sizeof(int32_t);
// Static sizes with implicit batch size: allocation sizes known to engine
if (mEnableVariableLen)
{
const size_t allocationSizes[] = {allocationSize, allocationSize,
sizeof(int32_t) * (maxBatchSize + 1),
sizeof(int32_t) * (mSeqLength)};
for (int i = 0; i < sizeof(allocationSizes) / sizeof(allocationSizes[0]); i++)
{
void* devBuf;
gpuErrChk(cudaMalloc(&devBuf, allocationSizes[i]));
gpuErrChk(cudaMemset(devBuf, 0, allocationSizes[i]));
mDeviceBuffers.emplace_back(devBuf);
mInputSizes.emplace_back(allocationSizes[i]);
}
}
else
{
for (int i = 0; i < kBERT_INPUT_NUM; i++)
{
void* devBuf;
gpuErrChk(cudaMalloc(&devBuf, allocationSize));
gpuErrChk(cudaMemset(devBuf, 0, allocationSize));
mDeviceBuffers.emplace_back(devBuf);
mInputSizes.emplace_back(allocationSize);
}
}
const size_t numOutputItems = maxBatchSize * mSeqLength * 2;
mOutputSize = numOutputItems * sizeof(float);
if (mEnableVariableLen)
{
mOutputDims = {maxBatchSize * mSeqLength * 2};
}
else
{
mOutputDims = {maxBatchSize, mSeqLength, 2, 1, 1};
}
void* devBuf;
gpuErrChk(cudaMalloc(&devBuf, mOutputSize));
gpuErrChk(cudaMemset(devBuf, 0, mOutputSize));
mDeviceBuffers.emplace_back(devBuf);
mHostOutput.resize(numOutputItems);
}
void prepare(int profIdx, int batchSize)
{
mContext->setOptimizationProfileAsync(profIdx, mStream);
if (mEnableVariableLen)
{
const int allocationSizes[] = {mSeqLength * batchSize, mSeqLength * batchSize, batchSize + 1, mSeqLength};
for (int i = 0; i < sizeof(allocationSizes)/sizeof(allocationSizes[0]); i++)
{
auto const tensorName = mEngine->getIOTensorName(i % mEngine->getNbIOTensors());
mContext->setInputShape(tensorName, Dims{1, {allocationSizes[i]}});
}
}
else
{
for (int i = 0; i < kBERT_INPUT_NUM; i++)
{
auto const tensorName = mEngine->getIOTensorName(i);
mContext->setInputShape(tensorName, Dims2(batchSize, mSeqLength));
}
}
if (!mContext->allInputDimensionsSpecified())
{
gLogError << "Not all input dimensions are specified for the exeuction context\n";
exit(-1);
}
if (mEnableGraph)
{
for (int32_t i = 0; i < mEngine->getNbIOTensors(); i++)
{
auto const& name = mEngine->getIOTensorName(i);
mContext->setTensorAddress(name, mDeviceBuffers[i]);
}
cudaGraph_t graph;
cudaGraphExec_t exec;
// warm up and let mContext do cublas initialization
bool status = mContext->enqueueV3(mStream);
if (!status)
{
gLogError << "Enqueue failed\n";
exit(-1);
}
gLogVerbose << "Capturing graph\n";
gpuErrChk(cudaStreamBeginCapture(mStream, cudaStreamCaptureModeRelaxed));
status = mContext->enqueueV3(mStream);
if (!status)
{
gLogError << "Enqueue failed\n";
exit(-1);
}
gpuErrChk(cudaStreamEndCapture(mStream, &graph));
gpuErrChk(cudaStreamSynchronize(mStream));
gpuErrChk(cudaGraphInstantiate(&exec, graph, NULL, NULL, 0));
mExecGraph = exec;
}
mCuSeqlens.resize(batchSize + 1);
std::generate(mCuSeqlens.begin(), mCuSeqlens.end(), [pos = -mSeqLength, this]() mutable{ pos += mSeqLength; return pos; });
}
void run(const void* const* inputBuffers, int warmUps, int iterations)
{
for (int i = 0; i < kBERT_INPUT_NUM; i++)
{
gpuErrChk(
cudaMemcpyAsync(mDeviceBuffers[i], inputBuffers[i], mInputSizes[i], cudaMemcpyHostToDevice, mStream));
}
gLogInfo << "Warming up " << warmUps << " iterations ...\n";
for (int it = 0; it < warmUps; it++)
{
if (mEnableGraph)
{
gpuErrChk(cudaGraphLaunch(mExecGraph, mStream));
}
else
{
bool status = mContext->enqueueV3(mStream);
if (!status)
{
gLogError << "Enqueue failed\n";
exit(-1);
}
}
}
gpuErrChk(cudaStreamSynchronize(mStream));
cudaEvent_t start, stop;
gpuErrChk(cudaEventCreate(&start));
gpuErrChk(cudaEventCreate(&stop));
std::vector<float> times;
gLogInfo << "Running " << iterations << " iterations ...\n";
for (int it = 0; it < iterations; it++)
{
gpuErrChk(cudaEventRecord(start, mStream));
if (mEnableGraph)
{
gpuErrChk(cudaGraphLaunch(mExecGraph, mStream));
}
else
{
bool status = mContext->enqueueV3(mStream);
if (!status)
{
gLogError << "Enqueue failed\n";
exit(-1);
}
}
gpuErrChk(cudaEventRecord(stop, mStream));
gpuErrChk(cudaStreamSynchronize(mStream));
float time;
gpuErrChk(cudaEventElapsedTime(&time, start, stop));
times.push_back(time);
}
gpuErrChk(cudaMemcpyAsync(
mHostOutput.data(), mDeviceBuffers[mEnableVariableLen ? kBERT_INPUT_NUM + 1 : kBERT_INPUT_NUM], mOutputSize, cudaMemcpyDeviceToHost, mStream));
gpuErrChk(cudaStreamSynchronize(mStream));
mTimes.push_back(times);
}
void run(const void* inputIds, const void* segmentIds, const void* inputMask, int warmUps, int iterations)
{
if (mEnableVariableLen)
{
const std::vector<const void*> inputBuffers = {inputIds, segmentIds, mCuSeqlens.data()};
run(inputBuffers.data(), warmUps, iterations);
}
else
{
const std::vector<const void*> inputBuffers = {inputIds, segmentIds, inputMask};
run(inputBuffers.data(), warmUps, iterations);
}
}
void run(int profIdx, int batchSize, const void* inputIds, const void* segmentIds, const void* inputMask,
int warmUps, int iterations)
{
prepare(profIdx, batchSize);
run(inputIds, segmentIds, inputMask, warmUps, iterations);
}
void reportTiming(int batchIndex, int batchSize)
{
std::vector<float>& times = mTimes[batchIndex];
const float totalTime = std::accumulate(times.begin(), times.end(), 0.0);
const float avgTime = totalTime / times.size();
sort(times.begin(), times.end());
const float percentile95 = times[(int) ((float) times.size() * 0.95)];
const float percentile99 = times[(int) ((float) times.size() * 0.99)];
const int throughput = (int) ((float) batchSize * (1000.0 / avgTime));
gLogInfo << "Running " << times.size() << " iterations with Batch Size: " << batchSize << "\n";
gLogInfo << "\tTotal Time: " << totalTime << " ms \n";
gLogInfo << "\tAverage Time: " << avgTime << " ms\n";
gLogInfo << "\t95th Percentile Time: " << percentile95 << " ms\n";
gLogInfo << "\t99th Percentile Time: " << percentile99 << " ms\n";
gLogInfo << "\tThroughput: " << throughput << " sentences/s\n";
}
~BertInference()
{
gpuErrChk(cudaStreamDestroy(mStream));
for (auto& buf : mDeviceBuffers)
{
gpuErrChk(cudaFree(buf));
}
}
static const int kBERT_INPUT_NUM = 3;
const int mSeqLength;
const bool mEnableGraph;
TrtUniquePtr<IRuntime> mRuntime{nullptr};
TrtUniquePtr<ICudaEngine> mEngine{nullptr};
TrtUniquePtr<IExecutionContext> mContext{nullptr};
bool mEnableVariableLen;
std::vector<int> mCuSeqlens;
cudaStream_t mStream{NULL};
std::vector<void*> mDeviceBuffers;
std::vector<float> mHostOutput;
std::vector<size_t> mInputSizes;
size_t mOutputSize;
std::vector<int> mOutputDims;
std::vector<std::vector<float>> mTimes;
cudaGraphExec_t mExecGraph;
};
#endif // INFER_C_BERT_INFER_H
demo/BERT/infer_c/common.h
-179
View File
@@ -1,179 +0,0 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef INFER_C_COMMON_H
#define INFER_C_COMMON_H
#include "logging.h"
#include <cuda_runtime_api.h>
#include <getopt.h>
#include <memory>
#include <vector>
struct Args
{
bool help{false};
std::string engine{};
std::vector<int> batchSize;
int sequenceLength{128};
int iterations{200};
int warmUpRuns{10};
int randomSeed{12345};
bool enableGraph{false};
};
//!
//! \brief Populates the Args struct with the provided command-line parameters.
//!
//! \throw invalid_argument if any of the arguments are not valid
//!
//! \return boolean If return value is true, execution can continue, otherwise program should exit
//!
inline bool parseArgs(Args& args, int argc, char* argv[])
{
while (1)
{
int arg;
// clang-format off
static struct option long_options[] =
{
{"help", no_argument, 0, 'h'},
{"engine", required_argument, 0, 'e'},
{"batch_size", required_argument, 0, 'b'},
{"sequence_length", no_argument, 0, 's'},
{"iterations", required_argument, 0, 'i'},
{"warm_up_runs", required_argument, 0, 'w'},
{"ramdon_seed", required_argument, 0, 'r'},
{"enable_graph", no_argument, 0, 'g'},
{nullptr, 0, nullptr, 0}
};
// clang-format on
int option_index = 0;
arg = getopt_long(argc, argv, "he:b:s:i:w:r:g", long_options, &option_index);
if (arg == -1)
{
break;
}
switch (arg)
{
case 'h': args.help = true; return false;
case 'e':
if (optarg)
{
args.engine = optarg;
}
else
{
std::cerr << "ERROR: --engine requires option argument" << std::endl;
return false;
}
break;
case 'b':
if (optarg)
{
args.batchSize.push_back(std::stoi(optarg));
}
else
{
std::cerr << "ERROR: --batch_size requires option argument" << std::endl;
return false;
}
break;
case 's':
if (optarg)
{
args.sequenceLength = std::stoi(optarg);
}
else
{
std::cerr << "ERROR: --sequence_length requires option argument" << std::endl;
return false;
}
break;
case 'i':
if (optarg)
{
args.iterations = std::stoi(optarg);
}
else
{
std::cerr << "ERROR: --iterations requires option argument" << std::endl;
return false;
}
break;
case 'w':
if (optarg)
{
args.warmUpRuns = std::stoi(optarg);
}
else
{
std::cerr << "ERROR: --warm_up_runs requires option argument" << std::endl;
return false;
}
break;
case 'r':
if (optarg)
{
args.randomSeed = std::stoi(optarg);
}
else
{
std::cerr << "ERROR: --random_seed requires option argument" << std::endl;
return false;
}
break;
case 'g': args.enableGraph = true; break;
default: return false;
}
}
return true;
}
// clang-format off
#define gpuErrChk(ans) \
{ \
gpuAssert((ans), __FILE__, __LINE__); \
}
// clang-format on
inline void gpuAssert(cudaError_t code, const char* file, int line, bool abort = true)
{
if (code != cudaSuccess)
{
gLogError << "GPUassert: " << cudaGetErrorString(code) << " " << file << " " << line << "\n";
if (abort)
{
exit(code);
}
}
}
template <typename T>
struct TrtDestroyer
{
void operator()(T* t)
{
delete t;
}
};
template <typename T>
using TrtUniquePtr = std::unique_ptr<T, TrtDestroyer<T>>;
#endif // INFER_C_COMMON_H
demo/BERT/infer_c/infer_c.cpp
-62
View File
@@ -1,62 +0,0 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "bert_infer.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
namespace py = pybind11;
struct BertInferenceRunner
{
BertInferenceRunner(
const std::string& enginePath, const int maxBatchSize, const int maxSeqLength, const bool enableGraph)
: bert{enginePath, maxBatchSize, maxSeqLength, enableGraph}
{
}
void prepare(const int batchSize)
{
bert.prepare(0, batchSize);
}
py::array_t<float> run(py::array_t<int> inputIds, py::array_t<int> segmentIds, py::array_t<int> inputMask)
{
const void* inputIdsPtr = inputIds.request().ptr;
const void* segmentIdsPtr = segmentIds.request().ptr;
const void* inputMaskPtr = inputMask.request().ptr;
bert.run(inputIdsPtr, segmentIdsPtr, inputMaskPtr, 0, 1);
auto output = py::array_t<float>(bert.mOutputDims, (float*) bert.mHostOutput.data());
return output;
}
BertInference bert;
};
PYBIND11_MODULE(infer_c, m)
{
m.doc() = "Pybind11 plugin for Bert inference";
py::class_<BertInferenceRunner>(m, "bert_inf")
.def(py::init<const std::string&, const int, const int, const bool>())
.def("prepare", &BertInferenceRunner::prepare)
.def("run", &BertInferenceRunner::run);
}
demo/BERT/infer_c/logging.h
-86
View File
@@ -1,86 +0,0 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef INFER_C_LOGGING_H
#define INFER_C_LOGGING_H
#include <NvInfer.h>
#include <cassert>
#include <iostream>
#include <map>
using namespace nvinfer1;
using Severity = nvinfer1::ILogger::Severity;
class Logger : public ILogger
{
public:
Logger(Severity severity)
: mOstream(&std::cout)
, mReportableSeverity(severity)
{
}
template <typename T>
Logger& operator<<(T const& obj)
{
if (mOstream != nullptr)
{
*mOstream << obj;
}
return *this;
}
Logger& report(Severity severity, const char* msg)
{
if (severity <= mReportableSeverity)
{
const std::map<Severity, std::string> prefixMapping = {{Severity::kINTERNAL_ERROR, "[DemoBERT][F] "},
{Severity::kERROR, "[DemoBERT][E] "}, {Severity::kWARNING, "[DemoBERT][W] "},
{Severity::kINFO, "[DemoBERT][I] "}, {Severity::kVERBOSE, "[DemoBERT][V] "}};
assert(prefixMapping.find(severity) != prefixMapping.end());
mOstream = &std::cout;
*this << prefixMapping.at(severity) << msg;
return *this;
}
mOstream = nullptr;
return *this;
}
private:
void log(Severity severity, const char* msg) noexcept override
{
report(severity, msg) << "\n";
}
std::ostream* mOstream;
Severity mReportableSeverity;
};
extern Logger gLogger;
#define gLogFatal gLogger.report(Severity::kINTERNAL_ERROR, "")
#define gLogError gLogger.report(Severity::kERROR, "")
#define gLogWarning gLogger.report(Severity::kWARNING, "")
#define gLogInfo gLogger.report(Severity::kINFO, "")
#define gLogVerbose gLogger.report(Severity::kVERBOSE, "")
#endif // INFER_C_LOGGING_H
demo/BERT/infer_c/perf.cpp
-135
View File
@@ -1,135 +0,0 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "bert_infer.h"
#include "common.h"
#include <limits>
#include <random>
using namespace nvinfer1;
void printHelpInfo()
{
std::cout << "usage: ./perf [-h] [-e ENGINE] [-b BATCH_SIZE] [-s SEQUENCE_LENGTH]\n";
std::cout << " [-i ITERATIONS] [-w WARM_UP_RUNS] [-r RANDOM_SEED] [--enable_graph]\n";
std::cout << "\n";
std::cout << "BERT Inference Benchmark\n";
std::cout << "\n";
std::cout << "optional arguments:\n";
std::cout << " -h, --help show this help message and exit\n";
std::cout << " -e ENGINE, --engine ENGINE\n";
std::cout << " Path to BERT TensorRT engine\n";
std::cout << " -b BATCH_SIZE, --batch_size BATCH_SIZE\n";
std::cout << " Batch size(s) to benchmark. Can be specified multiple\n";
std::cout << " times for more than one batch size. This script\n";
std::cout << " assumes that the engine has been built with one\n";
std::cout << " optimization profile for each batch size, and that\n";
std::cout << " these profiles are in order of increasing batch size.\n";
std::cout << " -s SEQUENCE_LENGTH, --sequence_length SEQUENCE_LENGTH\n";
std::cout << " Sequence length of the BERT model\n";
std::cout << " -i ITERATIONS, --iterations ITERATIONS\n";
std::cout << " Number of iterations to run when benchmarking.\n";
std::cout << " -w WARM_UP_RUNS, --warm_up_runs WARM_UP_RUNS\n";
std::cout << " Number of iterations to run prior to benchmarking.\n";
std::cout << " -r RANDOM_SEED, --random_seed RANDOM_SEED\n";
std::cout << " Random seed.\n";
std::cout << " --enable_graph\n";
std::cout << " Enable CUDA Graph.\n";
std::cout << std::endl;
}
void printDeviceInfo()
{
int32_t device{};
gpuErrChk(cudaGetDevice(&device));
cudaDeviceProp properties{};
gpuErrChk(cudaGetDeviceProperties(&properties, device));
std::cout << "=== Device Information ===" << std::endl;
std::cout << "Selected Device: " << properties.name << std::endl;
std::cout << "Compute Capability: " << properties.major << "." << properties.minor << std::endl;
std::cout << "SMs: " << properties.multiProcessorCount << std::endl;
std::cout << "Compute Clock Rate: " << properties.clockRate / 1000000.0F << " GHz" << std::endl;
std::cout << "Device Global Memory: " << (properties.totalGlobalMem >> 20) << " MiB" << std::endl;
std::cout << "Shared Memory per SM: " << (properties.sharedMemPerMultiprocessor >> 10) << " KiB" << std::endl;
std::cout << "Memory Bus Width: " << properties.memoryBusWidth << " bits"
<< " (ECC " << (properties.ECCEnabled != 0 ? "enabled" : "disabled") << ")" << std::endl;
std::cout << "Memory Clock Rate: " << properties.memoryClockRate / 1000000.0F << " GHz" << std::endl;
std::cout << "=== Software Information ===" << std::endl;
std::cout << "Build time TensorRT version: " << NV_TENSORRT_MAJOR << "." << NV_TENSORRT_MINOR << "." << NV_TENSORRT_PATCH << std::endl;
std::cout << "Runtime linked TensorRT version: " << getInferLibVersion() << std::endl;
}
int main(int argc, char* argv[])
{
Args args;
const bool argsOK = parseArgs(args, argc, argv);
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
if (!argsOK)
{
std::cerr << "Invalid arguments" << std::endl;
printHelpInfo();
return EXIT_FAILURE;
}
printDeviceInfo();
if (args.batchSize.empty())
{
args.batchSize.push_back(1);
}
const int maxBatchSize = *std::max_element(args.batchSize.begin(), args.batchSize.end());
BertInference bert(args.engine, maxBatchSize, args.sequenceLength, args.enableGraph);
std::default_random_engine generator(args.randomSeed);
std::uniform_int_distribution<int> distribution(0, std::numeric_limits<int>::max());
const int pseudoVocabSize = 30522;
const int pseudoTypeVocabSize = 2;
const int maxInputSize = args.sequenceLength * maxBatchSize;
std::vector<int> testWordIds(maxInputSize);
std::vector<int> testSegmentIds(maxInputSize);
std::vector<int> testInputMask(maxInputSize);
std::generate(
testWordIds.begin(), testWordIds.end(), [&] { return distribution(generator) % pseudoVocabSize; });
std::generate(testSegmentIds.begin(), testSegmentIds.end(),
[&] { return distribution(generator) % pseudoTypeVocabSize; });
std::generate(testInputMask.begin(), testInputMask.end(), [&] { return 1; });
for (int i = 0; i < args.batchSize.size(); i++)
{
bert.run(i, args.batchSize[i], (void*) (testWordIds.data()), (void*) (testSegmentIds.data()),
(void*) (testInputMask.data()), args.warmUpRuns, args.iterations);
}
for (int i = 0; i < args.batchSize.size(); i++)
{
bert.reportTiming(i, args.batchSize[i]);
}
return EXIT_SUCCESS;
}
demo/BERT/inference.ipynb
-364
View File
@@ -1,364 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Copyright 2021 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",
"# ==============================================================================\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"https://upload.wikimedia.org/wikipedia/en/6/6d/Nvidia_image_logo.svg\" style=\"width: 90px; float: right;\">\n",
"\n",
"# QA Inference on BERT using TensorRT"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Overview\n",
"\n",
"Bidirectional Embedding Representations from Transformers (BERT), is a method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks. \n",
"\n",
"The original paper can be found here: https://arxiv.org/abs/1810.04805.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 1.a Learning objectives\n",
"\n",
"This notebook demonstrates:\n",
"- Inference on Question Answering (QA) task with BERT Base/Large model\n",
"- The use fine-tuned NVIDIA BERT models\n",
"- Use of BERT model with TRT"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Requirements\n",
"\n",
"Please refer to the ReadMe file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. BERT Inference: Question Answering\n",
"\n",
"We can run inference on a fine-tuned BERT model for tasks like Question Answering.\n",
"\n",
"Here we use a BERT model fine-tuned on a [SQuaD 2.0 Dataset](https://rajpurkar.github.io/SQuAD-explorer/) which contains 100,000+ question-answer pairs on 500+ articles combined with over 50,000 new, unanswerable questions."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 3.a Paragraph and Queries\n",
"\n",
"The paragraph and the questions can be customized by changing the text below. Note that when using models with small sequence lengths, you should use a shorter paragraph:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Paragraph:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"paragraph_text = \"The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which accomplished landing the first humans on the Moon from 1969 to 1972. First conceived during Dwight D. Eisenhower's administration as a three-man spacecraft to follow the one-man Project Mercury which put the first Americans in space, Apollo was later dedicated to President John F. Kennedy's national goal of landing a man on the Moon and returning him safely to the Earth by the end of the 1960s, which he proposed in a May 25, 1961, address to Congress. Project Mercury was followed by the two-man Project Gemini. The first manned flight of Apollo was in 1968. Apollo ran from 1961 to 1972, and was supported by the two-man Gemini program which ran concurrently with it from 1962 to 1966. Gemini missions developed some of the space travel techniques that were necessary for the success of the Apollo missions. Apollo used Saturn family rockets as launch vehicles. Apollo/Saturn vehicles were also used for an Apollo Applications Program, which consisted of Skylab, a space station that supported three manned missions in 1973-74, and the Apollo-Soyuz Test Project, a joint Earth orbit mission with the Soviet Union in 1975.\"\n",
"\n",
"# Short paragraph version for BERT models with max sequence length of 128\n",
"short_paragraph_text = \"The Apollo program was the third United States human spaceflight program. First conceived as a three-man spacecraft to follow the one-man Project Mercury which put the first Americans in space, Apollo was dedicated to President John F. Kennedy's national goal of landing a man on the Moon. The first manned flight of Apollo was in 1968. Apollo ran from 1961 to 1972 followed by the Apollo-Soyuz Test Project a joint Earth orbit mission with the Soviet Union in 1975.\"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Question:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"question_text = \"What project put the first Americans into space?\"\n",
"#question_text = \"What year did the first manned Apollo flight occur?\"\n",
"#question_text = \"What President is credited with the original notion of putting Americans in space?\"\n",
"#question_text = \"Who did the U.S. collaborate with on an Earth orbit mission in 1975?\"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this example we ask our BERT model questions related to the following paragraph:\n",
"\n",
"**The Apollo Program**\n",
"_\"The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which accomplished landing the first humans on the Moon from 1969 to 1972. First conceived during Dwight D. Eisenhower's administration as a three-man spacecraft to follow the one-man Project Mercury which put the first Americans in space, Apollo was later dedicated to President John F. Kennedy's national goal of landing a man on the Moon and returning him safely to the Earth by the end of the 1960s, which he proposed in a May 25, 1961, address to Congress. Project Mercury was followed by the two-man Project Gemini. The first manned flight of Apollo was in 1968. Apollo ran from 1961 to 1972, and was supported by the two-man Gemini program which ran concurrently with it from 1962 to 1966. Gemini missions developed some of the space travel techniques that were necessary for the success of the Apollo missions. Apollo used Saturn family rockets as launch vehicles. Apollo/Saturn vehicles were also used for an Apollo Applications Program, which consisted of Skylab, a space station that supported three manned missions in 1973-74, and the Apollo-Soyuz Test Project, a joint Earth orbit mission with the Soviet Union in 1975.\"_\n",
"\n",
"The questions and relative answers expected are shown below:\n",
"\n",
" - **Q1:** \"What project put the first Americans into space?\" \n",
" - **A1:** \"Project Mercury\"\n",
" - **Q2:** \"What program was created to carry out these projects and missions?\"\n",
" - **A2:** \"The Apollo program\"\n",
" - **Q3:** \"What year did the first manned Apollo flight occur?\"\n",
" - **A3:** \"1968\"\n",
" - **Q4:** \"What President is credited with the original notion of putting Americans in space?\"\n",
" - **A4:** \"John F. Kennedy\"\n",
" - **Q5:** \"Who did the U.S. collaborate with on an Earth orbit mission in 1975?\"\n",
" - **A5:** \"Soviet Union\"\n",
" - **Q6:** \"How long did Project Apollo run?\"\n",
" - **A6:** \"1961 to 1972\"\n",
" - **Q7:** \"What program helped develop space travel techniques that Project Apollo used?\"\n",
" - **A7:** \"Gemini Mission\"\n",
" - **Q8:** \"What space station supported three manned missions in 1973-1974?\"\n",
" - **A8:** \"Skylab\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Data Preprocessing\n",
"Let's convert the paragraph and the question to BERT input with the help of the tokenizer:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import helpers.data_processing as dp\n",
"import helpers.tokenization as tokenization\n",
"\n",
"tokenizer = tokenization.FullTokenizer(vocab_file=\"/workspace/TensorRT/demo/BERT/models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt\", do_lower_case=True)\n",
"\n",
"# The maximum number of tokens for the question. Questions longer than this will be truncated to this length.\n",
"max_query_length = 64\n",
"\n",
"# When splitting up a long document into chunks, how much stride to take between chunks.\n",
"doc_stride = 128\n",
"\n",
"# The maximum total input sequence length after WordPiece tokenization. \n",
"# Sequences longer than this will be truncated, and sequences shorter \n",
"max_seq_length = 128\n",
"\n",
"# Extract tokens from the paragraph\n",
"doc_tokens = dp.convert_doc_tokens(short_paragraph_text)\n",
"\n",
"# Extract features from the paragraph and question\n",
"features = dp.convert_example_to_features(doc_tokens, question_text, tokenizer, max_seq_length, doc_stride, max_query_length)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TensorRT Inference"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import tensorrt as trt\n",
"TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import ctypes\n",
"import os\n",
"\n",
"ctypes.CDLL(\"libnvinfer_plugin.so\", mode=ctypes.RTLD_GLOBAL)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"import collections\n",
"import numpy as np\n",
"import time\n",
"\n",
"# Load the BERT-Large Engine\n",
"with open(\"/workspace/TensorRT/demo/BERT/engines/bert_large_128.engine\", \"rb\") as f, \\\n",
" trt.Runtime(TRT_LOGGER) as runtime, \\\n",
" runtime.deserialize_cuda_engine(f.read()) as engine, \\\n",
" engine.create_execution_context() as context:\n",
"\n",
" # We always use batch size 1.\n",
" input_shape = (1, max_seq_length)\n",
" input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n",
" \n",
" # Allocate device memory for inputs.\n",
" d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n",
"\n",
" # Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n",
" # Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
"\n",
" # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
" h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
" cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n",
" d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n",
"\n",
" with CudaStreamContext() as stream:\n",
" print(\"\\nRunning Inference...\")\n",
"\n",
" _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n",
" \"NetworkOutput\",\n",
" [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids = np.ascontiguousarray(feature.input_ids.ravel())\n",
" segment_ids = np.ascontiguousarray(feature.segment_ids.ravel())\n",
" input_mask = np.ascontiguousarray(feature.input_mask.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
" \n",
" print(\"-----------------------------\")\n",
" print(\"Running Inference at {:.3f} Sentences/Sec\".format(1.0/eval_time_elapsed))\n",
" print(\"-----------------------------\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Data Post-Processing"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have the inference results let's extract the actual answer to our question"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
" # The total number of n-best predictions to generate in the nbest_predictions.json output file\n",
" n_best_size = 20\n",
"\n",
" # The maximum length of an answer that can be generated. This is needed \n",
" # because the start and end predictions are not conditioned on one another\n",
" max_answer_length = 30\n",
"\n",
" prediction, nbest_json, scores_diff_json = dp.get_predictions(doc_tokens, features,\n",
" networkOutputs, n_best_size, max_answer_length)\n",
" \n",
" for index, output in enumerate(networkOutputs):\n",
" print(\"Processing output\")\n",
" print(\"Answer: '{}'\".format(prediction))\n",
" print(\"with prob: {:.3f}%\".format(nbest_json[0]['probability'] * 100.0))\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.8.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
demo/BERT/inference.py
-275
View File
@@ -1,275 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
This script uses a prebuilt TensorRT BERT QA Engine to answer a question
based on the provided passage. It additionally includes an interactive mode
where multiple questions can be asked.
"""
import sys
import time
import json
import ctypes
import argparse
import collections
import numpy as np
import tensorrt as trt
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async
from cuda.bindings import driver as cuda, runtime as cudart
import helpers.tokenization as tokenization
import helpers.data_processing as dp
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def parse_args():
"""
Parse command line arguments
"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('-e', '--engine',
help='Path to BERT TensorRT engine')
parser.add_argument("-b", "--batch-size", default=1, help="Batch size for inference.", type=int)
parser.add_argument('-p', '--passage', nargs='*',
help='Text for paragraph/passage for BERT QA',
default='')
parser.add_argument('-pf', '--passage-file',
help='File containing input passage',
default='')
parser.add_argument('-q', '--question', nargs='*',
help='Text for query/question for BERT QA',
default='')
parser.add_argument('-qf', '--question-file',
help='File containing input question',
default='')
parser.add_argument('-sq', '--squad-json',
help='SQuAD json file',
default='')
parser.add_argument('-o', '--output-prediction-file',
help='Output prediction file for SQuAD evaluation',
default='./predictions.json')
parser.add_argument('-v', '--vocab-file',
help='Path to file containing entire understandable vocab')
parser.add_argument('-s', '--sequence-length',
help='The sequence length to use. Defaults to 128',
default=128, type=int)
parser.add_argument('--max-query-length',
help='The maximum length of a query in number of tokens. Queries longer than this will be truncated',
default=64, type=int)
parser.add_argument('--max-answer-length',
help='The maximum length of an answer that can be generated',
default=30, type=int)
parser.add_argument('--n-best-size',
help='Total number of n-best predictions to generate in the nbest_predictions.json output file',
default=20, type=int)
parser.add_argument('--doc-stride',
help='When splitting up a long document into chunks, what stride to take between chunks',
default=128, type=int)
args, _ = parser.parse_known_args()
return args
if __name__ == '__main__':
args = parse_args()
paragraph_text = None
squad_examples = None
output_prediction_file = None
if not args.passage == '':
paragraph_text = ' '.join(args.passage)
elif not args.passage_file == '':
f = open(args.passage_file, 'r')
paragraph_text = f.read()
elif not args.squad_json == '':
squad_examples = dp.read_squad_json(args.squad_json)
output_prediction_file = args.output_prediction_file
else:
paragraph_text = input("Paragraph: ")
question_text = None
if not args.question == '':
question_text = ' '.join(args.question)
elif not args.question_file == '':
f = open(args.question_file, 'r')
question_text = f.read()
tokenizer = tokenization.FullTokenizer(vocab_file=args.vocab_file, do_lower_case=True)
# When splitting up a long document into chunks, how much stride to take between chunks.
doc_stride = args.doc_stride
# The maximum total input sequence length after WordPiece tokenization.
# Sequences longer than this will be truncated, and sequences shorter
max_seq_length = args.sequence_length
def question_features(tokens, question):
# Extract features from the paragraph and question
return dp.convert_example_to_features(tokens, question, tokenizer, max_seq_length, doc_stride, args.max_query_length)
# Import necessary plugins for demoBERT
plugin_lib_name = "nvinfer_plugin_10.dll" if sys.platform == "win32" else "libnvinfer_plugin.so"
env_name_to_add_path = "PATH" if sys.platform == "win32" else "LD_LIBRARY_PATH"
handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)
if not handle:
raise RuntimeError("Could not load plugin library. Is `{}` on your {}?".format(plugin_lib_name, env_name_to_add_path))
# The first context created will use the 0th profile. A new context must be created
# for each additional profile needed. Here, we only use batch size 1, thus we only need the first profile.
with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, \
runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
# select engine profile
selected_profile = -1
for idx in range(engine.num_optimization_profiles):
profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx)
if profile_shape[0][0] <= args.batch_size and profile_shape[2][0] >= args.batch_size and profile_shape[0][1] <= max_seq_length and profile_shape[2][1] >= max_seq_length:
selected_profile = idx
break
if selected_profile == -1:
raise RuntimeError("Could not find any profile that can run batch size {}.".format(args.batch_size))
# Create a stream in which to copy inputs/outputs and run inference.
with CudaStreamContext() as stream:
context.set_optimization_profile_async(selected_profile, stream.stream)
binding_idx_offset = selected_profile * engine.num_io_tensors
# Specify input shapes. These must be within the min/max bounds of the active profile
# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.
input_shape = (args.batch_size, max_seq_length)
input_nbytes = trt.volume(input_shape) * trt.int32.itemsize
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
# Allocate device memory for inputs.
d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = np.empty(tuple(context.get_tensor_shape("logits_out")), dtype=np.float32)
cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))
# Pin the memory for faster transfers
d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))
def inference(features, tokens):
global h_output
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), args.batch_size, axis=0)
segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), args.batch_size, axis=0)
input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), args.batch_size, axis=0)
input_ids = np.ascontiguousarray(input_ids_batch.ravel())
segment_ids = np.ascontiguousarray(segment_ids_batch.ravel())
input_mask = np.ascontiguousarray(input_mask_batch.ravel())
eval_start_time = time.time()
memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)
bindings = [0 for _ in range(binding_idx_offset)] + [int(d_inp) for d_inp in d_inputs] + [int(d_output)]
# allocate address for IO tensor
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
# Run inference
context.execute_async_v3(stream_handle=stream.stream)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Transfer predictions back from GPU
memcpy_device_to_host_async(h_output, d_output, stream.stream)
stream.synchronize()
# Only retrieve and post-process the first batch
batch = h_output[0]
networkOutputs.append(_NetworkOutput(
start_logits = np.array(batch.squeeze()[:, 0]),
end_logits = np.array(batch.squeeze()[:, 1]),
feature_index = feature_index
))
eval_time_elapsed /= len(features)
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
return eval_time_elapsed, prediction, nbest_json
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(args.batch_size/eval_time_elapsed))
print("------------------------")
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
if squad_examples:
all_predictions = collections.OrderedDict()
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# free allocated memory
for d_input in d_inputs:
cuda_call(cudart.cudaFree(d_input))
cuda_call(cudart.cudaFree(d_output))
cuda_call(cudart.cudaHostUnregister(h_output))
demo/BERT/inference_c.py
-224
View File
@@ -1,224 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
This script uses a prebuilt TensorRT BERT QA Engine to answer a question
based on the provided passage. It additionally includes an interactive mode
where multiple questions can be asked.
"""
import os
import sys
import time
import json
import argparse
import collections
import numpy as np
import helpers.tokenization as tokenization
import helpers.data_processing as dp
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'build'))
import infer_c
def parse_args():
"""
Parse command line arguments
"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('-e', '--engine',
help='Path to BERT TensorRT engine')
parser.add_argument('-p', '--passage', nargs='*',
help='Text for paragraph/passage for BERT QA',
default='')
parser.add_argument('-pf', '--passage-file',
help='File containing input passage',
default='')
parser.add_argument('-q', '--question', nargs='*',
help='Text for query/question for BERT QA',
default='')
parser.add_argument('-qf', '--question-file',
help='File containing input question',
default='')
parser.add_argument('-sq', '--squad-json',
help='SQuAD json file',
default='')
parser.add_argument('-o', '--output-prediction-file',
help='Output prediction file for SQuAD evaluation',
default='./predictions.json')
parser.add_argument('-v', '--vocab-file',
help='Path to file containing entire understandable vocab')
parser.add_argument('-s', '--sequence-length',
help='The sequence length to use. Defaults to 128',
default=128, type=int)
parser.add_argument('--max-query-length',
help='The maximum length of a query in number of tokens. Queries longer than this will be truncated',
default=64, type=int)
parser.add_argument('--max-answer-length',
help='The maximum length of an answer that can be generated',
default=30, type=int)
parser.add_argument('--n-best-size',
help='Total number of n-best predictions to generate in the nbest_predictions.json output file',
default=20, type=int)
parser.add_argument('--enable-graph',
help='Enable CUDA Graph support',
action='store_true',
default=False)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
paragraph_text = None
squad_examples = None
output_prediction_file = None
if not args.passage == '':
paragraph_text = ' '.join(args.passage)
elif not args.passage_file == '':
f = open(args.passage_file, 'r')
paragraph_text = f.read()
elif not args.squad_json == '':
squad_examples = dp.read_squad_json(args.squad_json)
output_prediction_file = args.output_prediction_file
else:
paragraph_text = input("Paragraph: ")
question_text = None
if not args.question == '':
question_text = ' '.join(args.question)
elif not args.question_file == '':
f = open(args.question_file, 'r')
question_text = f.read()
tokenizer = tokenization.FullTokenizer(vocab_file=args.vocab_file, do_lower_case=True)
# When splitting up a long document into chunks, how much stride to take between chunks.
doc_stride = 128
# The maximum total input sequence length after WordPiece tokenization.
# Sequences longer than this will be truncated, and sequences shorter
max_seq_length = args.sequence_length
def question_features(tokens, question):
# Extract features from the paragraph and question
return dp.convert_example_to_features(tokens, question, tokenizer, max_seq_length, doc_stride, args.max_query_length)
# The first context created will use the 0th profile. A new context must be created
# for each additional profile needed. Here, we only use batch size 1, thus we only need the first profile.
# We always use batch size 1.
# Specify input shapes as (max_seq_length, 1).
# These must be within the min/max bounds of the active profile (0th profile in this case)
# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.
bert = infer_c.bert_inf(args.engine, 1, max_seq_length, args.enable_graph)
bert.prepare(1)
def inference(features, tokens):
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
input_ids = np.ascontiguousarray(feature.input_ids.ravel())
segment_ids = np.ascontiguousarray(feature.segment_ids.ravel())
input_mask = np.ascontiguousarray(feature.input_mask.ravel())
eval_start_time = time.time()
# Run inference
h_output = bert.run(input_ids, segment_ids, input_mask)
eval_time_elapsed += (time.time() - eval_start_time)
# Data Post-processing
if len(h_output.shape) == 1:
S = int(h_output.shape[0] / 2)
networkOutputs.append(_NetworkOutput(
start_logits = np.array(h_output[0:S]),
end_logits = np.array(h_output[S:S*2]),
feature_index = feature_index
))
else:
for index, batch in enumerate(h_output):
networkOutputs.append(_NetworkOutput(
start_logits = np.array(batch.squeeze()[:, 0]),
end_logits = np.array(batch.squeeze()[:, 1]),
feature_index = feature_index
))
eval_time_elapsed /= len(features)
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
return eval_time_elapsed, prediction, nbest_json
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(1.0/eval_time_elapsed))
print("------------------------")
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
if squad_examples:
all_predictions = collections.OrderedDict()
for example_index, example in enumerate(squad_examples):
print("Processing example {} of {}".format(example_index+1, len(squad_examples)), end="\r")
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
demo/BERT/inference_varseqlen.py
-264
View File
@@ -1,264 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
This script uses a prebuilt TensorRT BERT QA Engine to answer a question
based on the provided passage. It additionally includes an interactive mode
where multiple questions can be asked.
"""
import sys
import time
import json
import ctypes
import argparse
import collections
import numpy as np
import tensorrt as trt
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async
from cuda.bindings import driver as cuda, runtime as cudart
import helpers.tokenization as tokenization
import helpers.data_processing as dp
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def parse_args():
"""
Parse command line arguments
"""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('-e', '--engine',
help='Path to BERT TensorRT engine')
parser.add_argument('-p', '--passage', nargs='*',
help='Text for paragraph/passage for BERT QA',
default='')
parser.add_argument('-pf', '--passage-file',
help='File containing input passage',
default='')
parser.add_argument('-q', '--question', nargs='*',
help='Text for query/question for BERT QA',
default='')
parser.add_argument('-qf', '--question-file',
help='File containing input question',
default='')
parser.add_argument('-sq', '--squad-json',
help='SQuAD json file',
default='')
parser.add_argument('-o', '--output-prediction-file',
help='Output prediction file for SQuAD evaluation',
default='./predictions.json')
parser.add_argument('-v', '--vocab-file',
help='Path to file containing entire understandable vocab')
parser.add_argument('-s', '--sequence-length',
help='The sequence length to use. Defaults to 128',
default=128, type=int)
parser.add_argument('--max-query-length',
help='The maximum length of a query in number of tokens. Queries longer than this will be truncated',
default=64, type=int)
parser.add_argument('--max-answer-length',
help='The maximum length of an answer that can be generated',
default=30, type=int)
parser.add_argument('--n-best-size',
help='Total number of n-best predictions to generate in the nbest_predictions.json output file',
default=20, type=int)
parser.add_argument('--doc-stride',
help='When splitting up a long document into chunks, what stride to take between chunks',
default=128, type=int)
args, _ = parser.parse_known_args()
return args
if __name__ == '__main__':
args = parse_args()
paragraph_text = None
squad_examples = None
output_prediction_file = None
if not args.passage == '':
paragraph_text = ' '.join(args.passage)
elif not args.passage_file == '':
f = open(args.passage_file, 'r')
paragraph_text = f.read()
elif not args.squad_json == '':
squad_examples = dp.read_squad_json(args.squad_json)
output_prediction_file = args.output_prediction_file
else:
paragraph_text = input("Paragraph: ")
question_text = None
if not args.question == '':
question_text = ' '.join(args.question)
elif not args.question_file == '':
f = open(args.question_file, 'r')
question_text = f.read()
tokenizer = tokenization.FullTokenizer(vocab_file=args.vocab_file, do_lower_case=True)
# When splitting up a long document into chunks, how much stride to take between chunks.
doc_stride = args.doc_stride
# The maximum total input sequence length after WordPiece tokenization.
# Sequences longer than this will be truncated, and sequences shorter
max_seq_length = args.sequence_length
def question_features(tokens, question):
# Extract features from the paragraph and question
return dp.convert_example_to_features(tokens, question, tokenizer, max_seq_length, doc_stride, args.max_query_length)
# Import necessary plugins for demoBERT
plugin_lib_name = "nvinfer_plugin_10.dll" if sys.platform == "win32" else "libnvinfer_plugin.so"
env_name_to_add_path = "PATH" if sys.platform == "win32" else "LD_LIBRARY_PATH"
handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)
if not handle:
raise RuntimeError("Could not load plugin library. Is `{}` on your {}?".format(plugin_lib_name, env_name_to_add_path))
# The first context created will use the 0th profile. A new context must be created
# for each additional profile needed. Here, we only use batch size 1, thus we only need the first profile.
with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, \
runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
# Create a stream in which to copy inputs/outputs and run inference.
with CudaStreamContext() as stream:
# select engine profile
context.set_optimization_profile_async(0, stream.stream)
input_nbytes = max_seq_length * trt.int32.itemsize
# Allocate device memory for inputs.
d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)]
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = np.empty((2 * max_seq_length), dtype=np.float32)
cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))
d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))
def inference(features, tokens):
global h_output
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
B = 1
S = np.sum(feature.input_mask)
input_ids = feature.input_ids[0:S]
segment_ids = feature.segment_ids[0:S]
cu_seq_lens = np.array([0, S], dtype=np.int32);
input_dim0_shape = {"input_ids": S, "segment_ids": S, "cu_seqlens": 2, "max_seqlen": S}
for name, val in input_dim0_shape.items():
if context.get_tensor_shape(name)[0] != val:
context.set_input_shape(name, (val,))
h_input_ids = np.ascontiguousarray(input_ids.ravel())
h_segment_ids = np.ascontiguousarray(segment_ids.ravel())
h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel())
eval_start_time = time.time()
memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream)
# Setup tensor address
bindings = [int(d_inputs[i]) for i in range(4)] + [int(d_output)]
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
# Run inference
context.execute_async_v3(stream_handle=stream.stream)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Transfer predictions back from GPU
memcpy_device_to_host_async(h_output, d_output, stream.stream)
stream.synchronize()
# Only retrieve and post-process the first batch
networkOutputs.append(_NetworkOutput(
start_logits = np.array(h_output[0:S]),
end_logits = np.array(h_output[S:S*2]),
feature_index = feature_index
))
eval_time_elapsed /= len(features)
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
return eval_time_elapsed, prediction, nbest_json
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(1/eval_time_elapsed))
print("------------------------")
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
if squad_examples:
all_predictions = collections.OrderedDict()
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# free allocated memory
for d_input in d_inputs:
cuda_call(cudart.cudaFree(d_input))
cuda_call(cudart.cudaFree(d_output))
cuda_call(cudart.cudaHostUnregister(h_output))
demo/BERT/notebooks/BERT-TRT-FP16.ipynb
-371
View File
@@ -1,371 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "herbal-royalty",
"metadata": {},
"outputs": [],
"source": [
"# Copyright 2021 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",
"id": "norwegian-dakota",
"metadata": {},
"source": [
"<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
"\n",
"# BERT QA Inference with TensorRT FP16\n",
"\n",
"\n",
"Bidirectional Encoder Representations from Transformers ([BERT](https://arxiv.org/abs/1810.04805)) is a method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks. \n",
"\n",
"BERT provided a leap in accuracy for NLU tasks that brought high-quality language-based services within the reach of companies across many industries. To use the model in production, you need to consider factors such as latency, in addition to accuracy, which influences end user satisfaction with a service. BERT requires significant compute during inference due to its 12/24-layer stacked multi-head attention network. This has posed a challenge for companies to deploy BERT as part of real-time applications until now.\n",
"\n",
"NVIDIA® [TensorRT](https://developer.nvidia.com/tensorrt)™ is an SDK for high-performance deep learning inference. TensorRT provides INT8 and FP16 optimizations for production deployments of deep learning inference applications such as video streaming, speech recognition, recommendation, fraud detection, and natural language processing.\n",
"TensorRT optimizations for BERT allows you to perform inference in 2.2 ms on T4 GPUs. This is 17x faster than CPU-only platforms and is well within the 10ms latency budget necessary for conversational AI applications. These optimizations make it practical to use BERT in production, for example, as part of a conversation AI service.\n",
"\n",
"This notebook demonstrates the inference of BERT models for question and answering applications with TensorRT in FP16 mode.\n",
"\n",
"## Pre-requisite\n",
"Follow the instruction at https://github.com/NVIDIA/TensorRT to build the TensorRT-OSS docker container required to run this notebook.\n",
"\n",
"\n",
"## Content\n",
"1. [Download data and model](#1)\n",
"1. [Building a FP16 TensorRT optimized BERT model](#2)\n",
"1. [Running inference examples](#3)\n",
"1. [Inference benchmarking](#4)\n"
]
},
{
"cell_type": "markdown",
"id": "nasty-frequency",
"metadata": {},
"source": [
"<a id=\"1\"></a>\n",
"\n",
"## 1. Download data and model\n",
"First, we download the \n",
"Stanford Question Answering Dataset ([SQuAD](https://rajpurkar.github.io/SQuAD-explorer/)) and a pre-trained BERT QA model from the NVIDIA GPU Cloud ([NGC](https://ngc.nvidia.com/catalog/models/nvidia:bert_pyt_ckpt_base_qa_squad11_amp)).\n",
"### SQuAD dataset\n",
"\n",
"Stanford Question Answering Dataset ([SQuAD](https://rajpurkar.github.io/SQuAD-explorer/)) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "studied-sheffield",
"metadata": {},
"outputs": [],
"source": [
"!bash ../scripts/download_squad.sh"
]
},
{
"cell_type": "markdown",
"id": "laughing-arthur",
"metadata": {},
"source": [
"### Fine-tuned BERT Large Model download\n",
"\n",
"Many AI applications have common needs: classification, object detection, language translation, text-to-speech, recommender engines, sentiment analysis, and more. When developing applications with these capabilities, it is much faster to start with a model that is pre-trained and then tune it for a specific use case. The NGC [catalog](https://ngc.nvidia.com/catalog/models) offers pre-trained models for a variety of common AI tasks that are optimized for NVIDIA Tensor Core GPUs, and can be easily re-trained by updating just a few layers, saving valuable time.\n",
"\n",
"Herein, we download a pretrained, fine-tuned BERT large model, trained with automatic mixed precision, from NGC."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "signed-symposium",
"metadata": {},
"outputs": [],
"source": [
"!bash ../scripts/download_model.sh large 384 v2"
]
},
{
"cell_type": "markdown",
"id": "cc5159a2",
"metadata": {},
"source": [
"### Install extra dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2ceba10e",
"metadata": {},
"outputs": [],
"source": [
"!cd /tmp && git clone https://github.com/vinhngx/transformers && cd transformers && pip install .\n",
"!pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "alive-slovakia",
"metadata": {},
"source": [
"<a id=\"2\"></a>\n",
"\n",
"## 2. Building a FP16 TensorRT optimized BERT model\n",
"\n",
"In this section, we will be optimizing the BERT model for inference with TensorRT using FP16. The overal workflow is as below.\n",
"\n",
"<img src=\"Figure-1-generating-bert-trt.png\">\n",
"\n",
"To optimize BERT with TensorRT, we focused on optimizing the transformer cell. Since several Transformer cells are stacked in BERT, we were able to achieve significant performance gains through this set of optimizations. We use custom plugins that accelerate key operations in the Transformer Encoder elements in a BERT model. The plugins fuse multiple operations into a sub-graph in a single CUDA kernel. Each sub-graph consists of several elementary computations, each of which requires a read and write to the global memory of the GPU (i.e. the slowest on-device memory). By fusing the elementary operations together into a single CUDA kernel we allow for the computation to happen on a larger sub-graph while visiting the global memory a minimal amount of times. \n",
"<img src=\"Figure-5-optimizations-through-trt.jpg\">\n",
"\n",
"For more information, see our developer [blog](https://developer.nvidia.com/blog/nlu-with-tensorrt-bert/)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "leading-reliance",
"metadata": {},
"outputs": [],
"source": [
"import tensorrt as trt;\n",
"TRT_VERSION = trt.__version__\n",
"\n",
"print(\"TensorRT version: {}\".format(TRT_VERSION))\n",
"!mkdir -p engines_$TRT_VERSION"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "studied-profession",
"metadata": {},
"outputs": [],
"source": [
"# Maximum TensorRT inference batch size\n",
"BATCH_SIZE = 128"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "unique-batman",
"metadata": {},
"outputs": [],
"source": [
"# Build BERT TensorRT FP16 model from NGC checkpoint\n",
"!python3 ../builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -w 40000 -o engines_$TRT_VERSION/bert_large_384.engine -b 1 -b $BATCH_SIZE -s 384 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1"
]
},
{
"cell_type": "markdown",
"id": "authorized-assignment",
"metadata": {},
"source": [
"<a id=\"3\"></a>\n",
"\n",
"## 3. Running inference examples\n",
"\n",
"Now that we've got a TensorRT engine, the inference workflow using the optimized network is as follows:\n",
"\n",
" - Start the TensorRT runtime with this engine.\n",
" - Feed a passage and a question to the TensorRT runtime and receive as output the answer predicted by the network.\n",
"\n",
"<img src=\"Figure-2-workflow-to-perform-inference-with-trt.png\">"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "incorporate-psychiatry",
"metadata": {},
"outputs": [],
"source": [
"PASSAGE = 'TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps'\\\n",
"'such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops'\\\n",
"'and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep'\\\n",
"'learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps.'\n",
"QUESTION=\"What is TensorRT?\"\n",
"\n",
"!python3 ../inference.py -e engines_$TRT_VERSION/bert_large_384.engine -s 384 -p $PASSAGE -q $QUESTION -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt"
]
},
{
"cell_type": "markdown",
"id": "340fa7fb-d997-4fbf-b3bf-684085584c94",
"metadata": {},
"source": [
"Let's ask a different question. Feel free to plugin your own question."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "legal-brief",
"metadata": {},
"outputs": [],
"source": [
"QUESTION=\"What is included in TensorRT?\"\n",
"\n",
"!python3 ../inference.py -e engines_$TRT_VERSION/bert_large_384.engine -s 384 -p $PASSAGE -q $QUESTION -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt"
]
},
{
"cell_type": "markdown",
"id": "korean-simpson",
"metadata": {},
"source": [
"## Validation on the SQuAD dev set\n",
"Next, we will assess the accuracy of the TensorRT-optimized FP16 BERT model on the SQuAD dev set. \n",
"\n",
"There are two dominant metrics used by many question answering datasets, including SQuAD: exact match (EM) and F1 score. These scores are computed on individual question+answer pairs. When multiple correct answers are possible for a given question, the maximum score over all possible correct answers is computed. Overall EM and F1 scores are computed for a model by averaging over the individual example scores.\n",
"\n",
"### Exact Match\n",
"\n",
"This metric is as simple as it sounds. For each question+answer pair, if the characters of the model's prediction exactly match the characters of (one of) the True Answer(s), EM = 1, otherwise EM = 0. This is a strict all-or-nothing metric; being off by a single character results in a score of 0. When assessing against a negative example, if the model predicts any text at all, it automatically receives a 0 for that example.\n",
"\n",
"### F1\n",
"\n",
"F1 score is a common metric for classification problems, and widely used in QA. It is appropriate when we care equally about precision and recall. In this case, it's computed over the individual words in the prediction against those in the True Answer. The number of shared words between the prediction and the truth is the basis of the F1 score: precision is the ratio of the number of shared words to the total number of words in the prediction, and recall is the ratio of the number of shared words to the total number of words in the ground truth.\n",
"\n",
"For more info, see [reference](https://qa.fastforwardlabs.com/no%20answer/null%20threshold/bert/distilbert/exact%20match/f1/robust%20predictions/2020/06/09/Evaluating_BERT_on_SQuAD.html#Metrics-for-QA).\n",
"\n",
"Herein, we verify that the TensorRT model achieves a state-of-the-art accuracy of 90% F1 score on the SQuAD development set."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "loose-musical",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../inference.py -e engines_$TRT_VERSION/bert_large_384.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions-bert_large_384.json\n",
"!python3 ../squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions-bert_large_384.json 90\n"
]
},
{
"cell_type": "markdown",
"id": "functional-smile",
"metadata": {},
"source": [
"<a id=\"4\"></a>\n",
"\n",
"## 4. Inference benchmarking\n",
"BERT can be applied both for online and offline use cases. Online NLU applications, such as conversational AI, place tight latency budgets during inference. Several models need to execute in a sequence in response to a single user query. When used as a service, the total time a customer experiences includes compute time as well as input and output network latency. Longer times lead to a sluggish performance and a poor customer experience.\n",
"\n",
"While the exact latency available for a single model can vary by application, several real-time applications need the language model to execute in under 10 ms. Using a Tesla T4 GPU, BERT optimized with TensorRT can perform inference in 2.2 ms for a QA task similar to available in SQuAD with batch size =1 and sequence length = 128. Using the TensorRT optimized sample, you can execute up to a batch size of 8 for BERT-base and even higher batch sizes for models with fewer Transformer layers within the 10 ms latency budget. It took 40 ms to execute the same task with highly optimized code on a CPU-only platform for batch size = 1, while higher batch sizes did not run to completion and exit with errors.\n",
"\n",
"<img src=\"Figure-6-Compute-latency.jpg\">"
]
},
{
"cell_type": "markdown",
"id": "17fe9312-7b20-4997-8f00-81b2d8d3ba9c",
"metadata": {},
"source": [
"Next, we will perform a couple of inference benchmarks with different batch sizes:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "attractive-binary",
"metadata": {},
"outputs": [],
"source": [
"BATCH_SIZE=1\n",
"!python3 ../perf.py -e ./engines_$TRT_VERSION/bert_large_384.engine -b $BATCH_SIZE -s 384 -i 100 -w 20"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "future-courage",
"metadata": {},
"outputs": [],
"source": [
"BATCH_SIZE = 64\n",
"!python3 ../perf.py -e ./engines_$TRT_VERSION/bert_large_384.engine -b $BATCH_SIZE -s 384 -i 100 -w 20"
]
},
{
"cell_type": "markdown",
"id": "floating-museum",
"metadata": {},
"source": [
"### BERT base model\n",
"\n",
"We repeat the same process with another BERT model, the BERT-base model (110M parameters) with sequence length 128."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "equivalent-niger",
"metadata": {},
"outputs": [],
"source": [
"!bash ../scripts/download_model.sh base 128 v2"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dressed-prophet",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../builder.py -m models/fine-tuned/bert_tf_ckpt_base_qa_squad2_amp_128_v19.03.1/model.ckpt -w 40000 -o engines_$TRT_VERSION/bert_base_128.engine -b 1 -b $BATCH_SIZE -s 128 --fp16 -c models/fine-tuned/bert_tf_ckpt_base_qa_squad2_amp_128_v19.03.1"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "handmade-contrary",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../perf.py -e ./engines_$TRT_VERSION/bert_base_128.engine -b 1 -s 128 -i 100 -w 20"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
demo/BERT/notebooks/BERT-TRT-INT8-QAT-sparse.ipynb
-339
View File
@@ -1,339 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "cb33810e-c6be-48cc-84a3-a79794c6a4b0",
"metadata": {},
"outputs": [],
"source": [
"# Copyright 2021 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",
"id": "104213a3-18c3-4384-b923-59e35a163093",
"metadata": {},
"source": [
"<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
"\n",
"# BERT QA Inference on TensorRT INT8: Quantization Aware Training (QAT) and Structured Sparsity\n",
"\n",
"This notebook demonstrates the use of BERT model with TensorRT in QAT INT8 and structured sparsity mode. These are two new features introduced since TensorRT 8.\n",
"\n",
"**Quantization Aware Training**: Using INT8 precision with quantization scales obtained from Post-Training Quantization (PTQ) can produce additional performance gains, but may also result in accuracy loss. Alternatively, for PyTorch-trained models, NVIDIA PyTorch-Quantization [toolkit](https://docs.nvidia.com/deeplearning/tensorrt/pytorch-quantization-toolkit/docs/index.html) can be leveraged to perform quantized fine tuning (a.k.a. Quantization Aware Training or QAT) and generate the INT8 quantization scales as part of training. This generally results in higher accuracy compared to PTQ.\n",
"\n",
"**Structured Sparsity**: Fine-grained 2:4 structured sparsity support introduced in NVIDIA Ampere GPUs can produce significant performance gains in BERT inference. The network is first trained using dense weights, then fine-grained structured pruning is applied, and finally the remaining non-zero weights are fine-tuned with additional training steps. This method results in virtually no loss in inferencing accuracy.\n",
"\n",
"For more information on sparsity and how to train sparse models, see the GTC [talk](https://gtc21.event.nvidia.com/media/Making%20the%20Most%20of%20Structured%20Sparsity%20in%20the%20NVIDIA%20Ampere%20Architecture%20%5BS31552%5D/1_0j8hi0r7) titled \"Making the Most of Structured Sparsity in the NVIDIA Ampere Architecture.\"\n",
"\n",
"TensorRT since version 8 supports both QAT and structured-sparsity trained networks. \n",
"\n",
"## Pre-requisite\n",
"Follow the instruction at https://github.com/NVIDIA/TensorRT to build the TensorRT-OSS docker container required to run this notebook.\n",
"\n",
"## Content\n",
"1. [Download data and model](#1)\n",
"1. [Building a INT8-Sparsity TensorRT optimized BERT model](#2)\n",
"1. [Running inference examples](#3)\n",
"1. [Inference benchmarking](#4)\n"
]
},
{
"cell_type": "markdown",
"id": "99997954-bb39-42ba-8247-c2e47d213b29",
"metadata": {},
"source": [
"<a id=\"1\"></a>\n",
"\n",
"## Download data and model\n",
"First, we download the \n",
"Stanford Question Answering Dataset ([SQuAD](https://rajpurkar.github.io/SQuAD-explorer/)) dataset and a pre-trained BERT QA model from NVIDIA GPU Cloud ([NGC](https://ngc.nvidia.com/catalog/models/nvidia:bert_pyt_ckpt_base_qa_squad11_amp)).\n",
"### SQUAD dataset\n",
"\n",
"Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "93cfa328-1d7e-49ea-aa4e-23005f30b460",
"metadata": {},
"outputs": [],
"source": [
"!bash ../scripts/download_squad.sh"
]
},
{
"cell_type": "markdown",
"id": "6a699716-e432-4520-8b8e-fc12e039f3b5",
"metadata": {
"jupyter": {
"outputs_hidden": true
},
"tags": []
},
"source": [
"### Fine-tuned BERT Large Model download\n",
"\n",
"Many AI applications have common needs: classification, object detection, language translation, text-to-speech, recommender engines, sentiment analysis, and more. When developing applications with these capabilities, it is much faster to start with a model that is pre-trained and then tune it for a specific use case. The NGC [catalog](https://ngc.nvidia.com/catalog/models) offers pre-trained models for a variety of common AI tasks that are optimized for NVIDIA Tensor Core GPUs, and can be easily re-trained by updating just a few layers, saving valuable time.\n",
"\n",
"Herein, we download a pretrained, fine-tuned BERT large model, trained with automatic mixed precision, from NGC.\n",
"\n",
"To demonstrate the potential speedups from these optimizations in demoBERT, we provide the Megatron-LM transformer model finetuned for SQuAD 2.0 task with sparsity and quantization.\n",
"The sparse weights are generated by finetuning with INT8 Quantization Aware Training recipe. "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "181f2505-5cf2-4ccf-9bc6-a08290b41578",
"metadata": {},
"outputs": [],
"source": [
"!bash ../scripts/download_model.sh 384 # BERT-large model checkpoint\n",
"!bash ../scripts/download_model.sh pyt megatron-large int8-qat sparse # Megatron-LM model weights"
]
},
{
"cell_type": "markdown",
"id": "5fa4218b",
"metadata": {},
"source": [
"### Install extra dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8c51b43",
"metadata": {},
"outputs": [],
"source": [
"!cd /tmp && git clone https://github.com/vinhngx/transformers && cd transformers && pip install .\n",
"!pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "d879fba2-1bfe-45d9-9fda-c816f545ca87",
"metadata": {},
"source": [
"<a id=\"2\"></a>\n",
"\n",
"## 2. Building an INT8-Sparsity TensorRT optimized BERT model\n",
"\n",
"In this section, we will be optimizing the BERT model for inference with TRT using INT8 precision while leveraging structured sparsity.\n",
"\n",
"\n",
"## What Is Sparsity in AI?\n",
"The brain cannot be fully connected: 10^11 nerve cells, but only upto 10^4 connections each.\n",
"In AI inference and machine learning, sparsity refers to a matrix of numbers that includes many zeros or values that will not significantly impact a calculation.\n",
"\n",
"### Fine-grained structured sparsity\n",
"NVIDIA Ampere GPU architecture introduces the concept of fine-grained structured sparsity. On the NVIDIA A100 GPU, the structure manifests as a 2:4 pattern: out of every four elements, at least two must be zero. This reduces the data footprint and bandwidth of one matrix multiply (also known as GEMM) operand by 2x and doubles throughput by skipping the computation of the zero values using new NVIDIA Sparse Tensor Cores.\n",
"\n",
"<img src=\"structured_spare_matrix.jpg\" style=\"width: 200px;\"/>\n",
"\n",
"Fine-grained structured sparsity results in even load balancing, regular memory accesses, and 2x math efficiency with no loss in network accuracy.\n",
"\n",
"<img src=\"sparsity-diagram-600x338-r3.jpg\">\n",
"\n",
"### Training recipe\n",
"[ASP](https://github.com/NVIDIA/apex/tree/master/apex/contrib/sparsity) (Automatic SParsity) is a tool that enables sparse training and inference for PyTorch models by adding 2 lines of Python.\n",
"\n",
"\n",
"See our GTC session titled [Integer Quantization for DNN Inference Acceleration](https://developer.nvidia.com/gtc/2020/video/s22075-vid) for further info.\n",
"\n",
"\n",
"In the following code, we will be downloading a pretrained BERT model from NGC that has been trained with QAT and ASP. This model is ready to be used in TRT INT8 mode while leveraging structured sparsity features on NVIDIA A100 GPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2599b25d-cbaf-4544-bb9d-64e378febccc",
"metadata": {},
"outputs": [],
"source": [
"import tensorrt as trt;\n",
"TRT_VERSION = trt.__version__\n",
"\n",
"print(\"TensorRT version: {}\".format(TRT_VERSION))\n",
"!mkdir -p engines_$TRT_VERSION"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1a13df18-2b61-48ba-a1e8-f047100a209a",
"metadata": {},
"outputs": [],
"source": [
"BATCH_SIZE = 128"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "52e7eddc-9536-44b8-9b4b-5c70ca040fc2",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"!export CKPT_PATH=models/fine-tuned/bert_pyt_statedict_megatron_sparse_int8qat_v21.03.0/bert_pyt_statedict_megatron_sparse_int8_qat\n",
"!python3 ../builder_varseqlen.py -w 40000 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -b $BATCH_SIZE -s 384 -o engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine --fp16 --int8 --strict -il --megatron --pickle models/fine-tuned/bert_pyt_statedict_megatron_sparse_int8qat_v21.03.0/bert_pyt_statedict_megatron_sparse_int8_qat -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -sp\n"
]
},
{
"cell_type": "markdown",
"id": "c944882f-d6ce-4719-8824-fb0b6fe682a2",
"metadata": {},
"source": [
"<a id=\"3\"></a>\n",
"## 3. Running inference examples\n",
"\n",
"Now that we've got a TensorRT engine, the inference workflow using the optimized network is as follows:\n",
"\n",
" - Start the TensorRT runtime with this engine.\n",
" - Feed a passage and a question to the TensorRT runtime and receive as output the answer predicted by the network.\n",
"\n",
"<img src=\"Figure-2-workflow-to-perform-inference-with-trt.png\">\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "be366873-b1d2-4b8f-9723-04f419a63263",
"metadata": {},
"outputs": [],
"source": [
"PASSAGE = 'TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps'\\\n",
"'such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops'\\\n",
"'and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep'\\\n",
"'learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps.'\n",
"QUESTION=\"What is TensorRT?\"\n",
"\n",
"!python3 ../inference_varseqlen.py -e engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine -p $PASSAGE -q $QUESTION -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "edfc1170-5da4-4204-ab71-3e51e54391a5",
"metadata": {},
"outputs": [],
"source": [
"QUESTION=\"What is included in TensorRT?\"\n",
"!python3 ../inference_varseqlen.py -e engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine -p $PASSAGE -q $QUESTION -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256"
]
},
{
"cell_type": "markdown",
"id": "d7751aba",
"metadata": {},
"source": [
"## Validation on the SQuAD dev set\n",
"Next, we will assess the accuracy of the TensorRT-optimized INT8 BERT model on the SQuAD dev set. \n",
"\n",
"There are two dominant metrics used by many question answering datasets, including SQuAD: exact match (EM) and F1 score. These scores are computed on individual question+answer pairs. When multiple correct answers are possible for a given question, the maximum score over all possible correct answers is computed. Overall EM and F1 scores are computed for a model by averaging over the individual example scores.\n",
"\n",
"### Exact Match\n",
"\n",
"This metric is as simple as it sounds. For each question+answer pair, if the characters of the model's prediction exactly match the characters of (one of) the True Answer(s), EM = 1, otherwise EM = 0. This is a strict all-or-nothing metric; being off by a single character results in a score of 0. When assessing against a negative example, if the model predicts any text at all, it automatically receives a 0 for that example.\n",
"\n",
"### F1\n",
"\n",
"F1 score is a common metric for classification problems, and widely used in QA. It is appropriate when we care equally about precision and recall. In this case, it's computed over the individual words in the prediction against those in the True Answer. The number of shared words between the prediction and the truth is the basis of the F1 score: precision is the ratio of the number of shared words to the total number of words in the prediction, and recall is the ratio of the number of shared words to the total number of words in the ground truth.\n",
"\n",
"For more info, see [reference](https://qa.fastforwardlabs.com/no%20answer/null%20threshold/bert/distilbert/exact%20match/f1/robust%20predictions/2020/06/09/Evaluating_BERT_on_SQuAD.html#Metrics-for-QA).\n",
"\n",
"Herein, we verify that the TensorRT INT8 model maintains a state-of-the-art accuracy of 90% F1 score on the SQuAD development set, comparable to the TensorRT FP16 model as well as the original model."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f729edf0-6b9d-4dfa-8656-3608b1b5fe60",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../inference_varseqlen.py -e engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json\n",
"!python3 ../squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90\n"
]
},
{
"cell_type": "markdown",
"id": "99df76c2-7f46-4918-8839-f5a059098af7",
"metadata": {},
"source": [
"<a id=\"4\"></a>\n",
"\n",
"## 4. Inference benchmarking\n",
"\n",
"BERT can be applied both for online and offline use cases. Online NLU applications, such as conversational AI, place tight latency budgets during inference. Several models need to execute in a sequence in response to a single user query. When used as a service, the total time a customer experiences includes compute time as well as input and output network latency. Longer times lead to a sluggish performance and a poor customer experience.\n",
"\n",
"While the exact latency available for a single model can vary by application, several real-time applications need the language model to execute in under 10 ms. Using a Tesla T4 GPU, BERT optimized with TensorRT can perform inference in 2.2 ms for a QA task similar to available in SQuAD with batch size =1 and sequence length = 128. Using the TensorRT optimized sample, you can execute up to a batch size of 8 for BERT-base and even higher batch sizes for models with fewer Transformer layers within the 10 ms latency budget. It took 40 ms to execute the same task with highly optimized code on a CPU-only platform for batch size = 1, while higher batch sizes did not run to completion and exit with errors.\n",
"\n",
"<img src=\"./Figure-6-Compute-latency.jpg\">\n",
"\n",
"Next, we will perform a couple of inference benchmarks with different batch sizes:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "21da1622-73c6-4123-bb3d-f33d5d0c4aff",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../perf_varseqlen.py -e ./engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine -b 1 -s 384 -i 1000 -w 500"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e32384f3-2e69-4c4d-ab81-c33fea9de5d4",
"metadata": {},
"outputs": [],
"source": [
"!python3 ../perf_varseqlen.py -e ./engines_$TRT_VERSION/megatron_large_seqlen384_int8qat_sparse.engine -b 64 -s 384 -i 1000 -w 500"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
demo/BERT/notebooks/Figure-1-generating-bert-trt.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 113 KiB

demo/BERT/notebooks/Figure-2-workflow-to-perform-inference-with-trt.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 45 KiB

demo/BERT/notebooks/Figure-5-optimizations-through-trt.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 66 KiB

demo/BERT/notebooks/Figure-6-Compute-latency.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 23 KiB

demo/BERT/notebooks/Q-and-A.ipynb
-483
View File
@@ -1,483 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "ceramic-encoding",
"metadata": {},
"outputs": [],
"source": [
"# Copyright 2021 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",
"# ==============================================================================\n"
]
},
{
"cell_type": "markdown",
"id": "legendary-prairie",
"metadata": {},
"source": [
"<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
"\n",
"# TensorRT: Q&A with BERT\n",
"\n",
"This notebook provides a playground for testing various BERT QA models on the CPU and GPU.\n",
"\n",
"For \"CPU - Framework (PyTorch)\" and \"GPU - Framework (PyTorch)\", a SpanBERT large model from HuggingFace model repository is employed. Inference is carried out with PyTorch in FP32 precision. All models run with batch size 1.\n",
"Average run time across 10 runs is reported.\n",
"\n",
"**Notes**: \n",
" - Prior to running this notebook, run [BERT-TRT-FP16.ipynb](BERT-TRT-FP16.ipynb) and [BERT-TRT-INT8-QAT-sparse.ipynb](BERT-TRT-INT8-QAT-sparse.ipynb) to generate the TensorRT engines."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "outdoor-indonesia",
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": []
},
"outputs": [],
"source": [
"import warnings\n",
"warnings.filterwarnings('ignore')\n",
"\n",
"import sys\n",
"sys.path.append('/workspace/TensorRT/demo/BERT')\n",
"import os\n",
"sys.path.append(os.path.abspath('..'))\n",
"\n",
"import ipywidgets as widgets\n",
"import tensorrt as trt;\n",
"TRT_VERSION = trt.__version__\n",
"print(\"TensorRT version: \", TRT_VERSION)\n",
"\n",
"import time\n",
"import json\n",
"import ctypes\n",
"import argparse\n",
"import collections\n",
"import numpy as np\n",
"import tensorrt as trt\n",
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"\n",
"from helpers import tokenization as tokenization\n",
"from helpers import data_processing as dp\n",
"\n",
"TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
"\n",
"################################################## PyTorch inference #######################################################\n",
"# Install a customized version of HuggingFace, adding DL model inference timing\n",
"#!pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html\n",
"#!rm -rf /tmp/transformers\n",
"#!cd /tmp && git clone https://github.com/vinhngx/transformers && cd transformers && pip install .\n",
"# SpanBERT large model (340M params): https://github.com/facebookresearch/SpanBERT#finetuned-models-squad-1120-relation-extraction-coreference-resolution\n",
"from transformers import BertForQuestionAnswering, AutoTokenizer\n",
"\n",
"#modelname = 'deepset/bert-base-cased-squad2'\n",
"modelname = 'mrm8488/spanbert-large-finetuned-squadv2'\n",
"model = BertForQuestionAnswering.from_pretrained(modelname)\n",
"\n",
"from transformers import pipeline\n",
"nlp = pipeline('question-answering', model=model, tokenizer=\"SpanBERT/spanbert-large-cased\")\n",
"\n",
"\n",
"model_gpu = BertForQuestionAnswering.from_pretrained(modelname).cuda()\n",
"nlp_gpu = pipeline('question-answering', model=model_gpu, tokenizer=\"SpanBERT/spanbert-large-cased\", device=0)\n",
"\n",
"################################################## TensorRT inference #######################################################\n",
"def inference_FP16(trt_context, d_inputs, h_output, d_output, features, tokens):\n",
" #global h_output\n",
" context = trt_context\n",
" \n",
" _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n",
" \"NetworkOutput\",\n",
" [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" with CudaStreamContext() as stream:\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n",
" segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n",
" input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n",
"\n",
" input_ids = np.ascontiguousarray(input_ids_batch.ravel())\n",
" segment_ids = np.ascontiguousarray(segment_ids_batch.ravel())\n",
" input_mask = np.ascontiguousarray(input_mask_batch.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
"\n",
" # The total number of n-best predictions to generate in the nbest_predictions.json output file\n",
" n_best_size = 20\n",
"\n",
" # The maximum length of an answer that can be generated. This is needed \n",
" # because the start and end predictions are not conditioned on one another\n",
" max_answer_length = 30\n",
"\n",
" prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,\n",
" networkOutputs, n_best_size, max_answer_length)\n",
"\n",
" return eval_time_elapsed, prediction, nbest_json\n",
"\n",
"def inference_INT8(trt_context, d_inputs, h_output, d_output, features, tokens):\n",
" #global h_output\n",
" context = trt_context\n",
" \n",
" _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n",
" \"NetworkOutput\",\n",
" [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" with CudaStreamContext() as stream:\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" B = 1\n",
" S = np.sum(feature.input_mask)\n",
" input_ids = feature.input_ids[0:S]\n",
" segment_ids = feature.segment_ids[0:S]\n",
" cu_seq_lens = np.array([0, S], dtype=np.int32)\n",
"\n",
" first_tensor_name = engine.get_tensor_name(0)\n",
" second_tensor_name = engine.get_tensor_name(1)\n",
" third_tensor_name = engine.get_tensor_name(2)\n",
" forth_tensor_name = engine.get_tensor_name(3)\n",
"\n",
" if context.get_tensor_shape(first_tensor_name)[0] != S:\n",
" context.set_input_shape(first_tensor_name, (S,))\n",
" if context.get_tensor_shape(second_tensor_name)[0] != S:\n",
" context.set_input_shape(second_tensor_name, (S,))\n",
" if context.get_tensor_shape(third_tensor_name)[0] != 2:\n",
" context.set_input_shape(third_tensor_name, (2,))\n",
" if context.get_tensor_shape(forth_tensor_name)[0] != S:\n",
" context.set_input_shape(forth_tensor_name, (S,))\n",
"\n",
" h_input_ids = np.ascontiguousarray(input_ids.ravel())\n",
" h_segment_ids = np.ascontiguousarray(segment_ids.ravel())\n",
" h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Only retrieve and post-process the first batch\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(h_output[0:S]),\n",
" end_logits = np.array(h_output[S:S*2]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
"\n",
" # Total number of n-best predictions to generate in the nbest_predictions.json output file\n",
" n_best_size = 20\n",
"\n",
" # The maximum length of an answer that can be generated. This is needed\n",
" # because the start and end predictions are not conditioned on one another\n",
" max_answer_length = 30\n",
"\n",
" prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,\n",
" networkOutputs, n_best_size, max_answer_length)\n",
"\n",
" return eval_time_elapsed, prediction, nbest_json\n",
" \n",
"def print_single_query(eval_time_elapsed, prediction, nbest_json):\n",
" print(\"Answer: '{}'\".format(prediction))\n",
" print(\"With probability: {:.2f}%\".format(nbest_json[0]['probability'] * 100.0))\n",
" \n",
"def question_features(tokens, question):\n",
" # Extract features from the paragraph and question\n",
" return dp.convert_example_to_features(tokens, question, tokenizer, max_seq_length, doc_stride, max_query_length)\n",
"\n",
"doc_stride = 128\n",
"max_query_length = 64\n",
"\n",
"vocab_file = \"models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt\"\n",
"tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)\n",
"\n",
"### FP16 TRT model\n",
"engine_path = \"engines_{}/bert_large_384.engine\".format(TRT_VERSION)\n",
"max_seq_length = 384\n",
"batch_size = 1\n",
"\n",
"runtime = trt.Runtime(TRT_LOGGER)\n",
"engine = runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
"context = engine.create_execution_context()\n",
"\n",
" # We always use batch size 1.\n",
"input_shape = (1, max_seq_length)\n",
"input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n",
"\n",
"# Allocate device memory for inputs.\n",
"d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n",
"\n",
"# Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n",
"# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n",
"for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
"assert context.all_binding_shapes_specified\n",
"\n",
"# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
"h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
"cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n",
"d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n",
"\n",
"# Create a stream in which to copy inputs/outputs and run inference.\n",
"with CudaStreamContext() as stream:\n",
" ### INT8 TRT model\n",
" engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
" max_seq_length = 384\n",
"\n",
" INT8_runtime = trt.Runtime(TRT_LOGGER)\n",
" INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
" INT8_context = INT8_engine.create_execution_context()\n",
"\n",
" # select engine profile\n",
" INT8_context.set_optimization_profile_async(0, stream.stream)\n",
"\n",
" input_nbytes = max_seq_length * trt.int32.itemsize\n",
"\n",
" # Allocate device memory for inputs.\n",
" INT8_d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)]\n",
"\n",
" # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
" INT8_h_output = np.empty((2 * max_seq_length), dtype=np.float32)\n",
" cuda_call(cudart.cudaHostRegister(INT8_h_output, INT8_h_output.nbytes, 0))\n",
" INT8_d_output = cuda_call(cudart.cudaMalloc(INT8_h_output.nbytes))\n",
"\n",
" # No separate INT8 stream; using context manager stream\n",
" INT8_stream = stream.stream\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "pressed-adobe",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"device = widgets.RadioButtons(\n",
" options=['CPU - Framework (PyTorch)', \n",
" 'GPU - Framework (PyTorch)', \n",
" 'GPU - TensorRT FP16',\n",
" 'GPU - TensorRT INT8'],\n",
" description='Device:',\n",
" disabled=False\n",
")\n",
"\n",
"paragraph_text = widgets.Textarea(\n",
" value='TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps'\\\n",
"'such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops'\\\n",
"'and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep'\\\n",
"'learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps.',\n",
" placeholder='Type something',\n",
" description='Passage:',\n",
" disabled=False,\n",
" layout=widgets.Layout(width=\"auto\"),\n",
" rows=10, \n",
")\n",
"\n",
"question_text = widgets.Textarea(\n",
" value='What is TensorRT?',\n",
" placeholder='Type something',\n",
" description='Question:',\n",
" disabled=False,\n",
" layout=widgets.Layout(width=\"auto\"),\n",
" rows=2,\n",
")\n",
"display(paragraph_text)\n",
"display(question_text)\n",
"\n",
"from IPython.display import display\n",
"box_layout = widgets.Layout(display='flex',\n",
" flex_flow='column',\n",
" align_items='center',\n",
" width='100%')\n",
"\n",
"button = widgets.Button(description=\"Answer Me!\")\n",
"output = widgets.Output()\n",
"box = widgets.HBox(children=[button],layout=box_layout)\n",
"\n",
"N_RUN = 10\n",
"\n",
"def answer(b):\n",
" progress_bar.value = 0\n",
" inference_time_arr = []\n",
" with output:\n",
" if device.value == 'GPU - TensorRT FP16':\n",
" output.clear_output()\n",
" for _ in range(N_RUN):\n",
" doc_tokens = dp.convert_doc_tokens(paragraph_text.value)\n",
" features = question_features(doc_tokens, question_text.value)\n",
" eval_time_elapsed, prediction, nbest_json = inference_FP16(context, d_inputs, h_output, d_output, features, doc_tokens)\n",
" progress_bar.value += 1 \n",
" inference_time_arr.append(eval_time_elapsed)\n",
"\n",
" print_single_query(eval_time_elapsed, prediction, nbest_json)\n",
" print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr))) \n",
" elif device.value == 'GPU - TensorRT INT8':\n",
" output.clear_output()\n",
" for _ in range(N_RUN):\n",
" doc_tokens = dp.convert_doc_tokens(paragraph_text.value)\n",
" features = question_features(doc_tokens, question_text.value)\n",
" eval_time_elapsed, prediction, nbest_json = inference_INT8(INT8_context, INT8_d_inputs, INT8_h_output, INT8_d_output, features, doc_tokens)\n",
" progress_bar.value += 1 \n",
" inference_time_arr.append(eval_time_elapsed)\n",
"\n",
" print_single_query(eval_time_elapsed, prediction, nbest_json)\n",
" print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr))) \n",
" \n",
" elif device.value == 'CPU - Framework (PyTorch)':\n",
" output.clear_output()\n",
" for _ in range(N_RUN):\n",
" inference_time = time.time()\n",
" answer = nlp({\n",
" 'question': question_text.value,\n",
" 'context': paragraph_text.value\n",
" })\n",
" progress_bar.value += 1 \n",
" inference_time_arr.append(time.time() - inference_time)\n",
" \n",
" print(\"Answer: '{}'\".format(answer['answer']))\n",
" print(\"With probability: {:.2f}%\".format(answer['score']*100))\n",
" print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr))) \n",
" elif device.value == 'GPU - Framework (PyTorch)': \n",
" output.clear_output()\n",
" for _ in range(N_RUN):\n",
" inference_time = time.time()\n",
" answer = nlp_gpu({\n",
" 'question': question_text.value,\n",
" 'context': paragraph_text.value\n",
" })\n",
" progress_bar.value += 1 \n",
" inference_time_arr.append(time.time() - inference_time)\n",
" \n",
" print(\"Answer: '{}'\".format(answer['answer']))\n",
" print(\"With probability: {:.2f}%\".format(answer['score']*100))\n",
" print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr))) \n",
" \n",
"button.on_click(answer)\n",
"display(device, box, output)\n",
"\n",
"progress_bar = widgets.IntProgress(\n",
" value=0,\n",
" min=0,\n",
" max=N_RUN,\n",
" description='Progress:',\n",
" bar_style='', # 'success', 'info', 'warning', 'danger' or ''\n",
" style={'bar_color': 'green'},\n",
" orientation='horizontal', \n",
" layout=widgets.Layout(width='100%', height='50px')\n",
")\n",
"display(progress_bar)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "musical-right",
"metadata": {},
"outputs": [],
"source": [
"\n",
"# free allocated memory\n",
"for d_input in d_inputs:\n",
" cuda_call(cudart.cudaFree(d_input))\n",
"cuda_call(cudart.cudaFree(d_output))\n",
"cuda_call(cudart.cudaHostUnregister(h_output))\n",
"\n",
"for INT8_d_input in INT8_d_inputs:\n",
" cuda_call(cudart.cudaFree(INT8_d_input))\n",
"cuda_call(cudart.cudaFree(INT8_d_output))\n",
"cuda_call(cudart.cudaHostUnregister(INT8_h_output))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
demo/BERT/notebooks/README.md
-11
View File
@@ -1,11 +0,0 @@
# TensorRT Demo with BERT
To run the demo Jupyter notebooks in this folder, follow the instructions in the [TRT setup guide](../../../README.md) to build and launch the docker container. Then, use your browswer to open the Jupyter lab interface at <host_name>:8888/lab using the password provided in the terminal.
Notebook list:
- [BERT-TRT-FP16.ipynb](BERT-TRT-FP16.ipynb): Step by step walkthrough for building BERT TensorRT FP16 engine.
- [BERT-TRT-INT8-QAT-sparse.ipynb](BERT-TRT-INT8-QAT-sparse.ipynb): Step by step walkthrough for building BERT TensorRT INT8 engine.
- [Q-and-A.ipynb](Q-and-A.ipynb): GUI for Q&A with BERT.
- [benchmark.ipynb](benchmark.ipynb): GUI for benchmarking TensorRT engines.
demo/BERT/notebooks/benchmark.ipynb
-488
View File
@@ -1,488 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 4,
"id": "81fd5338-bf4e-45b2-8c8c-a6e8a7da07b8",
"metadata": {},
"outputs": [],
"source": [
"# Copyright 2021 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",
"# ==============================================================================\n"
]
},
{
"cell_type": "markdown",
"id": "d4896018-e56e-464f-9edc-a10057ea45d9",
"metadata": {},
"source": [
"<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
"\n",
"# BERT TensorRT Benchmarking: FP16 vs. INT8 QAT+Sparsity\n",
"\n",
"In this notebook, we benchmark different BERT Large TensorRT engines at different batch sizes.\n",
"\n",
"**Notes**: \n",
" - Prior to running this notebook, run [BERT-TRT-FP16.ipynb](BERT-TRT-FP16.ipynb) and [BERT-TRT-INT8-QAT-sparse.ipynb](BERT-TRT-INT8-QAT-sparse.ipynb) to generate the TensorRT engines.\n",
" - This benchmarking focuses on the compute part using synthetic inputs, without taking into account pre and post processing time."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c89592ab-e50b-47c9-9b5b-6b5067d0b22a",
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": []
},
"outputs": [],
"source": [
"import warnings\n",
"warnings.filterwarnings('ignore')\n",
"\n",
"import sys\n",
"sys.path.append('/workspace/TensorRT/demo/BERT')\n",
"import os\n",
"sys.path.append(os.path.abspath('..'))\n",
"\n",
"import tensorrt as trt;\n",
"TRT_VERSION = trt.__version__\n",
"\n",
"import time\n",
"import argparse\n",
"import ctypes\n",
"import numpy as np\n",
"import tensorrt as trt\n",
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"import numpy as np\n",
"from tqdm import tqdm\n",
"\n",
"import ipywidgets as widgets\n",
"from ipywidgets import IntProgress\n",
"from ipywidgets import Button, Layout\n",
"from IPython.display import display\n",
"\n",
"import helpers.tokenization as tokenization\n",
"import helpers.data_processing as dp\n",
"\n",
"TRT_LOGGER = trt.Logger(trt.Logger.ERROR)\n",
"\n",
"\n",
"class DeviceBuffer(object):\n",
" def __init__(self, shape, dtype=trt.int32):\n",
" self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))\n",
"\n",
" def binding(self):\n",
" return int(self.buf)\n",
"\n",
" def free(self):\n",
" cuda_call(cudart.cudaFree(self.buf))\n",
"\n",
"doc_stride = 128\n",
"max_query_length = 64\n",
"max_seq_length = 384\n",
"\n",
"vocab_file = \"models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt\"\n",
"tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)\n",
"\n",
"parser = argparse.ArgumentParser(description='BERT Inference Benchmark')\n",
"parser.add_argument(\"-e\", \"--engine\", help='Path to BERT TensorRT engine', default='')\n",
"parser.add_argument('-b', '--batch-size', default=[], action=\"append\", help='Batch size(s) to benchmark. Can be specified multiple times for more than one batch size. This script assumes that the engine has been built with one optimization profile for each batch size, and that these profiles are in order of increasing batch size.', type=int)\n",
"parser.add_argument('-s', '--sequence-length', default=384, help='Sequence length of the BERT model', type=int)\n",
"parser.add_argument('-i', '--iterations', default=1000, help='Number of iterations to run when benchmarking each batch size.', type=int)\n",
"parser.add_argument('-w', '--warm-up-runs', default=10, help='Number of iterations to run prior to benchmarking.', type=int)\n",
"parser.add_argument('-r', '--random-seed', required=False, default=12345, help='Random seed.', type=int)\n",
"args, _ = parser.parse_known_args()\n",
"args.batch_size = args.batch_size or [1]\n",
"\n",
"# Import necessary plugins for BERT TensorRT\n",
"ctypes.CDLL(\"libnvinfer_plugin.so\", mode=ctypes.RTLD_GLOBAL)\n",
"\n",
"### INT8 TRT model\n",
"def run_benchmark_INT8(b):\n",
" engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
" with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:\n",
" with output:\n",
" #output.clear_output()\n",
" args.batch_size = [int(batchsize_selector.value)]\n",
"\n",
" # Allocate buffers large enough to store the largest batch size\n",
" max_input_shape = (args.sequence_length * max(args.batch_size), )\n",
" max_output_shape = (args.sequence_length * max(args.batch_size), 2, 1, 1)\n",
" buffers = [\n",
" DeviceBuffer(max_input_shape),\n",
" DeviceBuffer(max_input_shape),\n",
" DeviceBuffer((max(args.batch_size) + 1, )),\n",
" DeviceBuffer((args.sequence_length, )),\n",
" DeviceBuffer(max_output_shape)\n",
" ]\n",
"\n",
" # Prepare random input\n",
" pseudo_vocab_size = 30522\n",
" pseudo_type_vocab_size = 2\n",
" np.random.seed(args.random_seed)\n",
" test_word_ids = np.random.randint(0, pseudo_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)\n",
" test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)\n",
" test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)\n",
"\n",
" # Copy input h2d\n",
" memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n",
" memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n",
" memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel())\n",
"\n",
" bench_times = {}\n",
" with CudaStreamContext() as stream:\n",
"\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors-1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.stream)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
"\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
"\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda_call(cudart.cudaEventCreate())\n",
" end = cuda_call(cudart.cudaEventCreate())\n",
"\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" cuda_call(cudart.cudaEventRecord(start, stream.stream))\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" cuda_call(cudart.cudaEventRecord(end, stream.stream))\n",
" stream.synchronize()\n",
" elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))\n",
" times.append(elapsed_time)\n",
" progress_bar.value +=1\n",
"\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
"\n",
" [b.free() for b in buffers]\n",
"\n",
" for batch_size, times in bench_times.items():\n",
" total_time = sum(times)\n",
" avg_time = total_time / float(len(times))\n",
" times.sort()\n",
" percentile95 = times[int(len(times) * 0.95)]\n",
" percentile99 = times[int(len(times) * 0.99)]\n",
" print(\"BERT TRT INT8: Running {:} iterations with Batch Size: {:}\\n\\tTotal Time: {:.2f} ms \\tAverage Time: {:.2f} ms\\t95th Percentile Time: {:.2f} ms\\t99th Percentile Time: {:.2f}\".format(args.iterations, batch_size, total_time, avg_time, percentile95, percentile99))\n",
"\n",
"\n",
"### FP16 TRT model\n",
"def run_benchmark_FP16(b):\n",
" engine_path = \"engines_%s/bert_large_384.engine\"%TRT_VERSION\n",
" with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:\n",
" with output:\n",
" #output.clear_output()\n",
" args.batch_size = [int(batchsize_selector.value)]\n",
"\n",
" # Allocate buffers large enough to store the largest batch size\n",
" max_input_shape = (max(args.batch_size), args.sequence_length)\n",
" max_output_shape = (max(args.batch_size), args.sequence_length, 2, 1, 1)\n",
" buffers = [\n",
" DeviceBuffer(max_input_shape),\n",
" DeviceBuffer(max_input_shape),\n",
" DeviceBuffer(max_input_shape),\n",
" DeviceBuffer(max_output_shape)\n",
" ]\n",
"\n",
" # Prepare random input\n",
" pseudo_vocab_size = 30522\n",
" pseudo_type_vocab_size = 2\n",
" np.random.seed(args.random_seed)\n",
" test_word_ids = np.random.randint(0, pseudo_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
" test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
" test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
"\n",
" # Copy input h2d\n",
" memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n",
" memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n",
" memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel())\n",
"\n",
" bench_times = {}\n",
" with CudaStreamContext() as stream:\n",
"\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors - 1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.stream)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
"\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
"\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda_call(cudart.cudaEventCreate())\n",
" end = cuda_call(cudart.cudaEventCreate())\n",
"\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" start.record(stream.stream)\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" end.record(stream.stream)\n",
" stream.synchronize()\n",
" times.append(end.time_since(start))\n",
" progress_bar.value +=1\n",
"\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
"\n",
" [b.free() for b in buffers]\n",
"\n",
" for batch_size, times in bench_times.items():\n",
" total_time = sum(times)\n",
" avg_time = total_time / float(len(times))\n",
" times.sort()\n",
" percentile95 = times[int(len(times) * 0.95)]\n",
" percentile99 = times[int(len(times) * 0.99)]\n",
" print(\"BERT TRT FP16: Running {:} iterations with Batch Size: {:}\\n\\tTotal Time: {:.2f} ms \\tAverage Time: {:.2f} ms\\t95th Percentile Time: {:.2f} ms\\t99th Percentile Time: {:.2f}\".format(args.iterations, batch_size, total_time, avg_time, percentile95, percentile99))\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "c72b6e9d-3181-4c22-8161-c1a1a9444e2d",
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": []
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "fcd71564172944a198f873066199c7e1",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"RadioButtons(description='Engine:', options=('GPU - TensorRT FP16', 'GPU - TensorRT INT8'), value='GPU - Tenso…"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8b51334850f84db4834a1d0035556971",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"RadioButtons(description='Batch size:', options=('1', '32', '64', '128'), value='1')"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "44b1a0bd02834393b7d7204c18395578",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"IntSlider(value=500, continuous_update=False, description='Iterations:', max=1000, min=100)"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "f988290aab5243fb9d43a8581d104d64",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(Button(description='Run benchmark', style=ButtonStyle()),), layout=Layout(align_items='center',…"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "03b48cb93c834fe89f2180dc1562a4c7",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Output()"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "3a3756b664bc482b843283c2e6cff3fa",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"IntProgress(value=0, description='Progress:', layout=Layout(height='50px', width='100%'), max=1000, style=Prog…"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# UI elements\n",
"engine_selector = widgets.RadioButtons(\n",
" options=['GPU - TensorRT FP16',\n",
" 'GPU - TensorRT INT8'],\n",
" description='Engine:',\n",
" disabled=False\n",
")\n",
"\n",
"batchsize_selector = widgets.RadioButtons(\n",
" options=['1', '32', '64', '128'],\n",
" description='Batch size:',\n",
" disabled=False\n",
")\n",
"\n",
"iteration_selector = widgets.IntSlider(\n",
" value=500,\n",
" min=100,\n",
" max=1000,\n",
" step=1,\n",
" description='Iterations:',\n",
" disabled=False,\n",
" continuous_update=False,\n",
" orientation='horizontal',\n",
" readout=True,\n",
" readout_format='d'\n",
")\n",
"\n",
"button = widgets.Button(description=\"Run benchmark\")\n",
"output = widgets.Output()\n",
"#box = widgets.HBox(children=[button],layout=box_layout)\n",
"\n",
"def run_benchmark(b):\n",
" args.iterations = iteration_selector.value\n",
" progress_bar.max = iteration_selector.value\n",
" with output:\n",
" if engine_selector.value=='GPU - TensorRT FP16':\n",
" run_benchmark_FP16(b)\n",
" elif engine_selector.value=='GPU - TensorRT INT8':\n",
" run_benchmark_INT8(b)\n",
" \n",
"button.on_click(run_benchmark)\n",
"display(engine_selector, batchsize_selector, iteration_selector)\n",
"\n",
"from IPython.display import display\n",
"box_layout = widgets.Layout(display='flex',\n",
" flex_flow='column',\n",
" align_items='center',\n",
" width='100%')\n",
"box = widgets.HBox(children=[button],layout=box_layout)\n",
"display(box, output)\n",
"\n",
"progress_bar = widgets.IntProgress(\n",
" value=0,\n",
" min=0,\n",
" max=1000,\n",
" description='Progress:',\n",
" bar_style='',\n",
" style={'bar_color': 'green'},\n",
" orientation='horizontal', \n",
" layout=Layout(width='100%', height='50px')\n",
")\n",
"display(progress_bar)\n",
" "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1317f3b9-6bd0-4f7b-9956-b33e5948a898",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
demo/BERT/notebooks/sparsity-diagram-600x338-r3.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 13 KiB

demo/BERT/notebooks/structured_spare_matrix.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 64 KiB

demo/BERT/perf.py
-143
View File
@@ -1,143 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse
import ctypes
import time
import tensorrt as trt
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class DeviceBuffer(object):
def __init__(self, shape, dtype=trt.int32):
self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))
def binding(self):
return int(self.buf)
def free(self):
cuda_call(cudart.cudaFree(self.buf))
def main():
parser = argparse.ArgumentParser(description='BERT Inference Benchmark')
parser.add_argument("-e", "--engine", help='Path to BERT TensorRT engine')
parser.add_argument('-b', '--batch-size', default=[], action="append", help='Batch size(s) to benchmark. Can be specified multiple times for more than one batch size. This script assumes that the engine has been built with one optimization profile for each batch size, and that these profiles are in order of increasing batch size.', type=int)
parser.add_argument('-s', '--sequence-length', default=128, help='Sequence length of the BERT model', type=int)
parser.add_argument('-i', '--iterations', default=200, help='Number of iterations to run when benchmarking each batch size.', type=int)
parser.add_argument('-w', '--warm-up-runs', default=10, help='Number of iterations to run prior to benchmarking.', type=int)
parser.add_argument('-d', '--duration', default=0.0, help='Minimal number of seconds to run when benchmarking each batch size.', type=float)
parser.add_argument('-r', '--random-seed', required=False, default=12345, help='Random seed.', type=int)
args, _ = parser.parse_known_args()
args.batch_size = args.batch_size or [1]
# Import necessary plugins for BERT TensorRT
ctypes.CDLL("libnvinfer_plugin.so", mode=ctypes.RTLD_GLOBAL)
with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
# Allocate buffers large enough to store the largest batch size
max_input_shape = (max(args.batch_size), args.sequence_length)
max_output_shape = (max(args.batch_size), args.sequence_length, 2, 1, 1)
buffers = [
DeviceBuffer(max_input_shape),
DeviceBuffer(max_input_shape),
DeviceBuffer(max_input_shape),
DeviceBuffer(max_output_shape)
]
# Prepare random input
pseudo_vocab_size = 30522
pseudo_type_vocab_size = 2
np.random.seed(args.random_seed)
test_word_ids = np.random.randint(0, pseudo_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)
test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)
test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)
# Copy input h2d
memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())
memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())
memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel())
bench_times = {}
with CudaStreamContext() as stream:
for batch_size in sorted(args.batch_size):
# Select engine profile
selected_profile = -1
for idx in range(engine.num_optimization_profiles):
profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx)
if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size and profile_shape[0][1] <= args.sequence_length and profile_shape[2][1] >= args.sequence_length:
selected_profile = idx
break
if selected_profile == -1:
raise RuntimeError("None of the dynamic shape profiles meets the requirement batch = {} and sequence = {}.".format(batch_size, args.sequence_length))
context.set_optimization_profile_async(selected_profile, stream.stream)
# Each profile has unique bindings
binding_idx_offset = selected_profile * engine.num_io_tensors
bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]
input_shape = (batch_size, args.sequence_length)
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
# Inference
total_time = 0
start = cuda_call(cudart.cudaEventCreate())
end = cuda_call(cudart.cudaEventCreate())
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.stream)
stream.synchronize()
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
cuda_call(cudart.cudaEventRecord(start, stream.stream))
context.execute_async_v3(stream_handle=stream.stream)
cuda_call(cudart.cudaEventRecord(end, stream.stream))
stream.synchronize()
elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))
times.append(elapsed_time)
actual_iterations += 1
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
[b.free() for b in buffers]
for batch_size, times in bench_times.items():
total_time = sum(times)
avg_time = total_time / float(actual_iterations)
times.sort()
percentile95 = times[int(actual_iterations * 0.95)]
percentile99 = times[int(actual_iterations * 0.99)]
print("Running {:} iterations with Batch Size: {:}\n\tTotal Time: {:} ms \tAverage Time: {:} ms\t95th Percentile Time: {:} ms\t99th Percentile Time: {:}".format(actual_iterations, batch_size, total_time, avg_time, percentile95, percentile99))
if __name__ == '__main__':
main()
demo/BERT/perf_varseqlen.py
-142
View File
@@ -1,142 +0,0 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse
import ctypes
import time
import numpy as np
import tensorrt as trt
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class DeviceBuffer(object):
def __init__(self, shape, dtype=trt.int32):
self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))
def binding(self):
return int(self.buf)
def free(self):
cuda_call(cudart.cudaFree(self.buf))
def main():
parser = argparse.ArgumentParser(description='BERT Inference Benchmark')
parser.add_argument("-e", "--engine", help='Path to BERT TensorRT engine')
parser.add_argument('-b', '--batch-size', default=[], action="append", help='Batch size(s) to benchmark. Can be specified multiple times for more than one batch size. This script assumes that the engine has been built with one optimization profile for each batch size, and that these profiles are in order of increasing batch size.', type=int)
parser.add_argument('-s', '--sequence-length', default=128, help='Sequence length of the BERT model', type=int)
parser.add_argument('-i', '--iterations', default=200, help='Number of iterations to run when benchmarking each batch size.', type=int)
parser.add_argument('-w', '--warm-up-runs', default=10, help='Number of iterations to run prior to benchmarking.', type=int)
parser.add_argument('-d', '--duration', default=0.0, help='Minimal number of seconds to run when benchmarking each batch size.', type=float)
parser.add_argument('-r', '--random-seed', required=False, default=12345, help='Random seed.', type=int)
args, _ = parser.parse_known_args()
args.batch_size = args.batch_size or [1]
# Import necessary plugins for BERT TensorRT
ctypes.CDLL("libnvinfer_plugin.so", mode=ctypes.RTLD_GLOBAL)
with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
# Allocate buffers large enough to store the largest batch size
max_input_shape = (args.sequence_length * max(args.batch_size), )
max_output_shape = (args.sequence_length * max(args.batch_size), 2, 1, 1)
buffers = [
DeviceBuffer(max_input_shape),
DeviceBuffer(max_input_shape),
DeviceBuffer((max(args.batch_size) + 1, )),
DeviceBuffer((args.sequence_length, )),
DeviceBuffer(max_output_shape)
]
# Prepare random input
pseudo_vocab_size = 30522
pseudo_type_vocab_size = 2
np.random.seed(args.random_seed)
test_word_ids = np.random.randint(0, pseudo_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)
test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)
test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)
# Copy input h2d
memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())
memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())
memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel())
bench_times = {}
for idx, batch_size in enumerate(sorted(args.batch_size)):
with CudaStreamContext() as stream:
context.set_optimization_profile_async(0, stream.stream)
# Each profile has unique bindings
bindings = [buf.binding() for buf in buffers]
shapes = {
"input_ids": (args.sequence_length * batch_size, ),
"segment_ids": (args.sequence_length * batch_size, ),
"cu_seqlens": (batch_size + 1, ),
"max_seqlen": (args.sequence_length, ),
}
for binding, shape in shapes.items():
context.set_input_shape(binding, shape)
assert len(context.infer_shapes()) == 0
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
# Inference
total_time = 0
start = cuda_call(cudart.cudaEventCreate())
end = cuda_call(cudart.cudaEventCreate())
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.stream)
stream.synchronize()
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
cuda_call(cudart.cudaEventRecord(start, stream.stream))
context.execute_async_v3(stream_handle=stream.stream)
cuda_call(cudart.cudaEventRecord(end, stream.stream))
stream.synchronize()
elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))
times.append(elapsed_time)
actual_iterations += 1
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
[b.free() for b in buffers]
for batch_size, times in bench_times.items():
total_time = sum(times)
avg_time = total_time / float(actual_iterations)
times.sort()
percentile95 = times[int(actual_iterations * 0.95)]
percentile99 = times[int(actual_iterations * 0.99)]
print("Running {:} iterations with Batch Size: {:}\n\tTotal Time: {:} ms \tAverage Time: {:} ms\t95th Percentile Time: {:} ms\t99th Percentile Time: {:}".format(actual_iterations, batch_size, total_time, avg_time, percentile95, percentile99))
if __name__ == '__main__':
main()
demo/BERT/scripts/download_model.sh
-92
View File
@@ -1,92 +0,0 @@
#!/bin/bash
#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Setup default parameters (if no command-line parameters given)
SQUAD='2'
MODEL='large'
SEQ_LEN='128'
FW='tf'
WTYPE='dense'
PREC='fp16'
while test $# -gt 0
do
case "$1" in
-h) echo "Usage: sh download_model.sh [tf|pyt] [base|large|megatron-large] [128|384] [v2|v1_1] [sparse] [int8-qat]"
exit 0
;;
base) MODEL='base'
;;
large) MODEL='large'
;;
megatron-large) MODEL='megatron'
;;
128) SEQ_LEN='128'
;;
384) SEQ_LEN='384'
;;
v2) SQUAD='2'
;;
v1_1) SQUAD='11'
;;
tf) FW='tf'
;;
pyt) FW='pyt'
;;
int8-qat) PREC='int8qat'
;;
sparse) WTYPE='sparse'
;;
*) echo "Invalid argument $1...exiting"
exit 0
;;
esac
shift
done
# Prepare the download directory
mkdir -p models/fine-tuned
pushd models/fine-tuned
# Download the BERT fine-tuned model
echo "Downloading BERT-${FW} ${MODEL} checkpoints for sequence length ${SEQ_LEN} and fine-tuned for SQuAD ${SQUAD}."
if [ "${FW}" = 'tf' ]; then
CKPT=bert_${FW}_ckpt_${MODEL}_qa_squad${SQUAD}_amp_${SEQ_LEN}
CKPT_VERSION=19.03.1
elif [ "${FW}" = 'pyt' ]; then
if [ "${MODEL}" == 'megatron' ]; then
CKPT=bert_${FW}_statedict_megatron_${WTYPE}_${PREC}
CKPT_VERSION=21.03.0
elif [ "${MODEL}" != 'large' ] || [ "${SQUAD}" != '11' ]; then
echo "ERROR: Only BERT-large checkpoint fine-tuned for SQuAD v1.1 available in the QAT (PyTorch) workflow."
else
CKPT=bert_${FW}_onnx_${MODEL}_qa_squad${SQUAD}_amp_fake_quant
CKPT_VERSION=1
fi
else
echo "Invalid framework specified for checkpoint. Run download_model.sh -h for help."
fi
if [ -n "$CKPT" ]; then
if [ -d "${CKPT}_v${CKPT_VERSION}" ]; then
echo "Checkpoint directory ${PWD}/${CKPT}_v${CKPT_VERSION} already exists. Skip download."
else
ngc registry model download-version nvidia/${CKPT}:${CKPT_VERSION}
fi
fi
popd
demo/BERT/scripts/download_squad.sh
-44
View File
@@ -1,44 +0,0 @@
#!/bin/bash
#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Setup default parameters (if no command-line parameters given)
VERSION='v1.1'
while test $# -gt 0
do
case "$1" in
-h) echo "Usage: sh download_squad.sh [v2_0|v1_1]"
exit 0
;;
v2_0) VERSION='v2.0'
;;
v1_1) VERSION='v1.1'
;;
*) echo "Invalid argument $1...exiting"
exit 0
;;
esac
shift
done
# Download the SQuAD training and dev datasets
echo "Downloading SQuAD-${VERSION} training and dev datasets"
mkdir -p squad
pushd squad
wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-${VERSION}.json
wget https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-${VERSION}.json
popd
demo/BERT/scripts/inference_benchmark.sh
-280
View File
@@ -1,280 +0,0 @@
#!/bin/bash
#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Usage: run_benchmark(batch_sizes, model_variant: (base/large), precision: (int8/int8-qat/fp16/fp32), sequence_length, max_batch_size, gpu_arch)
run_benchmark() {
BATCH_SIZES="${1}"
MODEL_VARIANT="${2}"
PRECISION="${3}"
SEQUENCE_LENGTH="${4}"
MAX_BATCH="${5}"
GPU_ARCH="${6}"
CHECKPOINTS_DIR="models/fine-tuned/bert_tf_ckpt_${MODEL_VARIANT}_qa_squad2_amp_${SEQUENCE_LENGTH}_v19.03.1"
SQUAD_DIR="BERT/squad"
ENGINE_NAME="engines/bert_${MODEL_VARIANT}_${PRECISION}_bs${MAX_BATCH}_seqlen${SEQUENCE_LENGTH}_benchmark.engine"
# QAT Checkpoint - available only for BERT-Large
QAT_CHECKPOINT="models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx"
CUDAGRAPH_PERFBIN="build/perf"
TIMING_CACHE_FILE="build.tcf"
echo "==== Benchmarking BERT ${MODEL_VARIANT} ${PRECISION} SEQLEN ${SEQUENCE_LENGTH} on ${GPU_ARCH} ===="
if [ ! -f ${ENGINE_NAME} ]; then
if [ ! -d ${CHECKPOINTS_DIR} ]; then
echo "Downloading checkpoints: scripts/download_model.sh ${MODEL_VARIANT} ${SEQUENCE_LENGTH}"
scripts/download_model.sh "${MODEL_VARIANT}" "${SEQUENCE_LENGTH}"
fi;
if [ "${PRECISION}" == "int8-qat" ]; then
if [ ${MODEL_VARIANT} != "large" ]; then
echo "Skipping: BERT-base not supported for int8 (QAT)"
return
fi;
if [ ! -f ${QAT_CHECKPOINT} ]; then
echo "Downloading QAT checkpoint: scripts/download_model.sh pyt v1_1 ${MODEL_VARIANT}"
scripts/download_model.sh pyt v1_1 "${MODEL_VARIANT}"
fi;
PRECISION="int8"
BUILDER_ARGS="-x ${QAT_CHECKPOINT}"
else
BUILDER_ARGS="-m ${CHECKPOINTS_DIR}/model.ckpt"
fi;
BUILDER_ARGS="${BUILDER_ARGS} -tcf ${TIMING_CACHE_FILE} -o ${ENGINE_NAME} ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -c ${CHECKPOINTS_DIR} -v ${CHECKPOINTS_DIR}/vocab.txt --${PRECISION}"
if [ "${PRECISION}" == "int8" ]; then
BUILDER_ARGS="${BUILDER_ARGS} --fp16 --strict --calib-num 1"
if [ "${GPU_ARCH}" == "Ampere" ] || [ "${GPU_ARCH}" == "Turing" ]; then
BUILDER_ARGS="${BUILDER_ARGS} -iln -imh"
elif [ "${GPU_ARCH}" == "Xavier" ]; then
BUILDER_ARGS="${BUILDER_ARGS} -iln"
fi;
fi;
echo "Building engine: python3 builder.py ${BUILDER_ARGS}"
python3 builder.py ${BUILDER_ARGS}
fi;
if [ "${GPU_ARCH}" == "Ampere" ]; then
# Use more iterations for faster GPUs
NUM_ITERATIONS=2000
else
NUM_ITERATIONS=1000
fi;
if [ -f ${CUDAGRAPH_PERFBIN} ]; then
echo "Running benchmark with CUDA graph acceleration: perf ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS} --enable_graph"
${CUDAGRAPH_PERFBIN} ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS} --enable_graph
else
echo "Running benchmark: perf.py ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS}"
python3 perf.py ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS}
fi;
echo
}
arg_gpu="Volta"
arg_help=0
while [[ "$#" -gt 0 ]]; do case $1 in
--gpu) arg_gpu="$2"; shift;;
-h|--help) arg_help=1;;
*) echo "Unknown parameter passed: $1"; echo "For help type: $0 --help"; exit 1;
esac; shift; done
if [ "$arg_help" -eq "1" ]; then
echo "Usage: $0 [options]"
echo " --help or -h : Print this help menu."
echo " --gpu <arch> : GPU arch. Options: 'Volta', 'Xavier', 'Turing', 'Ampere'"
exit;
fi
mkdir -p /workspace/TensorRT/demo/BERT/engines
nvidia-smi -q
# BERT BASE
## INT8
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "int8" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "int8" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "int8" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "int8" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "int8" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "int8" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "int8" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "int8" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "int8" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "int8" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "int8" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "int8" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "int8" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "int8" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "int8" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "int8" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "int8" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "int8" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "int8" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "int8" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "int8" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "int8" "384" "128" "${arg_gpu}"
## FP16
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "fp16" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "fp16" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "fp16" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "fp16" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "fp16" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "fp16" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "fp16" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "fp16" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "fp16" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "fp16" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "fp16" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "fp16" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "fp16" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "fp16" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "fp16" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "fp16" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "fp16" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "fp16" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "fp16" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "fp16" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "fp16" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "fp16" "384" "128" "${arg_gpu}"
## FP32
#run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "fp32" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "fp32" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "fp32" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "fp32" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "fp32" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "fp32" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "fp32" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "fp32" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "fp32" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "fp32" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "fp32" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "base" "fp32" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "base" "fp32" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "base" "fp32" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "base" "fp32" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "base" "fp32" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "base" "fp32" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "base" "fp32" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "base" "fp32" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "base" "fp32" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "base" "fp32" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "base" "fp32" "384" "128" "${arg_gpu}"
# BERT LARGE
## INT8-QAT
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "int8-qat" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "int8-qat" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "int8-qat" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "int8-qat" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "int8-qat" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "int8-qat" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "int8-qat" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "int8-qat" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "int8-qat" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "int8-qat" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "int8-qat" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "int8-qat" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "int8-qat" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "int8-qat" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "int8-qat" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "int8-qat" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "int8-qat" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "int8-qat" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "int8-qat" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "int8-qat" "384" "128" "${arg_gpu}"
## INT8
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "int8" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "int8" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "int8" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "int8" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "int8" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "int8" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "int8" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "int8" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "int8" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "int8" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "int8" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "int8" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "int8" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "int8" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "int8" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "int8" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "int8" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "int8" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "int8" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "int8" "384" "128" "${arg_gpu}"
## FP16
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "fp16" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "fp16" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "fp16" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "fp16" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "fp16" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "fp16" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "fp16" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "fp16" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "fp16" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "fp16" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "fp16" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "fp16" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "fp16" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "fp16" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "fp16" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "fp16" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "fp16" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "fp16" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "fp16" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "fp16" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "fp16" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "fp16" "384" "128" "${arg_gpu}"
## FP32
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "fp32" "128" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "fp32" "128" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "fp32" "128" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "fp32" "128" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "fp32" "128" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "fp32" "128" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "fp32" "128" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "fp32" "128" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "fp32" "128" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "fp32" "128" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "fp32" "128" "128" "${arg_gpu}"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "fp32" "384" "32" "${arg_gpu}"
run_benchmark "-b 1" "large" "fp32" "384" "1" "${arg_gpu}"
run_benchmark "-b 2" "large" "fp32" "384" "2" "${arg_gpu}"
run_benchmark "-b 4" "large" "fp32" "384" "4" "${arg_gpu}"
run_benchmark "-b 8" "large" "fp32" "384" "8" "${arg_gpu}"
run_benchmark "-b 12" "large" "fp32" "384" "12" "${arg_gpu}"
run_benchmark "-b 16" "large" "fp32" "384" "16" "${arg_gpu}"
run_benchmark "-b 24" "large" "fp32" "384" "24" "${arg_gpu}"
run_benchmark "-b 32" "large" "fp32" "384" "32" "${arg_gpu}"
run_benchmark "-b 64" "large" "fp32" "384" "64" "${arg_gpu}"
run_benchmark "-b 128" "large" "fp32" "384" "128" "${arg_gpu}"
demo/BERT/scripts/inference_benchmark_megatron.sh
-164
View File
@@ -1,164 +0,0 @@
#!/bin/bash
#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Usage: run_benchmark(batch_sizes, model_variant: (large), precision: (int8-qat), sequence_length, max_batch_size, gpu_arch: Ampere, weights: (dense/sparse))
run_benchmark() {
BATCH_SIZES="${1}"
MODEL_VARIANT="${2}"
PRECISION="${3}"
SEQUENCE_LENGTH="${4}"
MAX_BATCH="${5}"
GPU_ARCH="${6}"
WEIGHTS="${7}"
CHECKPOINTS_DIR="models/fine-tuned/bert_tf_ckpt_${MODEL_VARIANT}_qa_squad2_amp_${SEQUENCE_LENGTH}_v19.03.1"
SQUAD_DIR="BERT/squad"
ENGINE_NAME="engines/bert_mt_${MODEL_VARIANT}_${PRECISION}_bs${MAX_BATCH}_seqlen${SEQUENCE_LENGTH}_weight${WEIGHTS}_benchmark.engine"
QAT_CHECKPOINT="models/fine-tuned/bert_pyt_statedict_megatron_sparse_int8qat_v21.03.0/bert_pyt_statedict_megatron_sparse_int8_qat"
CUDAGRAPH_PERFBIN="build/perf"
TIMING_CACHE_FILE="build_megatron.tcf"
if [ "${PRECISION}" != "int8-qat" ]; then
echo "Skipping: Megatron-BERT only supported for int8 (QAT)"
return
elif [ "${GPU_ARCH}" != "Ampere" ]; then
echo "Skipping: Sparsity only supported on Ampere GPUs"
return
fi
echo "==== Benchmarking BERT ${MODEL_VARIANT} ${PRECISION} SEQLEN ${SEQUENCE_LENGTH} on ${GPU_ARCH} ===="
if [ ! -f ${ENGINE_NAME} ]; then
if [ ! -d ${CHECKPOINTS_DIR} ]; then
echo "Downloading checkpoints: scripts/download_model.sh ${MODEL_VARIANT} ${SEQUENCE_LENGTH}"
scripts/download_model.sh "${MODEL_VARIANT}" "${SEQUENCE_LENGTH}"
fi;
if [ "${PRECISION}" == "int8-qat" ]; then
if [ ${MODEL_VARIANT} != "large" ]; then
echo "Skipping: Megatron-BERT-base not supported for int8 (QAT)"
return
fi;
if [ ! -f ${QAT_CHECKPOINT} ]; then
echo "Downloading QAT checkpoint: scripts/download_model.sh pyt megatron-${MODEL_VARIANT} ${PRECISION} sparse"
scripts/download_model.sh pyt megatron-${MODEL_VARIANT} ${PRECISION} sparse
fi;
PRECISION="int8"
BUILDER_ARGS="--pickle ${QAT_CHECKPOINT}"
fi;
BUILDER_ARGS="${BUILDER_ARGS} -tcf ${TIMING_CACHE_FILE} -o ${ENGINE_NAME} ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -c ${CHECKPOINTS_DIR} -v ${CHECKPOINTS_DIR}/vocab.txt --megatron"
if [ "${WEIGHTS}" == "sparse" ]; then
BUILDER_ARGS="${BUILDER_ARGS} --sp"
fi;
if [ "${PRECISION}" == "int8" ]; then
BUILDER_ARGS="${BUILDER_ARGS} --fp16 --int8 --strict"
if [ "${GPU_ARCH}" == "Ampere" ]; then
BUILDER_ARGS="${BUILDER_ARGS} -il"
fi;
fi;
echo "Building engine: python3 builder_varseqlen.py ${BUILDER_ARGS}"
python3 builder_varseqlen.py ${BUILDER_ARGS}
fi;
if [ "${GPU_ARCH}" == "Ampere" ]; then
# Use more iterations for faster GPUs
NUM_ITERATIONS=2000
else
NUM_ITERATIONS=1000
fi;
if [ -f ${CUDAGRAPH_PERFBIN} ]; then
echo "Running benchmark with CUDA graph acceleration: perf ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS} --enable_graph"
${CUDAGRAPH_PERFBIN} ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS} --enable_graph
else
echo "Running benchmark: perf.py ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS}"
python3 perf_varseqlen.py ${BATCH_SIZES} -s ${SEQUENCE_LENGTH} -e ${ENGINE_NAME} -w 100 -i ${NUM_ITERATIONS}
fi;
echo
}
arg_gpu="Ampere"
arg_help=0
while [[ "$#" -gt 0 ]]; do case $1 in
--gpu) arg_gpu="$2"; shift;;
-h|--help) arg_help=1;;
*) echo "Unknown parameter passed: $1"; echo "For help type: $0 --help"; exit 1;
esac; shift; done
if [ "$arg_help" -eq "1" ]; then
echo "Usage: $0 [options]"
echo " --help or -h : Print this help menu."
echo " --gpu <arch> : GPU arch. Options: 'Volta', 'Xavier', 'Turing', 'Ampere'"
exit;
fi;
mkdir -p /workspace/TensorRT/demo/BERT/engines
nvidia-smi -q
# BERT LARGE
## INT8-QAT (dense)
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "128" "32" "${arg_gpu}" "dense"
run_benchmark "-b 1" "large" "int8-qat" "128" "1" "${arg_gpu}" "dense"
run_benchmark "-b 2" "large" "int8-qat" "128" "2" "${arg_gpu}" "dense"
run_benchmark "-b 4" "large" "int8-qat" "128" "4" "${arg_gpu}" "dense"
run_benchmark "-b 8" "large" "int8-qat" "128" "8" "${arg_gpu}" "dense"
run_benchmark "-b 12" "large" "int8-qat" "128" "12" "${arg_gpu}" "dense"
run_benchmark "-b 16" "large" "int8-qat" "128" "16" "${arg_gpu}" "dense"
run_benchmark "-b 24" "large" "int8-qat" "128" "24" "${arg_gpu}" "dense"
run_benchmark "-b 32" "large" "int8-qat" "128" "32" "${arg_gpu}" "dense"
run_benchmark "-b 64" "large" "int8-qat" "128" "64" "${arg_gpu}" "dense"
run_benchmark "-b 128" "large" "int8-qat" "128" "128" "${arg_gpu}" "dense"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "384" "32" "${arg_gpu}" "dense"
run_benchmark "-b 1" "large" "int8-qat" "384" "1" "${arg_gpu}" "dense"
run_benchmark "-b 2" "large" "int8-qat" "384" "2" "${arg_gpu}" "dense"
run_benchmark "-b 4" "large" "int8-qat" "384" "4" "${arg_gpu}" "dense"
run_benchmark "-b 8" "large" "int8-qat" "384" "8" "${arg_gpu}" "dense"
run_benchmark "-b 12" "large" "int8-qat" "384" "12" "${arg_gpu}" "dense"
run_benchmark "-b 16" "large" "int8-qat" "384" "16" "${arg_gpu}" "dense"
run_benchmark "-b 24" "large" "int8-qat" "384" "24" "${arg_gpu}" "dense"
run_benchmark "-b 32" "large" "int8-qat" "384" "32" "${arg_gpu}" "dense"
run_benchmark "-b 64" "large" "int8-qat" "384" "64" "${arg_gpu}" "dense"
run_benchmark "-b 128" "large" "int8-qat" "384" "128" "${arg_gpu}" "dense"
## INT8-QAT (sparse)
if [ "${arg_gpu}" == "Ampere" ]; then
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "128" "32" "${arg_gpu}" "sparse"
run_benchmark "-b 1" "large" "int8-qat" "128" "1" "${arg_gpu}" "sparse"
run_benchmark "-b 2" "large" "int8-qat" "128" "2" "${arg_gpu}" "sparse"
run_benchmark "-b 4" "large" "int8-qat" "128" "4" "${arg_gpu}" "sparse"
run_benchmark "-b 8" "large" "int8-qat" "128" "8" "${arg_gpu}" "sparse"
run_benchmark "-b 12" "large" "int8-qat" "128" "12" "${arg_gpu}" "sparse"
run_benchmark "-b 16" "large" "int8-qat" "128" "16" "${arg_gpu}" "sparse"
run_benchmark "-b 24" "large" "int8-qat" "128" "24" "${arg_gpu}" "sparse"
run_benchmark "-b 32" "large" "int8-qat" "128" "32" "${arg_gpu}" "sparse"
run_benchmark "-b 64" "large" "int8-qat" "128" "64" "${arg_gpu}" "sparse"
run_benchmark "-b 128" "large" "int8-qat" "128" "128" "${arg_gpu}" "sparse"
# run_benchmark "-b 1 -b 2 -b 4 -b 8 -b 12 -b 16 -b 24 -b 32" "large" "int8-qat" "384" "32" "${arg_gpu}" "sparse"
run_benchmark "-b 1" "large" "int8-qat" "384" "1" "${arg_gpu}" "sparse"
run_benchmark "-b 2" "large" "int8-qat" "384" "2" "${arg_gpu}" "sparse"
run_benchmark "-b 4" "large" "int8-qat" "384" "4" "${arg_gpu}" "sparse"
run_benchmark "-b 8" "large" "int8-qat" "384" "8" "${arg_gpu}" "sparse"
run_benchmark "-b 12" "large" "int8-qat" "384" "12" "${arg_gpu}" "sparse"
run_benchmark "-b 16" "large" "int8-qat" "384" "16" "${arg_gpu}" "sparse"
run_benchmark "-b 24" "large" "int8-qat" "384" "24" "${arg_gpu}" "sparse"
run_benchmark "-b 32" "large" "int8-qat" "384" "32" "${arg_gpu}" "sparse"
run_benchmark "-b 64" "large" "int8-qat" "384" "64" "${arg_gpu}" "sparse"
run_benchmark "-b 128" "large" "int8-qat" "384" "128" "${arg_gpu}" "sparse"
else
echo "Sparsity only supported on Ampere GPUs. Skip benchmark."
fi;
demo/BERT/squad/evaluate-v1.1.py
-116
View File
@@ -1,116 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Obtained from https://rajpurkar.github.io/SQuAD-explorer/
""" Official evaluation script for v1.1 of the SQuAD dataset. """
from __future__ import print_function
from collections import Counter
import string
import re
import argparse
import json
import sys
def normalize_answer(s):
"""Lower text and remove punctuation, articles and extra whitespace."""
def remove_articles(text):
return re.sub(r'\b(a|an|the)\b', ' ', text)
def white_space_fix(text):
return ' '.join(text.split())
def remove_punc(text):
exclude = set(string.punctuation)
return ''.join(ch for ch in text if ch not in exclude)
def lower(text):
return text.lower()
return white_space_fix(remove_articles(remove_punc(lower(s))))
def f1_score(prediction, ground_truth):
prediction_tokens = normalize_answer(prediction).split()
ground_truth_tokens = normalize_answer(ground_truth).split()
common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
num_same = sum(common.values())
if num_same == 0:
return 0
precision = 1.0 * num_same / len(prediction_tokens)
recall = 1.0 * num_same / len(ground_truth_tokens)
f1 = (2 * precision * recall) / (precision + recall)
return f1
def exact_match_score(prediction, ground_truth):
return (normalize_answer(prediction) == normalize_answer(ground_truth))
def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
scores_for_ground_truths = []
for ground_truth in ground_truths:
score = metric_fn(prediction, ground_truth)
scores_for_ground_truths.append(score)
return max(scores_for_ground_truths)
def evaluate(dataset, predictions, f1_acc):
f1 = exact_match = total = 0
for article in dataset:
for paragraph in article['paragraphs']:
for qa in paragraph['qas']:
total += 1
if qa['id'] not in predictions:
message = 'Unanswered question ' + qa['id'] + \
' will receive score 0.'
print(message, file=sys.stderr)
continue
ground_truths = list(map(lambda x: x['text'], qa['answers']))
prediction = predictions[qa['id']]
exact_match += metric_max_over_ground_truths(
exact_match_score, prediction, ground_truths)
f1 += metric_max_over_ground_truths(
f1_score, prediction, ground_truths)
exact_match = 100.0 * exact_match / total
f1 = 100.0 * f1 / total
if (f1 < f1_acc - 0.5):
print("&&&& FAILED TensorRT BERT Squad Accuracy matches reference.")
else:
print("&&&& PASSED TensorRT BERT Squad Accuracy matches reference.")
return {'exact_match': exact_match, 'f1': f1}
if __name__ == '__main__':
expected_version = '1.1'
parser = argparse.ArgumentParser(
description='Evaluation for SQuAD ' + expected_version)
parser.add_argument('dataset_file', help='Dataset file')
parser.add_argument('prediction_file', help='Prediction File')
parser.add_argument('f1_acc', help='Reference Accuracy')
args = parser.parse_args()
with open(args.dataset_file) as dataset_file:
dataset_json = json.load(dataset_file)
if (dataset_json['version'] != expected_version):
print('Evaluation expects v-' + expected_version +
', but got dataset with v-' + dataset_json['version'],
file=sys.stderr)
dataset = dataset_json['data']
with open(args.prediction_file) as prediction_file:
predictions = json.load(prediction_file)
f1_acc = float(args.f1_acc)
print(json.dumps(evaluate(dataset, predictions, f1_acc)))
demo/BERT/squad/evaluate-v2.0.py
-296
View File
@@ -1,296 +0,0 @@
#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Obtained from https://rajpurkar.github.io/SQuAD-explorer/
"""Official evaluation script for SQuAD version 2.0.
In addition to basic functionality, we also compute additional statistics and
plot precision-recall curves if an additional na_prob.json file is provided.
This file is expected to map question ID's to the model's predicted probability
that a question is unanswerable.
"""
import argparse
import collections
import json
import numpy as np
import os
import re
import string
import sys
OPTS = None
def parse_args():
parser = argparse.ArgumentParser('Official evaluation script for SQuAD version 2.0.')
parser.add_argument('data_file', metavar='data.json', help='Input data JSON file.')
parser.add_argument('pred_file', metavar='pred.json', help='Model predictions.')
parser.add_argument('--out-file', '-o', metavar='eval.json',
help='Write accuracy metrics to file (default is stdout).')
parser.add_argument('--na-prob-file', '-n', metavar='na_prob.json',
help='Model estimates of probability of no answer.')
parser.add_argument('--na-prob-thresh', '-t', type=float, default=1.0,
help='Predict "" if no-answer probability exceeds this (default = 1.0).')
parser.add_argument('--out-image-dir', '-p', metavar='out_images', default=None,
help='Save precision-recall curves to directory.')
parser.add_argument('--verbose', '-v', action='store_true')
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
return parser.parse_args()
def make_qid_to_has_ans(dataset):
qid_to_has_ans = {}
for article in dataset:
for p in article['paragraphs']:
for qa in p['qas']:
qid_to_has_ans[qa['id']] = bool(qa['answers'])
return qid_to_has_ans
def normalize_answer(s):
"""Lower text and remove punctuation, articles and extra whitespace."""
def remove_articles(text):
regex = re.compile(r'\b(a|an|the)\b', re.UNICODE)
return re.sub(regex, ' ', text)
def white_space_fix(text):
return ' '.join(text.split())
def remove_punc(text):
exclude = set(string.punctuation)
return ''.join(ch for ch in text if ch not in exclude)
def lower(text):
return text.lower()
return white_space_fix(remove_articles(remove_punc(lower(s))))
def get_tokens(s):
if not s: return []
return normalize_answer(s).split()
def compute_exact(a_gold, a_pred):
return int(normalize_answer(a_gold) == normalize_answer(a_pred))
def compute_f1(a_gold, a_pred):
gold_toks = get_tokens(a_gold)
pred_toks = get_tokens(a_pred)
common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
num_same = sum(common.values())
if len(gold_toks) == 0 or len(pred_toks) == 0:
# If either is no-answer, then F1 is 1 if they agree, 0 otherwise
return int(gold_toks == pred_toks)
if num_same == 0:
return 0
precision = 1.0 * num_same / len(pred_toks)
recall = 1.0 * num_same / len(gold_toks)
f1 = (2 * precision * recall) / (precision + recall)
return f1
def get_raw_scores(dataset, preds):
exact_scores = {}
f1_scores = {}
for article in dataset:
for p in article['paragraphs']:
for qa in p['qas']:
qid = qa['id']
gold_answers = [a['text'] for a in qa['answers']
if normalize_answer(a['text'])]
if not gold_answers:
# For unanswerable questions, only correct answer is empty string
gold_answers = ['']
if qid not in preds:
print('Missing prediction for %s' % qid)
continue
a_pred = preds[qid]
# Take max over all gold answers
exact_scores[qid] = max(compute_exact(a, a_pred) for a in gold_answers)
f1_scores[qid] = max(compute_f1(a, a_pred) for a in gold_answers)
return exact_scores, f1_scores
def apply_no_ans_threshold(scores, na_probs, qid_to_has_ans, na_prob_thresh):
new_scores = {}
for qid, s in scores.items():
pred_na = na_probs[qid] > na_prob_thresh
if pred_na:
new_scores[qid] = float(not qid_to_has_ans[qid])
else:
new_scores[qid] = s
return new_scores
def make_eval_dict(exact_scores, f1_scores, qid_list=None):
if not qid_list:
total = len(exact_scores)
return collections.OrderedDict([
('exact', 100.0 * sum(exact_scores.values()) / total),
('f1', 100.0 * sum(f1_scores.values()) / total),
('total', total),
])
else:
total = len(qid_list)
return collections.OrderedDict([
('exact', 100.0 * sum(exact_scores[k] for k in qid_list) / total),
('f1', 100.0 * sum(f1_scores[k] for k in qid_list) / total),
('total', total),
])
def merge_eval(main_eval, new_eval, prefix):
for k in new_eval:
main_eval['%s_%s' % (prefix, k)] = new_eval[k]
def plot_pr_curve(precisions, recalls, out_image, title):
plt.step(recalls, precisions, color='b', alpha=0.2, where='post')
plt.fill_between(recalls, precisions, step='post', alpha=0.2, color='b')
plt.xlabel('Recall')
plt.ylabel('Precision')
plt.xlim([0.0, 1.05])
plt.ylim([0.0, 1.05])
plt.title(title)
plt.savefig(out_image)
plt.clf()
def make_precision_recall_eval(scores, na_probs, num_true_pos, qid_to_has_ans,
out_image=None, title=None):
qid_list = sorted(na_probs, key=lambda k: na_probs[k])
true_pos = 0.0
cur_p = 1.0
cur_r = 0.0
precisions = [1.0]
recalls = [0.0]
avg_prec = 0.0
for i, qid in enumerate(qid_list):
if qid_to_has_ans[qid]:
true_pos += scores[qid]
cur_p = true_pos / float(i+1)
cur_r = true_pos / float(num_true_pos)
if i == len(qid_list) - 1 or na_probs[qid] != na_probs[qid_list[i+1]]:
# i.e., if we can put a threshold after this point
avg_prec += cur_p * (cur_r - recalls[-1])
precisions.append(cur_p)
recalls.append(cur_r)
if out_image:
plot_pr_curve(precisions, recalls, out_image, title)
return {'ap': 100.0 * avg_prec}
def run_precision_recall_analysis(main_eval, exact_raw, f1_raw, na_probs,
qid_to_has_ans, out_image_dir):
if out_image_dir and not os.path.exists(out_image_dir):
os.makedirs(out_image_dir)
num_true_pos = sum(1 for v in qid_to_has_ans.values() if v)
if num_true_pos == 0:
return
pr_exact = make_precision_recall_eval(
exact_raw, na_probs, num_true_pos, qid_to_has_ans,
out_image=os.path.join(out_image_dir, 'pr_exact.png'),
title='Precision-Recall curve for Exact Match score')
pr_f1 = make_precision_recall_eval(
f1_raw, na_probs, num_true_pos, qid_to_has_ans,
out_image=os.path.join(out_image_dir, 'pr_f1.png'),
title='Precision-Recall curve for F1 score')
oracle_scores = {k: float(v) for k, v in qid_to_has_ans.items()}
pr_oracle = make_precision_recall_eval(
oracle_scores, na_probs, num_true_pos, qid_to_has_ans,
out_image=os.path.join(out_image_dir, 'pr_oracle.png'),
title='Oracle Precision-Recall curve (binary task of HasAns vs. NoAns)')
merge_eval(main_eval, pr_exact, 'pr_exact')
merge_eval(main_eval, pr_f1, 'pr_f1')
merge_eval(main_eval, pr_oracle, 'pr_oracle')
def histogram_na_prob(na_probs, qid_list, image_dir, name):
if not qid_list:
return
x = [na_probs[k] for k in qid_list]
weights = np.ones_like(x) / float(len(x))
plt.hist(x, weights=weights, bins=20, range=(0.0, 1.0))
plt.xlabel('Model probability of no-answer')
plt.ylabel('Proportion of dataset')
plt.title('Histogram of no-answer probability: %s' % name)
plt.savefig(os.path.join(image_dir, 'na_prob_hist_%s.png' % name))
plt.clf()
def find_best_thresh(preds, scores, na_probs, qid_to_has_ans):
num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k])
cur_score = num_no_ans
best_score = cur_score
best_thresh = 0.0
qid_list = sorted(na_probs, key=lambda k: na_probs[k])
for i, qid in enumerate(qid_list):
if qid not in scores: continue
if qid_to_has_ans[qid]:
diff = scores[qid]
else:
if preds[qid]:
diff = -1
else:
diff = 0
cur_score += diff
if cur_score > best_score:
best_score = cur_score
best_thresh = na_probs[qid]
return 100.0 * best_score / len(scores), best_thresh
def find_all_best_thresh(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans):
best_exact, exact_thresh = find_best_thresh(preds, exact_raw, na_probs, qid_to_has_ans)
best_f1, f1_thresh = find_best_thresh(preds, f1_raw, na_probs, qid_to_has_ans)
main_eval['best_exact'] = best_exact
main_eval['best_exact_thresh'] = exact_thresh
main_eval['best_f1'] = best_f1
main_eval['best_f1_thresh'] = f1_thresh
def main():
with open(OPTS.data_file) as f:
dataset_json = json.load(f)
dataset = dataset_json['data']
with open(OPTS.pred_file) as f:
preds = json.load(f)
if OPTS.na_prob_file:
with open(OPTS.na_prob_file) as f:
na_probs = json.load(f)
else:
na_probs = {k: 0.0 for k in preds}
qid_to_has_ans = make_qid_to_has_ans(dataset) # maps qid to True/False
has_ans_qids = [k for k, v in qid_to_has_ans.items() if v]
no_ans_qids = [k for k, v in qid_to_has_ans.items() if not v]
exact_raw, f1_raw = get_raw_scores(dataset, preds)
exact_thresh = apply_no_ans_threshold(exact_raw, na_probs, qid_to_has_ans,
OPTS.na_prob_thresh)
f1_thresh = apply_no_ans_threshold(f1_raw, na_probs, qid_to_has_ans,
OPTS.na_prob_thresh)
out_eval = make_eval_dict(exact_thresh, f1_thresh)
if has_ans_qids:
has_ans_eval = make_eval_dict(exact_thresh, f1_thresh, qid_list=has_ans_qids)
merge_eval(out_eval, has_ans_eval, 'HasAns')
if no_ans_qids:
no_ans_eval = make_eval_dict(exact_thresh, f1_thresh, qid_list=no_ans_qids)
merge_eval(out_eval, no_ans_eval, 'NoAns')
if OPTS.na_prob_file:
find_all_best_thresh(out_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans)
if OPTS.na_prob_file and OPTS.out_image_dir:
run_precision_recall_analysis(out_eval, exact_raw, f1_raw, na_probs,
qid_to_has_ans, OPTS.out_image_dir)
histogram_na_prob(na_probs, has_ans_qids, OPTS.out_image_dir, 'hasAns')
histogram_na_prob(na_probs, no_ans_qids, OPTS.out_image_dir, 'noAns')
if OPTS.out_file:
with open(OPTS.out_file, 'w') as f:
json.dump(out_eval, f)
else:
print(json.dumps(out_eval, indent=2))
if __name__ == '__main__':
OPTS = parse_args()
if OPTS.out_image_dir:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
main()
demo/Diffusion/README.md
+9 -10
View File
@@ -7,7 +7,7 @@ This demo application ("demoDiffusion") showcases the acceleration of Stable Dif
### Clone the TensorRT OSS repository
```bash
git clone git@github.com:NVIDIA/TensorRT.git -b release/10.16 --single-branch
git clone git@github.com:NVIDIA/TensorRT.git -b release/11.0 --single-branch
cd TensorRT
```
@@ -23,7 +23,7 @@ mkdir -p deps
docker run --rm -it --gpus all \
-v $PWD:/workspace \
-v $PWD/deps:/workspace/deps \
nvcr.io/nvidia/pytorch:25.09-py3 /bin/bash
nvcr.io/nvidia/pytorch:26.03-py3 /bin/bash
```
> **NOTE:** Mounting `/workspace/deps` as a volume ensures dependencies persist across container restarts. After initial installation, subsequent container launches will reuse the installed dependencies.
@@ -243,19 +243,19 @@ Note that a denosing-percentage is applied to the number of denoising-steps when
```bash
# Depth BF16
python3 demo_controlnet_sd35.py "a photo of a man" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 40 --guidance-scale 4.5 --bf16 --download-onnx-models
python3 demo_controlnet_sd35.py "a photo of a man" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 40 --guidance-scale 4.5 --bf16 --download-onnx-models --low-vram
# Depth FP8
python3 demo_controlnet_sd35.py "a photo of a man" --version=3.5-large --fp8 --controlnet-type depth --download-onnx-models --denoising-steps=40 --guidance-scale 4.5 --hf-token=$HF_TOKEN
python3 demo_controlnet_sd35.py "a photo of a man" --version=3.5-large --fp8 --controlnet-type depth --download-onnx-models --denoising-steps=40 --guidance-scale 4.5 --hf-token=$HF_TOKEN --low-vram
# Canny BF16
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models --low-vram
# Canny FP8
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --version=3.5-large --fp8 --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --download-onnx-models
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --version=3.5-large --fp8 --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --low-vram --download-onnx-models
# Blur
python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater" --controlnet-type blur --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models
python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater" --controlnet-type blur --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models --low-vram
```
### Generate a video guided by an initial image using Stable Video Diffusion
@@ -263,9 +263,8 @@ python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in
Download the pre-exported ONNX model
```bash
git lfs install
git clone https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1-tensorrt onnx-svd-xt-1-1
cd onnx-svd-xt-1-1 && git lfs pull && cd ..
pip install -U "huggingface_hub[cli]"
hf download stabilityai/stable-video-diffusion-img2vid-xt-1-1-tensorrt --local-dir onnx-svd-xt-1-1
```
SVD-XT-1.1 (25 frames at resolution 576x1024)
demo/Diffusion/demo_controlnet_sd35.py
+1 -1
View File
@@ -109,7 +109,7 @@ def process_demo_args(args):
raise ValueError(f"`--controlnet-scale` must be of type `float`, but is {type(args.controlnet_scale)}")
# Convert controlnet scales to tensor
controlnet_scale = torch.FloatTensor([args.controlnet_scale])
controlnet_scale = torch.tensor(args.controlnet_scale)
# Check images
input_images = []
demo/Diffusion/demo_diffusion/dd_argparse.py
+1 -1
View File
@@ -132,7 +132,7 @@ def add_arguments(parser):
"--onnx-opset",
type=int,
default=19,
choices=range(7, 20),
choices=range(7, 24),
help="Select ONNX opset version to target for exported models",
)
parser.add_argument("--onnx-dir", default="onnx", help="Output directory for ONNX export")
demo/Diffusion/demo_diffusion/deps.py
+22
View File
@@ -50,6 +50,21 @@ def _resolve_deps_root(deps_root: str | None) -> str:
return deps_root
return os.environ.get("TENSORRT_DIFFUSION_DEPS_ROOT", "/workspace/deps")
def _prepend_env_path(var: str, path: str, clean_root: str | None = None):
"""Prepend `path` to an os.pathsep-separated env var so child processes
(e.g. the `polygraphy` CLI) inherit the group's dependencies. If
`clean_root` is given, drop existing entries under it first."""
def _norm(p: str) -> str:
return os.path.abspath(os.path.expanduser(p))
existing = [p for p in os.environ.get(var, "").split(os.pathsep) if p]
if clean_root is not None:
root_abs = _norm(clean_root).rstrip(os.sep) + os.sep
existing = [p for p in existing if not _norm(p).startswith(root_abs)]
existing = [p for p in existing if _norm(p) != _norm(path)]
os.environ[var] = os.pathsep.join([path, *existing])
def _clean_diffusion_paths(deps_root: str = "/workspace/deps"):
"""Drop any existing paths under deps_root from sys.path."""
# Filter out any paths that live under deps_root (robust to path forms)
@@ -140,6 +155,13 @@ def configure(
# Insert at the beginning to override any system packages
sys.path.insert(0, deps_path)
# Mirror onto the environment so child processes (e.g. the polygraphy
# CLI that engine builds shell out to) use these deps too. When clean,
# drop other groups' entries under deps_root, matching sys.path above.
clean_root = deps_root if clean else None
_prepend_env_path("PYTHONPATH", deps_path, clean_root=clean_root)
_prepend_env_path("PATH", os.path.join(group_dir, "bin"), clean_root=clean_root)
if verbose:
description = GROUP_DESCRIPTIONS.get(group, group)
print(f"Configured dependencies: {description}")
demo/Diffusion/demo_diffusion/engine.py
+7 -22
View File
@@ -153,12 +153,7 @@ class Engine:
def build(
self,
onnx_path,
strongly_typed=False,
fp16=True,
bf16=False,
tf32=False,
int8=False,
fp8=False,
input_profile=None,
enable_refit=False,
enable_all_tactics=False,
@@ -172,32 +167,20 @@ class Engine:
):
print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
# Handle weight streaming case: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#streaming-weights.
if weight_streaming:
strongly_typed, fp16, bf16, int8, fp8 = True, False, False, False, False
# Base command
build_command = [f"polygraphy convert {onnx_path} --convert-to trt --output {self.engine_path}"]
# Precision flags
# Build arguments
build_args = [
"--fp16" if fp16 else "",
"--bf16" if bf16 else "",
"--strongly-typed",
"--tf32" if tf32 else "",
"--fp8" if fp8 else "",
"--int8" if int8 else "",
"--strongly-typed" if strongly_typed else "",
]
# Additional arguments
build_args.extend([
"--weight-streaming" if weight_streaming else "",
"--refittable" if enable_refit else "",
"--tactic-sources" if not enable_all_tactics else "",
"--onnx-flags native_instancenorm" if native_instancenorm else "",
f"--builder-optimization-level {builder_optimization_level}",
f"--precision-constraints {precision_constraints}",
])
]
# Timing cache
if timing_cache:
@@ -273,8 +256,10 @@ class Engine:
print(f"[W]: Unload an unloaded engine {self.engine_path}, skip unloading")
def activate(self, device_memory=None):
if device_memory:
self.context = self.engine.create_execution_context_without_device_memory()
if device_memory is not None:
self.context = self.engine.create_execution_context(
trt.ExecutionContextAllocationStrategy.USER_MANAGED
)
self.context.device_memory = device_memory
else:
self.context = self.engine.create_execution_context()
demo/Diffusion/demo_diffusion/model/base_model.py
+15 -4
View File
@@ -15,6 +15,7 @@
# limitations under the License.
#
import gc
import json
import os
@@ -170,8 +171,12 @@ class BaseModel:
else:
# WAR: Enable autocast for BF16 Stable Cascade pipeline
do_autocast = True if self.version == "cascade" and self.bf16 else False
model = self.get_model()
with torch.inference_mode(), torch.autocast("cuda", enabled=do_autocast):
export_onnx(self.get_model())
export_onnx(model)
del model
gc.collect()
torch.cuda.empty_cache()
else:
print(f"[I] Found cached ONNX model: {onnx_path}")
@@ -246,6 +251,12 @@ class BaseModel:
is_fp16_io = kwargs.get("is_fp16_io", True)
opt.modify_fp8_graph(is_fp16_io=is_fp16_io)
opt.info(self.name + ": modify fp8 graph")
elif self.bf16:
# Cast Resize I/O for strongly-typed TRT builds: BF16 -> FP32 inputs, FP32 -> BF16 outputs.
# TRT does not support BF16 for the Resize operator.
opt.infer_shapes()
opt.cast_resize_io(output_dtype=onnx.TensorProto.BFLOAT16)
opt.info(self.name + ": cast resize I/O for bf16")
if self.version.startswith("flux.1") and self.fp8:
opt.flux_convert_rope_weight_type()
opt.info(self.name + ": convert rope weight type for fp8 flux")
@@ -253,9 +264,9 @@ class BaseModel:
opt.info(self.name + ": fold constants")
opt.infer_shapes()
opt.info(self.name + ": shape inference")
if kwargs.get("fuse_mha_qkv_int8", False):
opt.fuse_mha_qkv_int8_sq()
opt.info(self.name + ": fuse QKV nodes")
if kwargs.get("modify_int8_graph", False):
opt.modify_int8_graph()
opt.info(self.name + ": modify int8 graph")
onnx_opt_graph = opt.cleanup(return_onnx=return_onnx)
opt.info(self.name + ": finished")
return onnx_opt_graph
demo/Diffusion/demo_diffusion/model/controlnet.py
+1 -3
View File
@@ -63,7 +63,6 @@ class SD3ControlNet(base_model.BaseModel):
int8=False,
fp8=False,
max_batch_size=16,
build_strongly_typed=False,
do_classifier_free_guidance=False,
):
super(SD3ControlNet, self).__init__(
@@ -94,7 +93,6 @@ class SD3ControlNet(base_model.BaseModel):
print(f"[I] Load SD3ControlNetModel config from: {self.controlnet_model_dir}")
self.config = SD3ControlNetModel.load_config(self.controlnet_model_dir)
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
self.build_strongly_typed = build_strongly_typed
def get_model(self, torch_inference=""):
model_opts = (
@@ -187,7 +185,7 @@ class SD3ControlNet(base_model.BaseModel):
"timestep": (self.xB * batch_size,),
"pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]),
"controlnet_cond": (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width),
"conditioning_scale": (1,),
"conditioning_scale": (),
"controlnet_block_samples": (
self.config["num_layers"],
self.xB * batch_size,
demo/Diffusion/demo_diffusion/model/diffusion_transformer.py
+3 -23
View File
@@ -158,7 +158,6 @@ class FluxTransformerModel(base_model.BaseModel):
bf16=False,
max_batch_size=16,
text_maxlen=77,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
kontext_resolution=None,
@@ -187,7 +186,6 @@ class FluxTransformerModel(base_model.BaseModel):
else:
print(f"[I] Load FluxTransformer2DModel config from: {self.transformer_model_dir}")
self.config = FluxTransformer2DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.out_channels = self.config.get("out_channels") or self.config["in_channels"]
@@ -235,6 +233,7 @@ class FluxTransformerModel(base_model.BaseModel):
"pooled_projections": {0: "B"},
"timestep": {0: "B"},
"img_ids": {0: "latent_dim"},
"txt_ids": {},
}
if self.config["guidance_embeds"]:
dynamic_axes["guidance"] = {0: "B"}
@@ -358,7 +357,7 @@ class FluxTransformerModel(base_model.BaseModel):
if self.fp8:
return super().optimize(onnx_graph)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph, modify_int8_graph=True)
return super().optimize(onnx_graph)
@@ -407,7 +406,6 @@ class SD3TransformerModel(base_model.BaseModel):
fp4=False,
max_batch_size=16,
text_maxlen=256,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
do_classifier_free_guidance=False,
@@ -437,7 +435,6 @@ class SD3TransformerModel(base_model.BaseModel):
else:
print(f"[I] Load SD3Transformer2DModel config from: {self.transformer_model_dir}")
self.config = SD3Transformer2DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.out_channels = self.config.get("out_channels")
@@ -631,7 +628,6 @@ class WanTransformerModel(base_model.BaseModel):
num_frames=81,
height=720,
width=1280,
build_strongly_typed=True,
weight_streaming=False,
weight_streaming_budget_percentage=None,
):
@@ -666,7 +662,6 @@ class WanTransformerModel(base_model.BaseModel):
print(f"[I] Load WanTransformer3DModel config from: {self.transformer_model_dir}")
self.config = WanTransformer3DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.do_constant_folding = False
@@ -709,20 +704,7 @@ class WanTransformerModel(base_model.BaseModel):
return ["denoised_latents"]
def get_dynamic_axes(self):
return {
"hidden_states": {
0: "batch",
2: "frames",
3: "latent_height",
4: "latent_width"
},
"timestep": {
0: "batch"
},
"encoder_hidden_states": {
0: "batch",
},
}
return {}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape, num_frames):
latent_height, latent_width, latent_frames = self.check_dims(
@@ -840,7 +822,6 @@ class CosmosTransformerModel(base_model.BaseModel):
bf16=False,
max_batch_size=16,
text_maxlen=77,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
):
@@ -868,7 +849,6 @@ class CosmosTransformerModel(base_model.BaseModel):
else:
print(f"[I] Load CosmosTransformer3DModel config from: {self.transformer_model_dir}")
self.config = CosmosTransformer3DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
demo/Diffusion/demo_diffusion/model/gan.py
+6
View File
@@ -106,6 +106,12 @@ class VQGANModel(base_model.BaseModel):
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
return torch.randn(batch_size, 4, latent_height, latent_width, dtype=dtype, device=self.device)
def optimize(self, onnx_graph, return_onnx=True, **kwargs):
onnx_opt_graph = super().optimize(onnx_graph, return_onnx=True, **kwargs)
opt = optimizer.Optimizer(onnx_opt_graph, verbose=self.verbose, version=self.version)
opt.cast_convtranspose_io()
return opt.cleanup(return_onnx=return_onnx)
def check_dims(self, batch_size, image_height, image_width):
latent_height, latent_width = super().check_dims(batch_size, image_height, image_width)
latent_height = int(latent_height * self.latent_dim_scale)
demo/Diffusion/demo_diffusion/model/optimizer.py
+14
View File
@@ -27,7 +27,9 @@ from polygraphy.backend.onnx.loader import fold_constants
from demo_diffusion.model import load
from demo_diffusion.utils_modelopt import (
cast_convtranspose_io,
cast_fp8_mha_io,
cast_layernorm_io,
cast_resize_io,
convert_fp16_io,
convert_zp_fp8,
@@ -183,6 +185,18 @@ class Optimizer:
print(f"Removed {removed} QDQ nodes")
return removed # expected 72 for L2.5
def modify_int8_graph(self):
# Cast LayerNorm scale/bias from FP16 to FP32 to match INT8 DQ activations.
cast_layernorm_io(self.graph)
# Fuse QKV QDQ nodes for INT8 SmoothQuant.
self.fuse_mha_qkv_int8_sq()
def cast_convtranspose_io(self):
cast_convtranspose_io(self.graph)
def cast_resize_io(self, output_dtype):
cast_resize_io(self.graph, output_dtype=output_dtype)
def modify_fp8_graph(self, is_fp16_io=True):
onnx_graph = gs.export_onnx(self.graph)
# Convert INT8 Zero to FP8.
demo/Diffusion/demo_diffusion/model/t5.py
-2
View File
@@ -43,7 +43,6 @@ class T5Model(base_model.BaseModel):
bf16=False,
subfolder="text_encoder",
text_maxlen=512,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
use_attention_mask=False,
@@ -71,7 +70,6 @@ class T5Model(base_model.BaseModel):
print(f"[I] Load T5Encoder Config from: {self.t5_model_dir}")
self.config = AutoConfig.from_pretrained(self.t5_model_dir)
self.is_umt5 = getattr(self.config, 'model_type', '') == 'umt5'
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.use_attention_mask = use_attention_mask
demo/Diffusion/demo_diffusion/model/unet.py
+2 -2
View File
@@ -287,7 +287,7 @@ class UNetModel(base_model.BaseModel):
if self.fp8:
return super().optimize(onnx_graph, modify_fp8_graph=True)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph, modify_int8_graph=True)
return super().optimize(onnx_graph)
@@ -428,7 +428,7 @@ class UNetXLModel(base_model.BaseModel):
if self.fp8:
return super().optimize(onnx_graph, modify_fp8_graph=True)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph, modify_int8_graph=True)
return super().optimize(onnx_graph)
demo/Diffusion/demo_diffusion/model/vae.py
+1
View File
@@ -64,6 +64,7 @@ class VAEModel(base_model.BaseModel):
bf16=bf16,
max_batch_size=max_batch_size,
)
self.do_constant_folding = False
self.subfolder = "vae"
self.vae_decoder_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
demo/Diffusion/demo_diffusion/pipeline/cosmos_pipeline.py
-2
View File
@@ -176,7 +176,6 @@ class CosmosPipeline(DiffusionPipeline):
tf32=self.tf32,
bf16=self.bf16,
text_maxlen=self.max_sequence_length,
build_strongly_typed=True,
use_attention_mask=True,
)
@@ -189,7 +188,6 @@ class CosmosPipeline(DiffusionPipeline):
fp8=fp8,
tf32=self.tf32,
text_maxlen=self.max_sequence_length,
build_strongly_typed=True,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage,
)
demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py
+41 -12
View File
@@ -595,6 +595,39 @@ class DiffusionPipeline(ABC):
# Native export is supported by default
return True
@staticmethod
def _fix_bf16_resize_nodes(onnx_opt_path):
"""Cast Resize node I/O for strongly-typed TRT engine builds.
TRT does not support BF16 for the Resize operator, so inputs are
cast to FP32 and outputs are cast back to BF16."""
import onnx
import onnx_graphsurgeon as gs
from demo_diffusion.model import load
from demo_diffusion.utils_modelopt import cast_resize_io
onnx_graph = onnx.load(onnx_opt_path, load_external_data=True)
graph = gs.import_onnx(onnx_graph)
resize_nodes = [n for n in graph.nodes if n.op == "Resize"]
if not resize_nodes:
return
print(f"[I] Fixing {len(resize_nodes)} BF16 Resize node(s) in downloaded model: {onnx_opt_path}")
cast_resize_io(graph, output_dtype=onnx.TensorProto.BFLOAT16)
graph.cleanup().toposort()
onnx_graph = gs.export_onnx(graph)
if load.onnx_graph_needs_external_data(onnx_graph):
onnx.save_model(
onnx_graph,
onnx_opt_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
convert_attribute=False,
)
else:
onnx.save(onnx_graph, onnx_opt_path)
def _export_onnx(
self,
obj,
@@ -620,11 +653,15 @@ class DiffusionPipeline(ABC):
if do_export_onnx:
if download_onnx_models:
self.download_onnx_models(model_name, model_config)
# Fix Resize nodes for strongly-typed TRT builds.
# Downloaded models bypass optimize(), so apply the fix here.
if obj.bf16:
self._fix_bf16_resize_nodes(model_config['onnx_opt_path'])
do_export_onnx = False
else:
self.is_native_export_supported(model_config)
dynamo = True if (self.pipeline_type.is_video2world() and model_name == "transformer") or (self.pipeline_type.is_txt2vid() and (model_name in ["transformer", "transformer_2"])) else False
dynamo = True if (self.pipeline_type.is_video2world() and model_name == "transformer") or (self.pipeline_type.is_txt2vid() and (model_name in ["transformer", "transformer_2"])) or (self.version.startswith("flux.1") and model_name == "transformer" and obj.fp16) else False
export_kwargs = {
"static_shape": static_shape,
@@ -675,13 +712,9 @@ class DiffusionPipeline(ABC):
def _build_engine(self, obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
fp16amp = False if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else obj.fp16
tf32amp = obj.tf32
bf16amp = False if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else obj.bf16
strongly_typed = True if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else False
weight_streaming = getattr(obj, 'weight_streaming', False)
int8amp = model_config.get('use_int8', False)
precision_constraints = 'prefer' if int8amp else 'none'
precision_constraints = 'none'
input_profile = obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape,
@@ -690,11 +723,7 @@ class DiffusionPipeline(ABC):
engine.build(
model_config["onnx_opt_path"],
strongly_typed=strongly_typed,
fp16=fp16amp,
tf32=tf32amp,
bf16=bf16amp,
int8=int8amp,
input_profile=input_profile,
enable_refit=model_config["do_engine_refit"],
enable_all_tactics=enable_all_tactics,
@@ -884,7 +913,7 @@ class DiffusionPipeline(ABC):
for model_name, engine in self.engine.items():
if self.low_vram:
engine.load()
max_device_memory = max(max_device_memory, engine.engine.device_memory_size)
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
if self.low_vram:
engine.unload()
return max_device_memory
@@ -893,7 +922,7 @@ class DiffusionPipeline(ABC):
device_memory_sizes = {}
for model_name, engine in self.engine.items():
engine.load()
device_memory_sizes[model_name] = engine.engine.device_memory_size
device_memory_sizes[model_name] = engine.engine.device_memory_size_v2
engine.unload()
return device_memory_sizes
demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py
+3 -4
View File
@@ -396,11 +396,12 @@ class FluxPipeline(DiffusionPipeline):
self.fp16 = True if not self.bf16 else False
self.tf32 = True
if "clip" in self.stages:
# BF16 CLIP ONNX export fails with ComplexDouble error in newer PyTorch; use FP16.
self.models["clip"] = CLIPModel(
**models_args,
fp16=self.fp16,
fp16=True,
tf32=self.tf32,
bf16=self.bf16,
bf16=False,
embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type),
keep_pooled_output=True,
subfolder="text_encoder",
@@ -415,7 +416,6 @@ class FluxPipeline(DiffusionPipeline):
bf16=self.bf16,
subfolder="text_encoder_2",
text_maxlen=self.max_sequence_length,
build_strongly_typed=True,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
)
@@ -456,7 +456,6 @@ class FluxPipeline(DiffusionPipeline):
"fp8": fp8,
"tf32": self.tf32,
"text_maxlen": self.max_sequence_length,
"build_strongly_typed": True,
"weight_streaming": self.weight_streaming,
"weight_streaming_budget_percentage": self.denoiser_weight_streaming_budget_percentage,
}
demo/Diffusion/demo_diffusion/pipeline/model_memory_manager.py
+4 -3
View File
@@ -37,6 +37,7 @@ class ModelMemoryManager:
self.parent = parent
self.model_names = model_names
self.low_vram = low_vram
self.device_memories = {}
def __enter__(self):
if not self.low_vram:
@@ -47,9 +48,9 @@ class ModelMemoryManager:
self.parent.engine[model_name].load()
# allocate device memory
_, shared_device_memory = cudart.cudaMalloc(self.parent.device_memory_sizes[model_name])
self.parent.shared_device_memory = shared_device_memory
self.device_memories[model_name] = shared_device_memory
# creating context
self.parent.engine[model_name].activate(device_memory=self.parent.shared_device_memory)
self.parent.engine[model_name].activate(device_memory=shared_device_memory)
# creating input and output buffer
self.parent.engine[model_name].allocate_buffers(
shape_dict=self.parent.shape_dicts[model_name], device=self.parent.device
@@ -66,7 +67,7 @@ class ModelMemoryManager:
self.parent.engine[model_name].deallocate_buffers()
self.parent.engine[model_name].deactivate()
self.parent.engine[model_name].unload()
cudart.cudaFree(self.parent.shared_device_memory)
cudart.cudaFree(self.device_memories.pop(model_name))
else:
print(f"[I] Offloading torch model {model_name} to cpu.")
self.parent.torch_models[model_name] = self.parent.torch_models[model_name].to("cpu")
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py
-2
View File
@@ -278,7 +278,6 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
tf32=self.tf32,
subfolder="text_encoder_3",
text_maxlen=self.max_sequence_length,
build_strongly_typed=True,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
)
@@ -293,7 +292,6 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
int8=self.int8,
fp4=self.fp4,
text_maxlen=self.models["t5"].text_maxlen + self.models["clip_g"].text_maxlen,
build_strongly_typed=False,
weight_streaming=self.weight_streaming,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py
+7 -7
View File
@@ -132,6 +132,8 @@ class StableDiffusion3Pipeline:
self.config = {}
self.config['clip_hidden_states'] = True
self.config['t5xxl_torch_fallback'] = True
self.config['vae_encoder_torch_fallback'] = True
self.torch_inference = torch_inference
if self.torch_inference:
torch._inductor.config.conv_1x1_as_mm = True
@@ -276,7 +278,7 @@ class StableDiffusion3Pipeline:
# Configure pipeline models to load
model_names = self.models.keys()
# Torch fallback
self.torch_fallback = dict(zip(model_names, [self.torch_inference or model_name in ('t5xxl') for model_name in model_names]))
self.torch_fallback = dict(zip(model_names, [self.torch_inference or self.config.get(model_name.replace('-','_')+'_torch_fallback', False) for model_name in model_names]))
onnx_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, opt=False) for model_name in model_names]))
onnx_opt_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir) for model_name in model_names]))
@@ -298,9 +300,7 @@ class StableDiffusion3Pipeline:
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
extra_build_args = {'verbose': self.verbose}
fp16amp = obj.fp16
engine.build(onnx_opt_path[model_name],
fp16=fp16amp,
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
@@ -326,7 +326,7 @@ class StableDiffusion3Pipeline:
def calculateMaxDeviceMemory(self):
max_device_memory = 0
for model_name, engine in self.engine.items():
max_device_memory = max(max_device_memory, engine.engine.device_memory_size)
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
return max_device_memory
def activateEngines(self, shared_device_memory=None):
@@ -446,7 +446,7 @@ class StableDiffusion3Pipeline:
sigma = torch.cat([timestep, timestep])
c_crossattn = torch.cat([cond["c_crossattn"], uncond["c_crossattn"]])
y = torch.cat([cond["y"], uncond["y"]])
if self.torch_inference:
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
batched = self.torch_models[model_name](sample, sigma, c_crossattn=c_crossattn, y=y)
else:
@@ -479,7 +479,7 @@ class StableDiffusion3Pipeline:
def encode_image(self, model_name='vae_encoder'):
self.input_image = self.input_image.to(self.device)
self.profile_start(model_name, color='orange')
if self.torch_inference:
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
latent = self.torch_models[model_name](self.input_image)
else:
@@ -491,7 +491,7 @@ class StableDiffusion3Pipeline:
def decode_latent(self, latent, model_name='vae_decoder'):
self.profile_start(model_name, color='red')
if self.torch_inference:
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
image = self.torch_models[model_name](latent)
else:
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py
+6 -10
View File
@@ -572,6 +572,10 @@ class StableDiffusionPipeline:
print(f"[I] Saving weights map: {weights_map_path[model_name]}")
obj.export_weights_map(onnx_opt_path[model_name], weights_map_path[model_name])
# Release temp GPU memory during onnx export to avoid OOM.
gc.collect()
torch.cuda.empty_cache()
# Build TensorRT engines
for model_name, obj in self.models.items():
if torch_fallback[model_name]:
@@ -579,19 +583,11 @@ class StableDiffusionPipeline:
engine = engine_module.Engine(engine_path[model_name])
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
fp16amp = obj.fp16 if not use_fp8[model_name] else False
bf16amp = obj.bf16 if not use_fp8[model_name] else False
# TF32 can be enabled for all precisions (including INT8/FP8)
tf32amp = obj.tf32
strongly_typed = False if not use_fp8[model_name] else True
int8amp = use_int8.get('model_name', False)
precision_constraints = 'prefer' if int8amp else 'none'
precision_constraints = 'none'
engine.build(onnx_opt_path[model_name],
strongly_typed=strongly_typed,
fp16=fp16amp,
bf16=bf16amp,
tf32=tf32amp,
int8=int8amp,
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
@@ -642,7 +638,7 @@ class StableDiffusionPipeline:
def calculateMaxDeviceMemory(self):
max_device_memory = 0
for model_name, engine in self.engine.items():
max_device_memory = max(max_device_memory, engine.engine.device_memory_size)
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
return max_device_memory
def activateEngines(self, shared_device_memory=None):
demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py
+1 -2
View File
@@ -153,7 +153,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
if not self.low_vram and not self.torch_inference:
for model_name in self.models.keys():
if not self.torch_fallback[model_name]:
self.max_shared_device_memory_size = max(self.max_shared_device_memory_size, self.engine[model_name].engine.device_memory_size)
self.max_shared_device_memory_size = max(self.max_shared_device_memory_size, self.engine[model_name].engine.device_memory_size_v2)
self.shared_device_memory = cudart.cudaMalloc(self.max_shared_device_memory_size)[1]
# Activate TensorRT engines
for model_name in self.models.keys():
@@ -337,7 +337,6 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
engine.build(onnx_opt_path[model_name],
fp16=True,
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
demo/Diffusion/demo_diffusion/pipeline/wan_pipeline.py
-1
View File
@@ -219,7 +219,6 @@ class WanPipeline(DiffusionPipeline):
fp16=False,
bf16=True,
text_maxlen=self.max_sequence_length,
build_strongly_typed=True,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
use_attention_mask=True,
demo/Diffusion/demo_diffusion/utils_modelopt.py
+53 -4
View File
@@ -448,6 +448,52 @@ def insert_cast(graph, input_tensor, attrs):
if next_input.name == input_tensor.name:
next_node.inputs[idx] = output_tensor
def cast_layernorm_io(graph):
"""
Cast LayerNormalization scale and bias inputs from FP16 to FP32.
In INT8 quantized graphs, DequantizeLinear outputs Float32 activations,
but LayerNorm scale/bias remain FP16 from the original model, causing
a type mismatch with --strongly-typed TensorRT builds.
"""
layernorm_nodes = [node for node in graph.nodes if node.op == "LayerNormalization"]
print(f"Found {len(layernorm_nodes)} LayerNormalization nodes to fix")
for node in layernorm_nodes:
# LayerNormalization inputs: 0=X (data), 1=Scale, 2=B (bias, optional)
for i in range(1, len(node.inputs)):
input_tensor = node.inputs[i]
if input_tensor.name and hasattr(input_tensor, 'dtype') and input_tensor.dtype == np.float16:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
def cast_convtranspose_io(graph):
"""
Fix ConvTranspose input/output type mismatches for strongly-typed TRT builds.
In mixed-precision graphs (e.g. BF16 Stable Cascade VQGAN), architectural FP16->FP32
casts can leave a ConvTranspose with a FP32 activation input but FP16 kernel weights.
We cast the activation to match the kernel dtype, then cast the output back to the
original activation dtype so surrounding FP32 ops (e.g. residual Add) are unaffected.
"""
convtranspose_nodes = [node for node in graph.nodes if node.op == "ConvTranspose"]
fixed = 0
for node in convtranspose_nodes:
if len(node.inputs) < 2:
continue
act_input = node.inputs[0]
kernel = node.inputs[1]
if act_input.dtype is None or kernel.dtype is None or act_input.dtype == kernel.dtype:
continue
orig_dtype = act_input.dtype # e.g. np.dtype('float32')
target_dtype = kernel.dtype.type # e.g. np.float16
insert_cast(graph, input_tensor=act_input, attrs={"to": target_dtype})
# Update the output dtype to match and cast back, so downstream FP32 ops are unaffected.
for out in node.outputs:
if out.name and out.dtype == orig_dtype:
out.dtype = target_dtype
insert_cast(graph, input_tensor=out, attrs={"to": orig_dtype.type})
fixed += 1
print(f"Fixed {fixed} ConvTranspose input/output type mismatches")
def convert_zp_fp8(onnx_graph):
"""
Convert Q/DQ zero datatype from INT8 to FP8.
@@ -468,22 +514,25 @@ def convert_zp_fp8(onnx_graph):
return onnx_graph
def cast_resize_io(graph):
def cast_resize_io(graph, output_dtype=np.float16):
"""
After all activations and weights are converted to fp16, we will
add cast nodes to Resize nodes I/O because Resize need to be run in fp32.
Add cast nodes to Resize nodes I/O because Resize needs to be run in fp32.
Inputs are cast to FP32, outputs are cast back to output_dtype (FP16 or BF16).
"""
resize_nodes = [node for node in graph.nodes if node.op == "Resize"]
print(f"Found {len(resize_nodes)} Resize nodes to fix")
for resize_node in resize_nodes:
# Skip Resize nodes whose data input is already FP32 — no casting needed.
if resize_node.inputs[0].dtype == np.float32:
continue
for i, input_tensor in enumerate(resize_node.inputs):
SIZES_INPUT_INDEX = 3 # Optional input "sizes" at index 3 must be in INT64. Skip cast for this input.
if i != SIZES_INPUT_INDEX and input_tensor.name:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
for output_tensor in resize_node.outputs:
if output_tensor.name:
insert_cast(graph, input_tensor=output_tensor, attrs={"to": np.float16})
insert_cast(graph, input_tensor=output_tensor, attrs={"to": output_dtype})
def cast_fp8_mha_io(graph):
r"""
demo/Diffusion/demo_txt2vid_wan.py
+7
View File
@@ -49,6 +49,13 @@ def parseArgs():
),
help="Negative prompt (Wan team default, English translation)"
)
parser.add_argument(
"--onnx-opset",
type=int,
default=23,
choices=range(7, 24),
help="Select ONNX opset version to target for exported models",
)
return parser.parse_args()
demo/Diffusion/pyproject.toml
-110
View File
@@ -1,110 +0,0 @@
[project]
name = "tensorrt-diffusion"
version = "0.1.0"
description = "TensorRT-accelerated implementations of Stable Diffusion and other diffusion models"
readme = "README.md"
requires-python = ">=3.10,<3.13"
# Core dependencies shared across all pipeline families
dependencies = [
"apex==0.9.10dev",
"accelerate==1.2.1",
"colored==2.3.1",
"controlnet-aux==0.0.6",
"cuda-python==13.0.2",
"ftfy==6.3.1",
"matplotlib==3.10.7",
"nvtx==0.2.13",
"opencv-python-headless==4.8.0.74",
"scipy==1.15.3",
"transformers==4.52.4",
"onnx==1.19.0",
"onnxscript==0.5.4",
"onnx-graphsurgeon==0.5.2",
"peft==0.17.0",
"polygraphy==0.49.22",
"sentencepiece==0.2.1",
"numpy==1.26.4",
"nvidia-modelopt[torch,onnx]==0.40.0",
]
[project.optional-dependencies]
# Stable Diffusion family (Stability AI)
# Pipelines: SD 1.4, SDXL, SD3, SD3.5, SVD (Stable Video Diffusion), Stable Cascade
sd = [
"diffusers==0.35.2",
"imageio-ffmpeg",
]
# Flux family (Black Forest Labs)
# Pipelines: Flux.1-dev, Flux.1-schnell, Flux.1-Canny, Flux.1-Depth, Flux.1-Kontext
# NOTE: Diffusers upgrade requires Dynamo export support
flux = [
"diffusers @ git+https://github.com/huggingface/diffusers.git@7298bdd8177c16eadb74f6166327f5984fd8c69d",
"flux @ git+https://github.com/black-forest-labs/flux.git",
]
# Cosmos family (NVIDIA)
# Pipelines: Cosmos-Predict2 Text2Image (2B, 14B), Cosmos-Predict2 Video2World (2B, 14B), Wan2.2
cosmos = [
"diffusers==0.35.2",
"imageio-ffmpeg",
"flux @ git+https://github.com/black-forest-labs/flux.git",
]
# Install all families
all = [
"tensorrt-diffusion[sd]",
"tensorrt-diffusion[flux]",
"tensorrt-diffusion[cosmos]",
]
# Development dependencies
dev = [
"pytest",
"pytest-cov",
"black",
"ruff",
]
[build-system]
requires = ["poetry-core>=1.3.0"]
build-backend = "poetry.core.masonry.api"
[tool.poetry]
# Tell Poetry where the actual Python package is located
packages = [{include = "demo_diffusion"}]
[tool.uv]
# Use PyPI as the main index
index-url = "https://pypi.org/simple"
# Add NVIDIA's PyPI index for nvidia-modelopt and related packages
[[tool.uv.index]]
name = "nvidia"
url = "https://pypi.nvidia.com"
explicit = true
[tool.uv.sources]
# Explicitly fetch these packages from NVIDIA's index
nvidia-modelopt = { index = "nvidia" }
onnx-graphsurgeon = { index = "nvidia" }
polygraphy = { index = "nvidia" }
[tool.ruff]
line-length = 120
target-version = "py310"
[tool.ruff.lint]
select = ["E", "F", "W", "I"]
ignore = ["E501"] # Line too long (handled by formatter)
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_functions = ["test_*"]
[tool.black]
line-length = 120
target-version = ["py310", "py311", "py312"]
demo/Diffusion/setup.py
+119 -56
View File
@@ -32,6 +32,7 @@ Options:
--skip-tensorrt Skip TensorRT upgrade/installation
--force Force reinstallation even if already installed
--deps-root DIR Root directory for dependencies
-q, --quiet Suppress informational output (errors are always shown)
Examples:
python setup.py # Install all
@@ -40,15 +41,25 @@ Examples:
python setup.py cosmos # Install Cosmos only
python setup.py --skip-tensorrt # Install all, skip TensorRT upgrade
python setup.py flux --skip-tensorrt # Install Flux only, skip TensorRT upgrade
python setup.py all --quiet # Install all, minimal output
"""
import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import List, Set
# Global quiet flag (set by parse_args)
_quiet = False
def log(msg: str = ""):
"""Print a message unless in quiet mode."""
if not _quiet:
print(msg)
# Group descriptions
GROUP_DESCRIPTIONS = {
@@ -72,11 +83,59 @@ DEFAULT_DEPS_ROOT = os.environ.get("TENSORRT_DIFFUSION_DEPS_ROOT", "/workspace/d
INSTALL_COMPLETE_MARKER = ".install_complete"
def _normalize_package_name(name: str) -> str:
"""Normalize a package name per PEP 503 (e.g. 'Nvidia_ModelOpt' -> 'nvidia-modelopt')."""
return re.sub(r"[-_.]+", "-", name).lower()
# Regex to extract the package name from a PEP 508 requirement string.
# Matches the leading identifier before any extras, version specifier, or URL marker.
# "nvidia-modelopt[torch,onnx]==0.40.0" -> "nvidia-modelopt"
# "flux @ git+https://..." -> "flux"
_REQ_NAME_RE = re.compile(r"^([A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?)")
def get_pyproject_package_names(pyproject_path: str, group: str) -> set[str]:
"""Return normalized names of packages explicitly listed in pyproject.toml."""
try:
import tomllib
except ModuleNotFoundError:
try:
import tomli as tomllib
except ModuleNotFoundError:
return set()
with open(pyproject_path, "rb") as f:
data = tomllib.load(f)
project = data.get("project", {})
reqs = list(project.get("dependencies", []))
reqs.extend(project.get("optional-dependencies", {}).get(group, []))
names = set()
for req in reqs:
match = _REQ_NAME_RE.match(req.strip())
if match:
names.add(_normalize_package_name(match.group(1)))
return names
def get_container_provided_packages() -> list[str]:
"""Discover torch/NVIDIA packages already installed in the container."""
from importlib.metadata import distributions
prefixes = ("torch", "torchvision", "triton", "nvidia-")
return sorted({
dist.metadata["Name"].lower()
for dist in distributions()
if dist.metadata["Name"].lower().startswith(prefixes)
})
def print_header():
"""Print setup header."""
print("=" * 60)
print(" TensorRT Diffusion - Dependency Setup")
print("=" * 60)
log("=" * 60)
log(" TensorRT Diffusion - Dependency Setup")
log("=" * 60)
def check_uv_installed() -> tuple[bool, bool]:
@@ -125,12 +184,13 @@ def upgrade_tensorrt(pip_spec: str) -> bool:
subprocess.run(
[sys.executable, "-m", "pip", "install", "--upgrade", "--pre", pip_spec],
check=True,
capture_output=_quiet,
)
# Print installed version for confirmation
try:
import importlib
trt = importlib.import_module("tensorrt")
print(" Installed TensorRT version:", getattr(trt, "__version__", "unknown"))
log(f" Installed TensorRT version: {getattr(trt, '__version__', 'unknown')}")
except Exception:
pass
return True
@@ -148,7 +208,7 @@ def install_libgl1() -> bool:
print("Checking for libgl1 (system dependency)...")
apt_get = shutil.which("apt-get")
if not apt_get:
print(" Skipping: apt-get not found. Please install 'libgl1' via your OS package manager.")
log(" Skipping: apt-get not found. Please install 'libgl1' via your OS package manager.")
return False
sudo = shutil.which("sudo")
@@ -161,11 +221,11 @@ def install_libgl1() -> bool:
try:
# Update package list first
cmd_update = ([sudo] if use_sudo else []) + [apt_get, "update"]
subprocess.run(cmd_update, check=True)
subprocess.run(cmd_update, check=True, capture_output=_quiet)
cmd_install = ([sudo] if use_sudo else []) + [apt_get, "install", "-y", "libgl1"]
subprocess.run(cmd_install, check=True)
print(" libgl1 installed (or already up to date)")
subprocess.run(cmd_install, check=True, capture_output=_quiet)
log(" libgl1 installed (or already up to date)")
return True
except subprocess.CalledProcessError as e:
print(f" Warning: Failed to install libgl1: {e}")
@@ -175,7 +235,7 @@ def install_libgl1() -> bool:
def install_uv():
"""Install uv package manager."""
print("Installing uv...")
log("Installing uv...")
try:
# Download and run uv installer
curl_cmd = [
@@ -208,7 +268,7 @@ def install_uv():
local_bin = os.path.expanduser("~/.local/bin")
if local_bin not in os.environ["PATH"]:
os.environ["PATH"] = f"{local_bin}:{os.environ['PATH']}"
print("uv installed successfully")
log("uv installed successfully")
except Exception as e:
print(f"Error installing uv: {e}")
@@ -216,7 +276,7 @@ def install_uv():
sys.exit(1)
def get_installed_groups(deps_root: str) -> Set[str]:
def get_installed_groups(deps_root: str) -> set[str]:
"""
Get set of already installed dependency groups.
@@ -263,7 +323,7 @@ def install_group(group: str, deps_root: str, project_root: str) -> bool:
install_path = os.path.join(deps_root, group)
print(f"Installing {description}...")
print(f" Location: {install_path}")
log(f" Location: {install_path}")
try:
# Determine Python version and site-packages path
@@ -273,44 +333,49 @@ def install_group(group: str, deps_root: str, project_root: str) -> bool:
# Create site-packages directory if it doesn't exist
site_packages.mkdir(parents=True, exist_ok=True)
# Install all dependencies (core + group-specific) to prefix
# Build an overrides file so uv never installs packages that are
# already installed in the container. Exclude packages explicitly
# listed in pyproject.toml so they get installed at the pinned version.
container_pkgs = get_container_provided_packages()
declared_pkgs = get_pyproject_package_names(
str(Path(project_root) / "pyproject.toml"), group
)
overrides_content = "\n".join(
f'{pkg} ; python_version < "0"'
for pkg in container_pkgs
if _normalize_package_name(pkg) not in declared_pkgs
)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".txt", prefix="uv_overrides_", delete=False
) as f:
f.write(overrides_content)
overrides_path = f.name
try:
cmd = [
"uv", "pip", "install",
"--python-preference", "only-system",
"--prefix", install_path,
f".[{group}]" # Install current project with specific extra
"--overrides", overrides_path,
f".[{group}]"
]
result = subprocess.run(
subprocess.run(
cmd,
cwd=project_root,
text=True,
capture_output=False,
capture_output=_quiet,
check=True
)
# Remove torch (use container's version instead)
# Container torch is optimized for NVIDIA hardware
print(" Removing installed torch (using container's version)...")
if site_packages.exists():
for item in site_packages.iterdir():
item_name = item.name.lower()
# Match torch, torch-*, torchvision, torchvision-*
if any(item_name.startswith(prefix) for prefix in ["torch", "torchvision"]):
try:
if item.is_dir():
shutil.rmtree(item)
else:
item.unlink()
except Exception as e:
print(f" Warning: Could not remove {item.name}: {e}")
finally:
os.unlink(overrides_path)
# Create marker file to indicate successful installation
marker_file = Path(install_path) / INSTALL_COMPLETE_MARKER
marker_file.write_text("Installation completed successfully\n")
print(f"{description} installed successfully")
print(f" {description} installed successfully")
return True
except subprocess.CalledProcessError as e:
@@ -319,7 +384,7 @@ def install_group(group: str, deps_root: str, project_root: str) -> bool:
# Clean up incomplete installation
if os.path.exists(install_path):
print(f" Cleaning up incomplete installation at {install_path}")
log(f" Cleaning up incomplete installation at {install_path}")
try:
shutil.rmtree(install_path)
except Exception as cleanup_error:
@@ -328,7 +393,7 @@ def install_group(group: str, deps_root: str, project_root: str) -> bool:
return False
def determine_groups_to_install(requested: str, already_installed: Set[str]) -> List[str]:
def determine_groups_to_install(requested: str, already_installed: set[str]) -> list[str]:
"""
Determine which groups need to be installed.
@@ -343,32 +408,31 @@ def determine_groups_to_install(requested: str, already_installed: Set[str]) ->
return sorted(VALID_GROUPS - already_installed)
elif requested in VALID_GROUPS:
return [] if requested in already_installed else [requested]
else:
return []
def print_summary(installed_groups: Set[str]):
def print_summary(installed_groups: set[str]):
"""
Print summary of installed groups.
Args:
installed_groups: Set of all installed groups
"""
print("=" * 60)
log("=" * 60)
print("Setup complete!")
print("=" * 60)
log("=" * 60)
if installed_groups:
print("Installed groups:")
log("Installed groups:")
for group in sorted(installed_groups):
description = GROUP_DESCRIPTIONS.get(group, group)
print(f" {description}")
log(f" {description}")
print("Each demo script automatically uses the correct dependencies.")
log("Each demo script automatically uses the correct dependencies.")
else:
print("No groups installed.")
log("No groups installed.")
print()
log()
def parse_args():
@@ -413,6 +477,8 @@ Groups:
parser.add_argument("--skip-tensorrt", action="store_true", help="Skip TensorRT upgrade/installation")
parser.add_argument("-q", "--quiet", action="store_true", help="Suppress informational output (errors are always shown)")
return parser.parse_args()
@@ -420,11 +486,8 @@ def main():
"""Main setup function."""
args = parse_args()
# Validate group argument
if args.group not in VALID_GROUPS and args.group != "all":
print(f"Unknown group: {args.group}")
print(f" Valid options: {', '.join(sorted(VALID_GROUPS))}, all")
sys.exit(1)
global _quiet
_quiet = args.quiet
print_header()
@@ -434,9 +497,9 @@ def main():
install_uv()
needs_path_setup = True # Fresh install always needs PATH setup
else:
print("uv is already installed")
log("uv is already installed")
if needs_path_setup:
print(" (temporarily added ~/.local/bin to PATH for this session)")
log(" (temporarily added ~/.local/bin to PATH for this session)")
# Get project root (where pyproject.toml is)
project_root = str(Path(__file__).parent.absolute())
@@ -455,10 +518,10 @@ def main():
if not args.skip_tensorrt:
upgrade_tensorrt("tensorrt-cu12")
else:
print("Skipping TensorRT upgrade (--skip-tensorrt specified)")
log("Skipping TensorRT upgrade (--skip-tensorrt specified)")
if installed_groups and not args.force:
print(f"Already installed: {', '.join(sorted(installed_groups))}")
log(f"Already installed: {', '.join(sorted(installed_groups))}")
# Determine what to install
if args.force and args.group == "all":
@@ -475,7 +538,7 @@ def main():
else:
description = GROUP_DESCRIPTIONS.get(args.group, args.group)
print(f"{description} is already installed!")
print(f" To reinstall, use --force flag or remove {args.deps_root}/{args.group}")
log(f" To reinstall, use --force flag or remove {args.deps_root}/{args.group}")
print_summary(installed_groups)
sys.exit(0)
docker/rockylinux8.Dockerfile
+28 -14
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.2.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux8
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.16.1.11
ENV TRT_VERSION 11.0.0.114
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -46,7 +46,11 @@ RUN dnf -y install \
libnccl \
libnccl-devel \
openmpi \
openmpi-devel
openmpi-devel \
zstd \
epel-release
RUN dnf -y install ccache
# Install python3
RUN dnf install -y python38 python38-devel &&\
@@ -55,15 +59,15 @@ RUN dnf install -y python38 python38-devel &&\
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp38-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& tar --use-compress-program=unzstd -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib64 \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp38-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& tar --use-compress-program=unzstd -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib64 \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp38-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
@@ -77,10 +81,20 @@ RUN pip install jupyter jupyterlab
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-Linux-x86_64.sh && \
chmod +x cmake-3.31.11-Linux-x86_64.sh && \
./cmake-3.31.11-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-Linux-x86_64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set
docker/rockylinux9.Dockerfile
+28 -14
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.2.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux9
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.16.1.11
ENV TRT_VERSION 11.0.0.114
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -56,19 +56,23 @@ RUN dnf -y install \
libnccl \
libnccl-devel \
openmpi \
openmpi-devel
openmpi-devel \
zstd \
epel-release
RUN dnf -y install ccache
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp39-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& tar --use-compress-program=unzstd -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib64 \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp39-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp39-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& tar --use-compress-program=unzstd -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib64 \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp39-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
@@ -77,10 +81,20 @@ RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-Linux-x86_64.sh && \
chmod +x cmake-3.31.11-Linux-x86_64.sh && \
./cmake-3.31.11-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-Linux-x86_64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set
docker/ubuntu-22.04.Dockerfile
+26 -14
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.2.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.16.1.11
ENV TRT_VERSION 11.0.0.114
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -59,7 +59,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
libnccl2 \
libnccl-dev \
openmpi-bin \
libopenmpi-dev
libopenmpi-dev \
zstd \
ccache
# Install python3
RUN apt-get install -y --no-install-recommends \
@@ -73,15 +75,15 @@ RUN apt-get install -y --no-install-recommends \
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp310-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp310-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp310-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp310-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
@@ -97,10 +99,20 @@ RUN pip3 install --upgrade numpy
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-Linux-x86_64.sh && \
chmod +x cmake-3.31.11-Linux-x86_64.sh && \
./cmake-3.31.11-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-Linux-x86_64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set
docker/ubuntu-24.04-aarch64.Dockerfile
+26 -14
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.2.0
# Multi-arch container support available in non-cudnn containers.
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04
ENV TRT_VERSION 10.16.1.11
ENV TRT_VERSION 11.0.0.114
SHELL ["/bin/bash", "-c"]
# Setup user account and edit default account
@@ -64,7 +64,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
libnccl2 \
libnccl-dev \
openmpi-bin \
libopenmpi-dev
libopenmpi-dev \
zstd \
ccache
# Install python3
RUN apt-get install -y --no-install-recommends \
@@ -83,15 +85,15 @@ ENV PATH="/opt/venv/bin:$PATH"
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-13.2.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-13.2.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp312-none-linux_aarch64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-13.2-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-13.2-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp312-none-linux_aarch64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp312-none-linux_aarch64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-12.9-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-12.9-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp312-none-linux_aarch64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
@@ -99,10 +101,20 @@ RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-linux-aarch64.sh && \
chmod +x cmake-3.27.9-linux-aarch64.sh && \
./cmake-3.27.9-linux-aarch64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-linux-aarch64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-linux-aarch64.sh && \
chmod +x cmake-3.31.11-linux-aarch64.sh && \
./cmake-3.31.11-linux-aarch64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-linux-aarch64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Install PyPI packages
RUN pip3 install --upgrade pip
docker/ubuntu-24.04.Dockerfile
+26 -14
View File
@@ -21,7 +21,7 @@ ARG CUDA_VERSION=13.2.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION=10.16.1.11
ENV TRT_VERSION=11.0.0.114
SHELL ["/bin/bash", "-c"]
# Setup user account and edit default account
@@ -63,7 +63,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
libnccl2 \
libnccl-dev \
openmpi-bin \
libopenmpi-dev
libopenmpi-dev \
zstd \
ccache
# Install python3
RUN apt-get install -y --no-install-recommends \
@@ -82,15 +84,15 @@ ENV PATH="/opt/venv/bin:$PATH"
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-13.2.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp312-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-13.2-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp312-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.16.1.11.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-10.16.1.11/python/tensorrt-10.16.1.11-cp312-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& tar -xf TensorRT-Enterprise-11.0.0.114-Linux-x86_64-cuda-12.9-Release-external.tar.zst \
&& cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-11.0.0.114/python/tensorrt-11.0.0.114-cp312-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
@@ -105,10 +107,20 @@ RUN pip3 install --upgrade numpy
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-Linux-x86_64.sh && \
chmod +x cmake-3.31.11-Linux-x86_64.sh && \
./cmake-3.31.11-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-Linux-x86_64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set
docker/ubuntu-cross-aarch64.Dockerfile
+21 -9
View File
@@ -21,7 +21,7 @@ ARG OS_VERSION=24.04
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${OS_VERSION}
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.16.1.11
ENV TRT_VERSION 11.0.0.114
ENV DEBIAN_FRONTEND=noninteractive
# Setup user account and edit default account
@@ -54,7 +54,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
libnccl2 \
libnccl-dev \
openmpi-bin \
libopenmpi-dev
libopenmpi-dev \
zstd \
ccache
# Install python3
RUN apt-get install -y --no-install-recommends \
@@ -73,10 +75,20 @@ ENV PATH="/opt/venv/bin:$PATH"
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
wget https://github.com/Kitware/CMake/releases/download/v3.31.11/cmake-3.31.11-Linux-x86_64.sh && \
chmod +x cmake-3.31.11-Linux-x86_64.sh && \
./cmake-3.31.11-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.31.11-Linux-x86_64.sh
# Install gtest
RUN cd /tmp && \
git clone https://github.com/google/googletest.git -b v1.14.0 && \
cd googletest && \
mkdir build && cd build && \
cmake .. && \
make -j4 && \
make install && \
rm -rf /tmp/googletest
# Install CUDA cross compile toolchain
RUN wget https://developer.download.nvidia.com/compute/cuda/13.2.0/local_installers/cuda-repo-cross-sbsa-ubuntu2404-13-2-local_13.2.0-1_all.deb && \
@@ -87,9 +99,9 @@ RUN wget https://developer.download.nvidia.com/compute/cuda/13.2.0/local_install
# Unpack libnvinfer.
RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-13.2.tar.gz && \
tar -xf TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-13.2.tar.gz && \
cp -a TensorRT-10.16.1.11/lib/*.so* /usr/lib/aarch64-linux-gnu
RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/11.0.0/tars/TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-13.2-Release-external.tar.zst && \
tar -xf TensorRT-Enterprise-11.0.0.114-Linux-aarch64-cuda-13.2-Release-external.tar.zst && \
cp -a TensorRT-11.0.0.114/lib/*.so* /usr/lib/aarch64-linux-gnu
# Link required library
RUN cd /usr/aarch64-linux-gnu/lib && ln -sf librt.so.1 librt.so
documents/import_workflows.md
+466
View File
@@ -0,0 +1,466 @@
# TensorRT Import Workflows — Step-by-Step Guide
This guide, together with [`supported_models.md`](./supported_models.md), centralizes TensorRT import-path guidance that was previously fragmented across release notes, samples, blogs, and forum posts.
TensorRT supports several paths for bringing a trained model into an optimized inference engine. This guide walks through each path end to end — install, export, build, verify — with runnable commands and the most common pitfalls called out inline.
## Table of Contents
- [Choosing a Path](#choosing-a-path)
- [Common Prerequisites](#common-prerequisites)
- [Path 1: ONNX → TensorRT](#path-1-onnx--tensorrt)
- [Path 2: Torch-TensorRT (PyTorch native)](#path-2-torch-tensorrt-pytorch-native)
- [Path 3: Hugging Face Hub Models → TensorRT](#path-3-hugging-face-hub-models--tensorrt)
- [Path 4: Direct Network Definition API (C++/Python)](#path-4-direct-network-definition-api-cpython)
- [Adding a Custom Operator / Plugin](#adding-a-custom-operator--plugin)
- [AI-Assisted Model Rewriting for Export](#ai-assisted-model-rewriting-for-export)
- [Verifying an Engine](#verifying-an-engine)
- [Tools Reference](#tools-reference)
- [Troubleshooting & Insights](#troubleshooting--insights)
---
## Choosing a Path
| You have… | Recommended path | Notes |
|---------------------------------------------|------------------------------------|-------|
| An ONNX file from any framework | [ONNX → TensorRT](#path-1-onnx--tensorrt) | Most portable path. Build via `trtexec`, Python API, or Polygraphy. |
| A trained PyTorch model, want fastest onboarding | [Torch-TensorRT](#path-2-torch-tensorrt-pytorch-native) | Python-first, stays in PyTorch. Best for iterative development. |
| A Hugging Face Hub model (LLM, diffusion, etc.) | [Hugging Face Hub Models](#path-3-hugging-face-hub-models--tensorrt) | Export → ONNX → TRT for most models; use TensorRT-LLM directly for LLM generation. |
| A model architecture authored in C++ or a custom research stack | [Network Definition API](#path-4-direct-network-definition-api-cpython) | Maximum control, maximum effort. |
| An existing TRT plan that just needs to be run | See [Verifying an Engine](#verifying-an-engine) | Not covered by this guide — see the Developer Guide on deserialization. |
---
## Common Prerequisites
All paths below assume:
1. A supported NVIDIA GPU (see the [Support Matrix](https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/support-matrix.html)).
2. NVIDIA driver + CUDA matching your TensorRT release. For TRT 11.x: **CUDA 13.x**.
3. Python 3.10+ if using Python APIs. C++ paths need a C++17 compiler.
### Install TensorRT (Python, pip)
```bash
# Python TRT runtime + Python bindings
pip install --extra-index-url https://pypi.nvidia.com tensorrt-cu13
```
> **Note:** Always use `-cu13` packages with TRT 11.x. Do not mix `-cu12` wheels.
### Install TensorRT (system packages)
Follow the [Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/overview.html) for `.deb` / `.tar` / container options. The NGC container `nvcr.io/nvidia/tensorrt:<tag>` is the fastest way to get a known-good environment.
### Verify the install
```bash
python3 -c "import tensorrt; print(tensorrt.__version__)"
trtexec --help | head -5
```
---
## Path 1: ONNX → TensorRT
The ONNX path is the most portable way to bring models from PyTorch, TensorFlow, JAX, or any framework with an ONNX exporter.
> **ONNX compatibility note:** TensorRT does not support every ONNX operator or every ONNX opset version. Operator and opset coverage depends on the TensorRT release, so validate exported models with `trtexec` or Polygraphy and be prepared to update the exporter/opset, rewrite unsupported subgraphs, or provide custom plugins.
### 1. Export to ONNX
**PyTorch (dynamo exporter, preferred for TRT 11+):**
```python
import torch
model = MyModel().eval().cuda()
example = torch.randn(1, 3, 224, 224, device="cuda")
onnx_program = torch.onnx.export(
model,
(example,),
"model.onnx",
dynamo=True, # Use the dynamo exporter
dynamic_shapes=None, # Or specify with torch.export.Dim
)
```
**TensorFlow / Keras:** use `tf2onnx`:
```bash
python -m tf2onnx.convert --saved-model ./saved_model --output model.onnx --opset 20
```
### 2. (Optional) Simplify & sanitize
```bash
pip install onnx onnxsim polygraphy
python -m onnxsim model.onnx model.sim.onnx
polygraphy surgeon sanitize model.sim.onnx -o model.clean.onnx --fold-constants
```
### 3. Build a TensorRT engine
An ONNX file must be converted to a serialized TensorRT plan (`.plan` / `.engine`) before the runtime can execute it — the TensorRT runtime deserializes plans, it does **not** parse ONNX. The three options below are interchangeable front-ends that call the same `IBuilder` + `nvonnxparser::IParser` underneath; pick whichever fits your workflow.
> **Note on ONNX Runtime:** If you've seen "TensorRT executes ONNX directly," that refers to ONNX Runtime's `TensorrtExecutionProvider`, which lazy-builds a TRT engine internally on first call. That's ORT integrating TRT, not the TRT runtime itself.
**Option A — `trtexec` (CLI, fastest to try):**
```bash
trtexec \
--onnx=model.clean.onnx \
--saveEngine=model.plan \
--memPoolSize=workspace:4096 \
--fp16 # or --bf16, --int8, --fp8 (platform-dependent)
```
For dynamic shapes, add:
```bash
--minShapes=input:1x3x224x224 \
--optShapes=input:8x3x224x224 \
--maxShapes=input:16x3x224x224
```
**Option B — Python (`tensorrt.Builder` + `OnnxParser`):**
```python
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
network = builder.create_network(flags)
parser = trt.OnnxParser(network, logger)
with open("model.clean.onnx", "rb") as f:
assert parser.parse(f.read()), [parser.get_error(i) for i in range(parser.num_errors)]
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
serialized = builder.build_serialized_network(network, config)
open("model.plan", "wb").write(serialized)
```
**Option C — Polygraphy (scriptable, good for pipelines):**
```bash
polygraphy convert model.clean.onnx \
--convert-to trt \
--fp16 \
--workspace 4G \
--trt-min-shapes input:[1,3,224,224] \
--trt-opt-shapes input:[8,3,224,224] \
--trt-max-shapes input:[16,3,224,224] \
-o model.plan
```
C++ users: see `samples/sampleOnnxMNIST/` for the equivalent `IBuilder` + `IParser` flow.
### 4. Run the engine
See [Verifying an Engine](#verifying-an-engine).
### Common pitfalls
- **Unsupported op.** `trtexec` will name the op. Options: update your exporter/opset, rewrite the subgraph, or write a [custom plugin](#adding-a-custom-operator--plugin).
- **Shape inference failures.** Run `polygraphy inspect model model.onnx --show attrs` to confirm every tensor has a known rank.
- **Constant-folding surprises.** `polygraphy surgeon sanitize --fold-constants` often removes spurious dynamic axes introduced during export.
- **`INT64` tensors.** TRT will warn and cast to `INT32`; if values exceed `INT32` range, sanitize first.
---
## Path 2: Torch-TensorRT (PyTorch native)
The Dynamo frontend (`torch.compile(backend="tensorrt")`) is the **active, preferred** path. JIT/tracing-based `torch_tensorrt.compile` still works but receives minimal new investment.
### 1. Install
```bash
pip install --extra-index-url https://pypi.nvidia.com torch-tensorrt tensorrt-cu13
```
### 2. Compile (AOT — produces a standalone artifact)
```python
import torch
import torch_tensorrt as torch_trt
model = MyModel().eval().cuda().to(torch.float16)
example = torch.randn(1, 3, 224, 224, device="cuda", dtype=torch.float16)
trt_gm = torch_trt.dynamo.compile(
torch.export.export(model, (example,)),
inputs=[example],
enabled_precisions={torch.float16},
workspace_size=4 << 30,
)
# Save and reload
torch_trt.save(trt_gm, "model.ep", inputs=[example])
loaded = torch.export.load("model.ep").module()
```
### 3. Compile (JIT — first call triggers compilation)
```python
import torch
compiled = torch.compile(model, backend="tensorrt", options={"enabled_precisions": {torch.float16}})
out = compiled(example) # Triggers TRT compilation + cache on first call
```
### 4. Run
```python
with torch.no_grad():
y = trt_gm(example) # or compiled(example)
```
### Common pitfalls
- **Graph breaks** fall back to eager. Inspect with `TORCH_LOGS="graph_breaks"`. Eliminate them by lifting Python conditionals, avoiding `.item()` calls, and using `torch.cond` where possible.
- **Dynamic shapes** need explicit `torch.export.Dim(...)` annotations for AOT. JIT handles them but may recompile per shape.
- **Custom ops / plugins.** Torch-TRT converters live at `core/conversion/converters/` in the Torch-TensorRT repo. To add one, see the [Torch-TensorRT converter guide](https://docs.pytorch.org/TensorRT/contributors/writing_converters.html).
### Insights from historical issues
- `torch.export` in PyTorch 2.4+ is required for stable Dynamo AOT. Earlier versions fall back to torchscript tracing, which has been deprecated.
- Mixed precision: prefer `enabled_precisions={torch.float16}` over `torch.float32` unless a specific layer loses accuracy. For BF16 targets (Blackwell, Hopper), use `{torch.bfloat16}` and cast inputs accordingly.
---
## Path 3: Hugging Face Hub Models → TensorRT
> **For LLM generation, use [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) directly.** It is NVIDIA's active, production-grade path for Hugging Face LLMs — handles KV-cache, batching, paged attention, FP8/INT4 quantization, speculative decoding, and multi-GPU tensor/pipeline parallelism. The `optimum-nvidia` wrapper described in Option C below has not seen a release in over a year as of 2026-Q2 and is not recommended for new work.
For non-LLM Hugging Face Hub models (encoders, vision, diffusion components, speech), prefer Option A.
### Option A — Export to ONNX, then Path 1 (recommended default)
Most HF models export cleanly through `optimum`'s ONNX exporter:
```bash
pip install optimum-onnx # ONNX integration moved out of the `optimum` package in v2
optimum-cli export onnx \
--model google-bert/bert-base-uncased \
--task feature-extraction \
bert_onnx/
trtexec --onnx=bert_onnx/model.onnx --saveEngine=bert.plan --fp16
```
Then run through the standard [Path 1 build](#3-build-a-tensorrt-engine). This is the most durable HF → TRT path because it depends only on actively-maintained pieces (`optimum-onnx`, `trtexec`/Python builder).
### Option B — Torch-TensorRT
Covered in [Path 2](#path-2-torch-tensorrt-pytorch-native). Load with `transformers`, move to CUDA, and compile via `torch_tensorrt.dynamo.compile` or `torch.compile(backend="tensorrt")`. Good fit when you want to stay inside PyTorch and iterate quickly.
### Option C — `optimum-nvidia` (convenience wrapper; upstream is stale)
```bash
pip install optimum-nvidia
```
```python
from optimum.nvidia import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.2-1B",
use_fp8=False, # Or True on Hopper/Blackwell
)
out = model.generate(input_ids, max_new_tokens=128)
```
> **Status warning:** The last `optimum-nvidia` release (`v0.1.0b9`) shipped on 2025-01-21 and there have been no releases since. It still pins an older `tensorrt-llm` via `third-party/`. Verify the pinned versions match your TRT/CUDA stack before adopting, and prefer TensorRT-LLM directly (see callout at the top of this section) for anything production-facing.
### Common pitfalls
- **Tokenizer padding.** Hugging Face defaults to right-padding; some decoder models expect left-padding for generation. Mismatch produces silently-wrong logits.
- **KV-cache shapes.** For generative models, dynamic shapes along the sequence axis are mandatory. Use `optimum-nvidia` or hand-author shape profiles.
- **Diffusion pipelines** must be split by component (text encoder, UNet/DiT, VAE) — TRT cannot ingest the whole pipeline. See [supported_models.md](./supported_models.md) for component-by-component support.
---
## Path 4: Direct Network Definition API (C++/Python)
Use this only when neither ONNX nor PyTorch can express what you need (e.g., custom research architectures, ultra-tight control over layer choice).
```python
import numpy as np
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
x = network.add_input("x", trt.float16, (-1, 3, 224, 224))
w = trt.Weights(np.random.randn(64, 3, 7, 7).astype(np.float16))
conv = network.add_convolution_nd(x, 64, (7, 7), w, trt.Weights())
conv.stride_nd = (2, 2)
network.mark_output(conv.get_output(0))
config = builder.create_builder_config()
profile = builder.create_optimization_profile()
profile.set_shape("x", (1, 3, 224, 224), (8, 3, 224, 224), (16, 3, 224, 224))
config.add_optimization_profile(profile)
plan = builder.build_serialized_network(network, config)
open("model.plan", "wb").write(plan)
```
The C++ equivalent follows the same structure; see `samples/sampleINT8API` and `samples/python/refactored/2_construct_network_with_layer_apis/` for a runnable reference.
---
## Adding a Custom Operator / Plugin
When the importer reports an unsupported op:
1. **Check `tensorrt.IPluginRegistry`** — the op may already have a plugin you haven't loaded.
2. **Write a plugin** implementing `IPluginV3` (preferred for TRT 10+).
3. **Register it** via `REGISTER_TENSORRT_PLUGIN` (C++) or `trt.get_plugin_registry().register_creator(...)` (Python).
4. **Wire into ONNX** by naming the op `mydomain::MyPlugin` during export and supplying a matching plugin name.
5. **Torch-TensorRT custom converters** live in `core/conversion/converters/` — see the Torch-TRT docs.
A runnable reference: `samples/python/aliased_io_plugin/` in this repo.
For migration details, see the TensorRT Developer Guide section on [migrating V2 plugins to IPluginV3](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/extending-custom-layers.html#migrating-v2-plugins-to-ipluginv3).
---
## AI-Assisted Model Rewriting for Export
Hugging Face Hub models often don't export cleanly on the first try. Modern library code uses patterns — complex-number arithmetic, data-dependent control flow, non-tensor forward arguments, variable-length outputs — that `torch.export` / `torch.onnx.export` / Torch-TensorRT cannot trace directly. The usual workaround is *not* to modify the upstream library on disk, but to **monkey-patch equivalent, export-friendly variants at runtime** before export, then transparently swap the compiled module back into the pipeline.
This is repetitive, carefully-scoped work: read an upstream implementation, identify the one pattern that breaks the exporter, write a behaviorally-equivalent replacement, preserve everything the rest of the library expects. **It is exactly the kind of task where an AI coding agent pays off** — the agent reads hundreds of lines of upstream source (diffusers, transformers), proposes an equivalent formulation, and iterates against tracer errors without losing the thread. The worked example below is the output of this process for a non-trivial diffusion pipeline.
### Worked example: Qwen-Image (`diffusers`) → Torch-TensorRT AOT
The condensed pattern below comes from a Qwen-Image Torch-TensorRT AOT validation script and is written to be self-contained for OSS users.
The script compiles all three heavy components of `QwenImagePipeline` — text encoder, MMDiT transformer, VAE decoder — via `torch_tensorrt.dynamo.compile` and re-injects the compiled modules back into the pipeline. Five distinct export blockers had to be fixed, each representative of a broader class:
#### 1. Complex-number RoPE math → pre-compute real-valued cos/sin
Diffusers' `QwenEmbedRope` stores rotary-embedding frequencies as `torch.complex64` buffers and calls `torch.view_as_real(...)` inside the forward path. Torch-TensorRT's complex-graph detection cannot handle residual complex ops and segfaults.
**Fix:** pre-materialize real-valued `cos`/`sin` buffers on the module, patch the forward to read from them, **but keep the original complex buffers intact** (the exporter still probes them on entry):
```python
pos = torch.view_as_real(module.pos_freqs)
module._real_pos_cos = pos[..., 0].repeat_interleave(2, dim=-1).contiguous()
module._real_pos_sin = pos[..., 1].repeat_interleave(2, dim=-1).contiguous()
# Do NOT overwrite pos_freqs/neg_freqs — complex_graph_detection still reads them.
```
Where an agent helps: reading `diffusers/models/transformers/transformer_qwenimage.py` (≈1.5k lines), pinpointing the two `forward` methods that touch complex tensors, and deriving the real-valued equivalent without changing numerics.
#### 2. Non-tensor forward arguments → bake them into a wrapper
The transformer takes `img_shapes: list[list[tuple[int, int, int]]]`, which `torch.export` refuses to trace.
**Fix:** a thin wrapper stores the shape list as a constructor arg so the exported forward signature is pure tensors:
```python
class QwenImageTransformerAOTWrapper(nn.Module):
def __init__(self, transformer, img_shapes):
super().__init__()
self.transformer = transformer
self.img_shapes = img_shapes
def forward(self, hidden_states, encoder_hidden_states, encoder_hidden_states_mask, timestep):
return self.transformer(
hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states,
encoder_hidden_states_mask=encoder_hidden_states_mask, timestep=timestep,
img_shapes=self.img_shapes, return_dict=False)[0]
```
#### 3. HF-style output dataclasses → unwrap before export, re-wrap on the way back
Pipelines expect outputs like `Transformer2DModelOutput(sample=...)` or objects with `.hidden_states`. Torch-TRT needs plain tensors. Solve in two layers:
- **Export wrappers** (`TextEncoderAOTWrapper`, `VaeDecoderAOTWrapper`) return a bare tensor.
- **Reinjection proxies** (`CompiledTextEncoderProxy`, `CompiledTransformerProxy`, `CompiledVAEProxy`) expose every attribute the pipeline touches on the original module (`config`, `dtype`, `device`, `cache_context()`, etc.) and re-wrap outputs into the expected dataclass so the surrounding pipeline code does not notice the swap.
#### 4. Variable-length tokenization → force a static prompt shape
The default `encode_prompt` path slices per-sample by attention mask, producing variable-length hidden states. Torch-TRT requires static shapes for AOT.
**Fix:** monkey-patch `_get_qwen_prompt_embeds` onto the pipeline so tokenization always produces the same `[B, S]` shape the TRT text encoder was compiled with:
```python
pipe._get_qwen_prompt_embeds = MethodType(_get_qwen_prompt_embeds_fixed, pipe)
```
Where an agent helps: locating the (undocumented) method the pipeline dispatches to, reproducing the trimming/padding logic with a fixed `max_seq_len`, and ensuring dtype/device alignment with the compiled module.
#### 5. Memory-bound compilation → hint the resource partitioner
The full pipeline exceeds a single GPU's working set if everything is compiled greedily:
```python
import torch_tensorrt as torch_trt
# cpu_memory_budget is in bytes — tune to your host RAM headroom.
CPU_MEMORY_BUDGET_BYTES = 32 * 1024**3 # 32 GiB
torch_trt.compile(module, ir="dynamo", arg_inputs=inputs,
require_full_compilation=False,
enable_resource_partitioning=True,
cpu_memory_budget=CPU_MEMORY_BUDGET_BYTES,
truncate_double=True, optimization_level=1)
```
### Pattern takeaways
Apply this recipe to any non-trivial HF export:
1. **Run the naive export first.** Let the tracer/exporter fail and read the error carefully — the failing op/pattern tells you what to patch.
2. **Patch at runtime, not on disk.** Monkey-patch upstream modules from your export script so you never fork the library.
3. **Wrap for the exporter; proxy for the pipeline.** A `*AOTWrapper` unwraps HF outputs for export; a `Compiled*Proxy` re-wraps them on the way back and carries every attribute the downstream code reads.
4. **Preserve unobserved invariants.** When the exporter probes a buffer (e.g., `complex_graph_detection` reading `pos_freqs`), don't overwrite that buffer — add a parallel real-valued one.
5. **Iterate with an agent in the loop.** Each of the fixes above took one or two read-diagnose-patch cycles against upstream source; an agent can execute those cycles faster than a human skimming unfamiliar library code, while you review the diffs.
The end result for Qwen-Image: a pipeline whose heavy components all run on Torch-TensorRT, with zero changes to installed `diffusers` / `transformers` / `torch_tensorrt`, and a generated image that matches the eager pipeline's output on a fixed seed.
---
## Verifying an Engine
```bash
# Sanity-check performance and numerics
trtexec --loadEngine=model.plan --shapes=input:1x3x224x224 --verbose
# Side-by-side accuracy against the ONNX source
polygraphy run model.onnx --trt --onnxrt \
--atol 1e-3 --rtol 1e-3 --input-shapes input:[1,3,224,224]
```
For LLM-style generation, compare token-by-token against the reference implementation on a deterministic seed before trusting a new engine.
---
## Tools Reference
| Tool | What it does | Install |
|-------------|---------------------------------------------------------------|---------|
| `trtexec` | Build + run + profile engines from ONNX or serialized plans | Bundled with TRT |
| `polygraphy`| Inspect, sanitize, compare, and debug models at every stage | `pip install polygraphy` |
| `onnxsim` | Fold constants and simplify ONNX graphs | `pip install onnxsim` |
| `onnx-graphsurgeon` | Programmatic ONNX graph edits | `pip install onnx-graphsurgeon` |
| `nsys` / `ncu` | Runtime profiling and kernel analysis | NVIDIA CUDA Toolkit |
---
## Troubleshooting & Insights
Shared learnings from customer-reported issues. Extend this list liberally — it is the single most valuable section of this guide.
- **"Engine plan file is generated on an incompatible device"** — plans are not portable across compute capabilities. Rebuild on the deployment GPU, or target multiple SMs at build time.
- **Accuracy gap vs. framework** — start with `polygraphy run ... --onnxrt --trt --atol ...` to localize. If an FP16 engine diverges, try `--stronglyTyped` + explicit FP32 cast on the offending subgraph.
- **OOM during build** — lower `--memPoolSize=workspace:N` or disable tactic sources you don't need (`--tacticSources=-CUBLAS_LT`).
- **Slow first inference** — CUDA kernel JIT + plan deserialization cost is one-time. Warm up with ≥3 iterations before timing.
- **`IShapeLayer` / data-dependent shapes** — some patterns (e.g. `where(cond, x, y)` with dynamic output shape) require `IShapeLayer` + profile-shaped tensors. See the Developer Guide chapter on dynamic shapes.
documents/supported_models.md
+148
View File
@@ -0,0 +1,148 @@
# TensorRT Supported Model List
This verified model matrix pairs with [`import_workflows.md`](./import_workflows.md). For each model family, it lists the dtype(s) used during validation.
## Scope & Reading Guide
TensorRT is a general-purpose neural-network graph execution engine, not a model zoo. In principle **any NN architecture** can run on TensorRT as long as it is expressible through the workflows described in the [Import Workflows Guide](./import_workflows.md). The [Custom Plugin](./import_workflows.md#adding-a-custom-operator--plugin) section covers the escape hatch for ops TensorRT does not yet implement natively.
The table below is **not** an exhaustive support list. It is the subset of models NVIDIA has verified and benchmarked; we publish it so you know which configurations have a known-good baseline and where the current rough edges are. If your model is not listed, the expectation is still that it works — please file an issue if it does not.
### Reading the Tables
- **Dtype** lists the precision used for the verified baseline. Other precisions may also work.
- Component-split models (diffusion pipelines, speech models with encoder/decoder) list one row per validated component.
## Table of Contents
- [LLMs / Text Generation](#llms--text-generation)
- [Encoder-only NLP (BERT family, embeddings)](#encoder-only-nlp-bert-family-embeddings)
- [Vision Classification & Embeddings](#vision-classification--embeddings)
- [Speech / Audio](#speech--audio)
- [Diffusion Models](#diffusion-models)
- [Multimodal](#multimodal)
- [Legacy / TRT Sample Models](#legacy--trt-sample-models)
- [Requesting New Model Coverage](#requesting-new-model-coverage)
---
## LLMs / Text Generation
> **Preferred path for LLM generation:** [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) (KV-cache, paged attention, FP8/INT4, speculative decoding, tensor/pipeline parallelism). For production LLM serving, use TensorRT-LLM.
| Model | Dtype |
|---------------------------------|----------|
| `meta-llama/Llama-3.1-8B` | bfloat16 |
| `meta-llama/Llama-3.2-1B` | bfloat16 |
| `Qwen/Qwen3-0.6B` | bfloat16 |
| `deepseek-ai/Janus-Pro-7B` | bfloat16 |
> For TensorRT-LLM's own coverage, see the [TensorRT-LLM model support matrix](https://github.com/NVIDIA/TensorRT-LLM#model-zoo).
---
## Encoder-only NLP (BERT family, embeddings)
| Model | Dtype |
|----------------------------------------------------|---------|
| `google-bert/bert-base-uncased` | float32 |
| `google-bert/bert-base-multilingual-cased` | float16 |
| `FacebookAI/roberta-base` | float32 |
| `FacebookAI/roberta-large` | float32 |
| `FacebookAI/xlm-roberta-base` | float32 |
| `distilbert/distilbert-base-uncased` | float32 |
| `sentence-transformers/all-MiniLM-L6-v2` | float32 |
| `sentence-transformers/all-mpnet-base-v2` | float32 |
| `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` | float32 |
| `BAAI/bge-base-en-v1.5` | float32 |
| `nlpaueb/legal-bert-base-uncased` | float32 |
---
## Vision Classification & Embeddings
| Model | Dtype |
|---------------------------------------------|---------|
| `torchvision/resnet50` | float32 |
| `timm/mobilenetv3_small_100.lamb_in1k` | float32 |
| `trpakov/vit-face-expression` | float32 |
| `openai/clip-vit-base-patch32` | float32 |
| `openai/clip-vit-large-patch14` | float32 |
| `facebook/dinov2-base` | float32 |
| `Falconsai/nsfw_image_detection` | float32 |
| `dima806/fairface_age_image_detection` | float32 |
---
## Speech / Audio
| Model (Component) | Dtype |
|---------------------------------------------------|---------|
| `openai/whisper-large-v3-turbo` (Encoder) | float32 |
| `openai/whisper-large-v3-turbo` (Decoder) | float32 |
| `openai/whisper-large-v3` (Encoder) | float32 |
| `openai/whisper-large-v3` (Decoder) | float32 |
| `laion/clap-htsat-fused` | float32 |
| `sesame/csm-1b` (Backbone) | float32 |
| `neuphonic/neutts-air` | float32 |
| `LiquidAI/LFM2-Audio-1.5B` | float32 |
---
## Diffusion Models
Diffusion pipelines are evaluated per component (Text Encoder / UNet or DiT / VAE) because TRT does not ingest the pipeline object directly.
| Pipeline (Component) | Dtype |
|------------------------------------------------------------|----------|
| `stabilityai/sd-turbo` | float16 |
| `stabilityai/sdxl-turbo` (UNet) | float16 |
| `stabilityai/sdxl-turbo` (VAE / Text Encoders) | mixed |
| `stabilityai/stable-diffusion-xl-base-1.0` | float16 |
| `CompVis/stable-diffusion-v1-4` | float16 |
| `stable-diffusion-v1-5/stable-diffusion-v1-5` | float16 |
| `stabilityai/stable-diffusion-2-1` | float16 |
| `playgroundai/playground-v2.5-1024px-aesthetic` | float16 |
| `dataautogpt3/ProteusV0.3` | float16 |
| `black-forest-labs/FLUX.2-dev` (Text Encoder) | bfloat16 |
| `black-forest-labs/FLUX.2-dev` (DiT) | bfloat16 |
| `black-forest-labs/FLUX.2-dev` (VAE) | float16 |
| `black-forest-labs/FLUX.1-schnell` (DiT / TextEnc / VAE) | mixed |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (Text Encoder) | float16 |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (VAE) | float16 |
| `Qwen/Qwen-Image` (Text Encoder) | bfloat16 |
| `Qwen/Qwen-Image` (DiT / VAE) | bfloat16 |
| `stabilityai/stable-diffusion-3-medium-diffusers` | bfloat16 |
| `stabilityai/stable-diffusion-3.5-medium` / `3.5-large` | mixed |
| `HiDream-ai/HiDream-I1-Full` | bfloat16 |
| `stabilityai/stable-video-diffusion-img2vid-xt` | float16 |
---
## Multimodal
| Model | Dtype |
|-----------------------------------|----------|
| `openai/clip-vit-base-patch32` | float32 |
| `deepseek-ai/Janus-Pro-7B` | bfloat16 |
| `Datadog/Toto-Open-Base-1.0` | float32 |
---
## Legacy / TRT Sample Models
TensorRT ships hand-validated C++/Python samples for these classic architectures and workflows:
- MNIST digit classifiers, model parsing, dynamic-shape, plugin, and safe-runtime samples — see `samples/` in this repo.
---
## Requesting New Model Coverage
File a GitHub issue with:
1. The Hugging Face ID or model source URL.
2. The target dtype (fp32 / fp16 / bf16 / fp8 / int8 / int4).
3. Any framework-level working example (helps us reproduce quickly).
The maintainers will benchmark the model and extend this table — no external contributor action needed for the benchmark step.
documents/tensorrt_roadmap_2026q3.pdf
BIN
View File
Binary file not shown.
include/NvInfer.h
+761 -1487
View File
File diff suppressed because it is too large Load Diff
include/NvInferImpl.h
+24 -115
View File
@@ -26,22 +26,14 @@
namespace nvinfer1
{
class ILogger;
namespace v_1_0
{
class ILogger;
class IProgressMonitor;
} // namespace v_1_0
using ILogger = v_1_0::ILogger;
using IProgressMonitor = v_1_0::IProgressMonitor;
#if !STRIP_TRT_RTX_INTERNAL_API
namespace v_1_0
{
class IAlgorithmSelector;
} // namespace v_1_0
using IAlgorithmSelector = v_1_0::IAlgorithmSelector;
#endif // !STRIP_TRT_RTX_INTERNAL_API
namespace v_1_0
{
class IProfiler;
@@ -61,12 +53,6 @@ class IDebugListener;
using IDebugListener = v_1_0::IDebugListener;
class IActivationLayer;
#if !STRIP_TRT_RTX_INTERNAL_API
class IAlgorithm;
class IAlgorithmContext;
class IAlgorithmIOInfo;
class IAlgorithmVariant;
#endif // !STRIP_TRT_RTX_INTERNAL_API
class IAssertionLayer;
class IAttention;
class IBuilder;
@@ -94,9 +80,6 @@ class ICastLayer;
class IIfConditional;
class IIfConditionalInputLayer;
class IIfConditionalOutputLayer;
#if !STRIP_TRT_RTX_INTERNAL_API
class IInt8Calibrator;
#endif // !STRIP_TRT_RTX_INTERNAL_API
class IIteratorLayer;
class IKVCacheUpdateLayer;
class ILayer;
@@ -180,12 +163,11 @@ struct Permutation;
class Weights;
enum class ActivationType : int32_t;
enum class AttentionIOForm : int32_t;
enum class AttentionNormalizationOp : int32_t;
enum class BoundingBoxFormat : int32_t;
enum class CausalMaskKind : int32_t;
enum class BuilderFlag : int32_t;
#if !STRIP_TRT_RTX_INTERNAL_API
enum class CalibrationAlgoType : int32_t;
#endif // !STRIP_TRT_RTX_INTERNAL_API
enum class CumulativeOperation : int32_t;
enum class DeviceType : int32_t;
enum class DimensionOperation : int32_t;
@@ -205,7 +187,6 @@ enum class OptProfileSelector : int32_t;
enum class PaddingMode : int32_t;
enum class PoolingType : int32_t;
enum class ProfilingVerbosity : int32_t;
enum class QuantizationFlag : int32_t;
enum class ReduceOperation : int32_t;
enum class CollectiveOperation : int32_t;
enum class ResizeCoordinateTransformation : int32_t;
@@ -234,7 +215,6 @@ using TacticSources = uint32_t;
using TensorFormats = uint32_t;
using BuilderFlags = uint32_t;
using NetworkDefinitionCreationFlags = uint32_t;
using QuantizationFlags = uint32_t;
using TempfileControlFlags = uint32_t;
using SerializationFlags = uint32_t;
@@ -288,7 +268,6 @@ class VRuntime : public VRoot
public:
virtual IRuntime* getPImpl() noexcept = 0;
virtual nvinfer1::ICudaEngine* deserializeCudaEngine(void const* blob, std::size_t size) noexcept = 0;
virtual nvinfer1::ICudaEngine* deserializeCudaEngine(IStreamReader& streamReader) noexcept = 0;
virtual void setDLACore(int32_t dlaCore) noexcept = 0;
virtual int32_t getDLACore() const noexcept = 0;
virtual int32_t getNbDLACores() const noexcept = 0;
@@ -319,10 +298,6 @@ public:
virtual bool refitCudaEngine() noexcept = 0;
virtual int32_t getMissing(int32_t size, char const** layerNames, WeightsRole* roles) noexcept = 0;
virtual int32_t getAll(int32_t size, char const** layerNames, WeightsRole* roles) noexcept = 0;
virtual bool setDynamicRange(char const* tensorName, float min, float max) noexcept = 0;
virtual float getDynamicRangeMin(char const* tensorName) const noexcept = 0;
virtual float getDynamicRangeMax(char const* tensorName) const noexcept = 0;
virtual int32_t getTensorsWithDynamicRange(int32_t size, char const** tensorNames) const noexcept = 0;
virtual void setErrorRecorder(IErrorRecorder* recorder) noexcept = 0;
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
virtual bool setNamedWeights(char const* name, Weights weights) noexcept = 0;
@@ -346,10 +321,7 @@ class VOptimizationProfile : public VRoot
public:
virtual bool setDimensions(char const* inputName, OptProfileSelector select, Dims const& dims) noexcept = 0;
virtual Dims getDimensions(char const* inputName, OptProfileSelector select) const noexcept = 0;
virtual bool setShapeValues(
char const* inputName, OptProfileSelector select, int32_t const* values, int32_t nbValues) noexcept = 0;
virtual int32_t getNbShapeValues(char const* inputName) const noexcept = 0;
virtual int32_t const* getShapeValues(char const* inputName, OptProfileSelector select) const noexcept = 0;
virtual bool setExtraMemoryTarget(float target) noexcept = 0;
virtual float getExtraMemoryTarget() const noexcept = 0;
virtual bool isValid() const noexcept = 0;
@@ -367,17 +339,12 @@ public:
virtual int32_t getNbLayers() const noexcept = 0;
virtual IHostMemory* serialize() const noexcept = 0;
virtual IExecutionContext* createExecutionContext(ExecutionContextAllocationStrategy strategy) noexcept = 0;
virtual IExecutionContext* createExecutionContextWithoutDeviceMemory() noexcept = 0;
virtual size_t getDeviceMemorySize() const noexcept = 0;
virtual bool isRefittable() const noexcept = 0;
virtual char const* getName() const noexcept = 0;
virtual int32_t getNbOptimizationProfiles() const noexcept = 0;
virtual int32_t const* getProfileTensorValues(
char const* tensorName, int32_t profileIndex, OptProfileSelector select) const noexcept = 0;
virtual EngineCapability getEngineCapability() const noexcept = 0;
virtual void setErrorRecorder(IErrorRecorder* recorder) noexcept = 0;
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
virtual bool hasImplicitBatchDimension() const noexcept = 0;
virtual TacticSources getTacticSources() const noexcept = 0;
virtual ProfilingVerbosity getProfilingVerbosity() const noexcept = 0;
virtual IEngineInspector* createEngineInspector() const noexcept = 0;
@@ -407,12 +374,8 @@ public:
virtual ISerializationConfig* createSerializationConfig() noexcept = 0;
virtual IHostMemory* serializeWithConfig(ISerializationConfig& config) const noexcept = 0;
virtual size_t getDeviceMemorySizeForProfile(int32_t profileIndex) const noexcept = 0;
virtual IRefitter* createRefitter(ILogger& logger) noexcept = 0;
virtual bool setWeightStreamingBudget(int64_t gpuMemoryBudget) noexcept = 0;
virtual int64_t getWeightStreamingBudget() const noexcept = 0;
virtual int64_t getMinimumWeightStreamingBudget() const noexcept = 0;
virtual int64_t getStreamableWeightsSize() const noexcept = 0;
virtual bool isDebugTensor(char const* name) const noexcept = 0;
@@ -449,7 +412,6 @@ public:
virtual void setDeviceMemory(void* memory) noexcept = 0;
virtual int32_t getOptimizationProfile() const noexcept = 0;
virtual bool allInputDimensionsSpecified() const noexcept = 0;
virtual bool allInputShapesSpecified() const noexcept = 0;
virtual void setErrorRecorder(IErrorRecorder* recorder) noexcept = 0;
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
virtual bool executeV2(void* const* bindings) noexcept = 0;
@@ -511,19 +473,13 @@ public:
virtual char const* getName() const noexcept = 0;
virtual void setDimensions(Dims const& dimensions) noexcept = 0;
virtual Dims getDimensions() const noexcept = 0;
virtual void setType(DataType type) noexcept = 0;
virtual DataType getType() const noexcept = 0;
virtual bool setDynamicRange(float min, float max) noexcept = 0;
virtual bool isNetworkInput() const noexcept = 0;
virtual bool isNetworkOutput() const noexcept = 0;
virtual void setBroadcastAcrossBatch(bool broadcastAcrossBatch) noexcept = 0;
virtual bool getBroadcastAcrossBatch() const noexcept = 0;
virtual TensorLocation getLocation() const noexcept = 0;
virtual void setLocation(TensorLocation location) noexcept = 0;
virtual bool dynamicRangeIsSet() const noexcept = 0;
virtual void resetDynamicRange() noexcept = 0;
virtual float getDynamicRangeMin() const noexcept = 0;
virtual float getDynamicRangeMax() const noexcept = 0;
virtual void setAllowedFormats(TensorFormats formats) noexcept = 0;
virtual TensorFormats getAllowedFormats() const noexcept = 0;
virtual bool isShapeTensor() const noexcept = 0;
@@ -543,14 +499,7 @@ public:
virtual int32_t getNbOutputs() const noexcept = 0;
virtual ITensor* getOutput(int32_t index) const noexcept = 0;
virtual void setInput(int32_t index, ITensor& tensor) noexcept = 0;
virtual void setPrecision(DataType dataType) noexcept = 0;
virtual DataType getPrecision() const noexcept = 0;
virtual bool precisionIsSet() const noexcept = 0;
virtual void resetPrecision() noexcept = 0;
virtual void setOutputType(int32_t index, DataType dataType) noexcept = 0;
virtual DataType getOutputType(int32_t index) const noexcept = 0;
virtual bool outputTypeIsSet(int32_t index) const noexcept = 0;
virtual void resetOutputType(int32_t index) noexcept = 0;
virtual void setMetadata(char const* docString) noexcept = 0;
virtual char const* getMetadata() const noexcept = 0;
virtual bool setNbRanks(int32_t nbRanks) noexcept = 0;
@@ -965,8 +914,8 @@ public:
TRT_NODISCARD virtual char const* getName() const noexcept = 0;
TRT_NODISCARD virtual bool setNormalizationOperation(AttentionNormalizationOp op) noexcept = 0;
TRT_NODISCARD virtual AttentionNormalizationOp getNormalizationOperation() const noexcept = 0;
TRT_NODISCARD virtual bool setCausal(bool isCausal) noexcept = 0;
TRT_NODISCARD virtual bool getCausal() const noexcept = 0;
TRT_DEPRECATED virtual bool setCausal(bool isCausal) noexcept = 0;
TRT_DEPRECATED virtual bool getCausal() const noexcept = 0;
TRT_NODISCARD virtual bool setMask(ITensor& mask) noexcept = 0;
TRT_NODISCARD virtual ITensor* getMask() const noexcept = 0;
TRT_NODISCARD virtual bool setDecomposable(bool decomposable) noexcept = 0;
@@ -979,6 +928,16 @@ public:
TRT_NODISCARD virtual char const* getMetadata() const noexcept = 0;
TRT_NODISCARD virtual bool setNbRanks(int32_t nbRanks) noexcept = 0;
TRT_NODISCARD virtual int32_t getNbRanks() const noexcept = 0;
TRT_NODISCARD virtual bool setCausalKind(CausalMaskKind kind) noexcept = 0;
TRT_NODISCARD virtual CausalMaskKind getCausalKind() const noexcept = 0;
TRT_NODISCARD virtual bool setQueryForm(AttentionIOForm form) noexcept = 0;
TRT_NODISCARD virtual AttentionIOForm getQueryForm() const noexcept = 0;
TRT_NODISCARD virtual bool setKeyValueForm(AttentionIOForm form) noexcept = 0;
TRT_NODISCARD virtual AttentionIOForm getKeyValueForm() const noexcept = 0;
TRT_NODISCARD virtual bool setQueryLengths(ITensor* lengths) noexcept = 0;
TRT_NODISCARD virtual ITensor* getQueryLengths() const noexcept = 0;
TRT_NODISCARD virtual bool setKeyValueLengths(ITensor* lengths) noexcept = 0;
TRT_NODISCARD virtual ITensor* getKeyValueLengths() const noexcept = 0;
}; // class VAttention
class VSelectLayer : public VRoot
@@ -1113,8 +1072,6 @@ public:
virtual uint32_t getAxes() const noexcept = 0;
virtual void setNbGroups(int64_t nbGroups) noexcept = 0;
virtual int64_t getNbGroups() const noexcept = 0;
virtual void setComputePrecision(DataType type) noexcept = 0;
virtual DataType getComputePrecision() const noexcept = 0;
virtual bool isV2() const noexcept = 0;
}; // class VNormalizationLayer
@@ -1152,6 +1109,10 @@ class VKVCacheUpdateLayer : public VRoot
public:
TRT_NODISCARD virtual bool setCacheMode(KVCacheMode cacheMode) noexcept = 0;
TRT_NODISCARD virtual KVCacheMode getCacheMode() const noexcept = 0;
TRT_NODISCARD virtual bool setUpdateForm(AttentionIOForm form) noexcept = 0;
TRT_NODISCARD virtual AttentionIOForm getUpdateForm() const noexcept = 0;
TRT_NODISCARD virtual bool setUpdateLengths(ITensor* lengths) noexcept = 0;
TRT_NODISCARD virtual ITensor* getUpdateLengths() const noexcept = 0;
}; // class VKVCacheUpdateLayer
class VMoELayer : public VRoot
@@ -1235,13 +1196,10 @@ public:
virtual IResizeLayer* addResize(ITensor& input) noexcept = 0;
virtual ILoop* addLoop() noexcept = 0;
virtual ISelectLayer* addSelect(ITensor& condition, ITensor& thenInput, ITensor& elseInput) noexcept = 0;
virtual IFillLayer* addFill(Dims const& dimensions, FillOperation op) noexcept = 0;
virtual IPaddingLayer* addPaddingNd(ITensor& input, Dims const& prePadding, Dims const& postPadding) noexcept = 0;
virtual bool setWeightsName(Weights weights, char const* name) noexcept = 0;
virtual void setErrorRecorder(IErrorRecorder* recorder) noexcept = 0;
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
virtual IDequantizeLayer* addDequantize(ITensor& input, ITensor& scale) noexcept = 0;
virtual IQuantizeLayer* addQuantize(ITensor& input, ITensor& scale) noexcept = 0;
virtual IGatherLayer* addGatherV2(ITensor& data, ITensor& indices, GatherMode mode) noexcept = 0;
virtual IIfConditional* addIfConditional() noexcept = 0;
virtual IScatterLayer* addScatter(ITensor& data, ITensor& indices, ITensor& updates, ScatterMode mode) noexcept = 0;
@@ -1280,7 +1238,7 @@ public:
virtual INonZeroLayer* addNonZeroV2(ITensor& input, DataType indicesType) noexcept = 0;
virtual INMSLayer* addNMSV2(
ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass, DataType indicesType) noexcept = 0;
virtual IAttention* addAttention(
TRT_DEPRECATED virtual IAttention* addAttention(
ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp, bool isCausal) noexcept = 0;
virtual IRotaryEmbeddingLayer* addRotaryEmbedding(ITensor& input, ITensor& cosCache, ITensor& sinCache,
bool interleaved, int32_t rotaryEmbeddingDim) noexcept = 0;
@@ -1294,45 +1252,10 @@ public:
ITensor& hiddenStates, ITensor& selectedExpertsForTokens, ITensor& scoresForSelectedExperts) noexcept = 0;
virtual IDistCollectiveLayer* addDistCollective(ITensor& input, CollectiveOperation distCollectiveOp,
ReduceOperation reduceOp, int64_t root, int64_t* groups, int64_t groupSize) noexcept = 0;
virtual IAttention* addAttentionV2(ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp,
CausalMaskKind causalKind) noexcept = 0;
};
#if !STRIP_TRT_RTX_INTERNAL_API
class VAlgorithmIOInfo : public VRoot
{
public:
virtual DataType getDataType() const noexcept = 0;
virtual Dims getStrides() const noexcept = 0;
virtual int64_t getVectorizedDim() const noexcept = 0;
virtual int64_t getComponentsPerElement() const noexcept = 0;
};
class VAlgorithmVariant : public VRoot
{
public:
virtual int64_t getImplementation() const noexcept = 0;
virtual int64_t getTactic() const noexcept = 0;
};
class VAlgorithmContext : public VRoot
{
public:
virtual char const* getName() const noexcept = 0;
virtual Dims getDimensions(int32_t index, OptProfileSelector select) const noexcept = 0;
virtual int32_t getNbInputs() const noexcept = 0;
virtual int32_t getNbOutputs() const noexcept = 0;
};
class VAlgorithm : public VRoot
{
public:
virtual IAlgorithmVariant const& getAlgorithmVariant() const noexcept = 0;
virtual float getTimingMSec() const noexcept = 0;
virtual std::size_t getWorkspaceSize() const noexcept = 0;
virtual IAlgorithmIOInfo const* getAlgorithmIOInfoByIndex(int32_t index) const noexcept = 0;
};
#endif // !STRIP_TRT_RTX_INTERNAL_API
class VTimingCache : public VRoot
{
public:
@@ -1351,8 +1274,6 @@ public:
virtual int32_t getAvgTimingIterations() const noexcept = 0;
virtual void setEngineCapability(EngineCapability capability) noexcept = 0;
virtual EngineCapability getEngineCapability() const noexcept = 0;
virtual void setInt8Calibrator(IInt8Calibrator* calibrator) noexcept = 0;
virtual IInt8Calibrator* getInt8Calibrator() const noexcept = 0;
virtual void setFlags(BuilderFlags builderFlags) noexcept = 0;
virtual BuilderFlags getFlags() const noexcept = 0;
virtual void clearFlag(BuilderFlag builderFlag) noexcept = 0;
@@ -1374,15 +1295,6 @@ public:
virtual int32_t getNbOptimizationProfiles() const noexcept = 0;
virtual void setProfilingVerbosity(ProfilingVerbosity verbosity) noexcept = 0;
virtual ProfilingVerbosity getProfilingVerbosity() const noexcept = 0;
virtual void setAlgorithmSelector(IAlgorithmSelector* selector) noexcept = 0;
virtual IAlgorithmSelector* getAlgorithmSelector() const noexcept = 0;
virtual bool setCalibrationProfile(IOptimizationProfile const* profile) noexcept = 0;
virtual IOptimizationProfile const* getCalibrationProfile() noexcept = 0;
virtual void setQuantizationFlags(QuantizationFlags flags) noexcept = 0;
virtual QuantizationFlags getQuantizationFlags() const noexcept = 0;
virtual void clearQuantizationFlag(QuantizationFlag flag) noexcept = 0;
virtual void setQuantizationFlag(QuantizationFlag flag) noexcept = 0;
virtual bool getQuantizationFlag(QuantizationFlag flag) const noexcept = 0;
virtual bool setTacticSources(TacticSources tacticSources) noexcept = 0;
virtual TacticSources getTacticSources() const noexcept = 0;
virtual nvinfer1::ITimingCache* createTimingCache(void const* blob, std::size_t size) const noexcept = 0;
@@ -1399,7 +1311,7 @@ public:
virtual void setPluginsToSerialize(char const* const* paths, int32_t nbPaths) noexcept = 0;
virtual char const* getPluginToSerialize(int32_t index) const noexcept = 0;
virtual int32_t getNbPluginsToSerialize() const noexcept = 0;
virtual void setMaxAuxStreams(int32_t nbStreams) noexcept = 0;
virtual bool setMaxAuxStreams(int32_t nbStreams) noexcept = 0;
virtual int32_t getMaxAuxStreams() const noexcept = 0;
virtual void setProgressMonitor(IProgressMonitor* monitor) noexcept = 0;
virtual IProgressMonitor* getProgressMonitor() const noexcept = 0;
@@ -1428,8 +1340,6 @@ public:
class VBuilder : public VRoot
{
public:
virtual bool platformHasFastFp16() const noexcept = 0;
virtual bool platformHasFastInt8() const noexcept = 0;
virtual int32_t getMaxDLABatchSize() const noexcept = 0;
virtual int32_t getNbDLACores() const noexcept = 0;
virtual void setGpuAllocator(IGpuAllocator* allocator) noexcept = 0;
@@ -1439,7 +1349,6 @@ public:
virtual void setErrorRecorder(IErrorRecorder* recorder) noexcept = 0;
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
virtual void reset() noexcept = 0;
virtual bool platformHasTf32() const noexcept = 0;
virtual nvinfer1::IHostMemory* buildSerializedNetwork(
INetworkDefinition& network, IBuilderConfig& config) noexcept = 0;
virtual bool isNetworkSupported(INetworkDefinition const& network, IBuilderConfig const& config) const noexcept = 0;
include/NvInferPluginBase.h
+1 -1
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
include/NvInferPluginUtils.h
+1 -70
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -96,52 +96,6 @@ struct GridAnchorParameters
float variance[4]; //!< Variance for adjusting the prior boxes.
};
//!
//! \enum CodeTypeSSD
//!
//! \brief The type of encoding used for decoding the bounding boxes and loc_data.
//!
//! \deprecated Deprecated in TensorRT 10.0. DetectionOutput plugin is deprecated.
//!
enum class CodeTypeSSD : int32_t
{
CORNER TRT_DEPRECATED_ENUM = 0, //!< Use box corners.
CENTER_SIZE TRT_DEPRECATED_ENUM = 1, //!< Use box centers and size.
CORNER_SIZE TRT_DEPRECATED_ENUM = 2, //!< Use box centers and size.
TF_CENTER TRT_DEPRECATED_ENUM = 3 //!< Use box centers and size but flip x and y coordinates.
};
//!
//! \struct DetectionOutputParameters
//!
//! \brief The DetectionOutput plugin layer generates the detection output
//! based on location and confidence predictions by doing non maximum suppression.
//!
//! This plugin first decodes the bounding boxes based on the anchors generated.
//! It then performs non_max_suppression on the decoded bounding boxes.
//! DetectionOutputParameters defines a set of parameters for creating the DetectionOutput plugin layer.
//!
//! \deprecated Deprecated in TensorRT 10.0. DetectionOutput plugin is deprecated.
//!
struct TRT_DEPRECATED DetectionOutputParameters
{
bool shareLocation; //!< If true, bounding box are shared among different classes.
bool varianceEncodedInTarget; //!< If true, variance is encoded in target.
//!< Otherwise we need to adjust the predicted offset accordingly.
int32_t backgroundLabelId; //!< Background label ID. If there is no background class, set it as -1.
int32_t numClasses; //!< Number of classes to be predicted.
int32_t topK; //!< Number of boxes per image with top confidence scores that are fed
//!< into the NMS algorithm.
int32_t keepTopK; //!< Number of total bounding boxes to be kept per image after NMS step.
float confidenceThreshold; //!< Only consider detections whose confidences are larger than a threshold.
float nmsThreshold; //!< Threshold to be used in NMS.
CodeTypeSSD codeType; //!< Type of coding method for bbox.
int32_t inputOrder[3]; //!< Specifies the order of inputs {loc_data, conf_data, priorbox_data}.
bool confSigmoid; //!< Set to true to calculate sigmoid of confidence scores.
bool isNormalized; //!< Set to true if bounding box data is normalized by the network.
bool isBatchAgnostic{true}; //!< Defaults to true. Set to false if prior boxes are unique per batch.
};
//!
//! \brief When performing yolo9000, softmaxTree is helping to do softmax on confidence scores,
//! for element to get the precise classification through word-tree structured classification definition.
@@ -175,29 +129,6 @@ struct RegionParameters
softmaxTree* smTree; //!< Helping structure to do softmax on confidence scores.
};
//!
//! \brief The NMSParameters are used by the BatchedNMSPlugin for performing
//! the non_max_suppression operation over boxes for object detection networks.
//!
//! \deprecated Deprecated in TensorRT 10.0. BatchedNMSPlugin plugin is deprecated.
//!
struct TRT_DEPRECATED NMSParameters
{
bool shareLocation; //!< If set to true, the boxes inputs are shared across all classes.
//!< If set to false, the boxes input should account for per class box data.
int32_t backgroundLabelId; //!< Label ID for the background class.
//!< If there is no background class, set it as -1
int32_t numClasses; //!< Number of classes in the network.
int32_t topK; //!< Number of bounding boxes to be fed into the NMS step.
int32_t keepTopK; //!< Number of total bounding boxes to be kept per image after NMS step.
//!< Should be less than or equal to the topK value.
float scoreThreshold; //!< Scalar threshold for score (low scoring boxes are removed).
float iouThreshold; //!< A scalar threshold for IOU (new boxes that have high IOU overlap
//!< with previously selected boxes are removed).
bool isNormalized; //!< Set to false, if the box coordinates are not normalized,
//!< i.e. not in the range [0,1]. Defaults to false.
};
} // namespace plugin
} // namespace nvinfer1
include/NvInferRuntime.h
+157 -517
View File
File diff suppressed because it is too large Load Diff
include/NvInferRuntimeBase.h
+13 -10
View File
@@ -115,21 +115,24 @@ using char_t = char;
//! This type is widely used in automotive safety context.
using AsciiChar = char_t;
//! Forward declare IErrorRecorder for use in other interfaces.
//! Forward declare IErrorRecorder and ILogger for use in other interfaces.
namespace v_1_0
{
class IErrorRecorder;
class ILogger;
} // namespace v_1_0
using IErrorRecorder = v_1_0::IErrorRecorder;
using ILogger = v_1_0::ILogger;
namespace impl
{
//! Declaration of EnumMaxImpl struct to store maximum number of elements in an enumeration type.
//! Declaration of EnumMaxImpl struct to store the exclusive upper bound of an enumeration type.
template <typename T>
struct EnumMaxImpl;
} // namespace impl
//! Maximum number of elements in an enumeration type.
//! One greater than the maximum value of enumeration type T.
//! For example, if the highest enumerator in T has value 5, then EnumMax<T>() returns 6.
template <typename T>
constexpr int32_t EnumMax() noexcept
{
@@ -196,11 +199,11 @@ enum class DataType : int32_t
namespace impl
{
//! Maximum number of elements in DataType enum. \see DataType
//! One greater than the maximum value of DataType enum. \see DataType
template <>
struct EnumMaxImpl<DataType>
{
//! Declaration of kVALUE that represents the maximum number of elements in the DataType enum.
//! One greater than the maximum value of DataType enum.
static constexpr int32_t kVALUE = 12;
};
} // namespace impl
@@ -261,11 +264,11 @@ enum class APILanguage : int32_t
namespace impl
{
//! Maximum number of elements in APILanguage enum. \see APILanguage
//! One greater than the maximum value of APILanguage enum. \see APILanguage
template <>
struct EnumMaxImpl<APILanguage>
{
//! Declaration of kVALUE that represents the maximum number of elements in the APILanguage enum.
//! One greater than the maximum value of APILanguage enum.
static constexpr int32_t kVALUE = 2;
};
} // namespace impl
@@ -401,7 +404,7 @@ enum class ErrorCode : int32_t
namespace impl
{
//! Maximum number of elements in ErrorCode enum. \see ErrorCode
//! One greater than the maximum value of ErrorCode enum. \see ErrorCode
template <>
struct EnumMaxImpl<ErrorCode>
{
@@ -669,11 +672,11 @@ enum class TensorIOMode : int32_t
namespace impl
{
//! Maximum number of elements in TensorIOMode enum. \see TensorIOMode
//! One greater than the maximum value of TensorIOMode enum. \see TensorIOMode
template <>
struct EnumMaxImpl<TensorIOMode>
{
// Declaration of kVALUE that represents maximum number of elements in TensorIOMode enum
// One greater than the maximum value of TensorIOMode enum
static constexpr int32_t kVALUE = 3;
};
} // namespace impl
include/NvInferRuntimeCommon.h
+4 -73
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -60,58 +60,6 @@ public:
//!
using PluginLibraryHandle = void*;
//!
//! \brief Register a plugin creator implementing IPluginCreator. Returns false if any plugin creator with the same
//! name, version or namespace is already registered.
//!
//! \warning The string pluginNamespace must be 1024 bytes or less including the NULL terminator and must be NULL
//! terminated.
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: Yes; calls to this method will be synchronized by a mutex.
//!
//! \deprecated Deprecated in TensorRT 10.0. Superseded by
//! IPluginRegistry::registerCreator(IPluginCreatorInterface&, AsciiChar const* const).
//!
TRT_DEPRECATED virtual bool registerCreator(
IPluginCreator& creator, AsciiChar const* const pluginNamespace) noexcept = 0;
//!
//! \brief Return all the registered plugin creators and the number of
//! registered plugin creators. Returns nullptr if none found.
//!
//! \warning If any plugin creators are registered or deregistered after calling this function, the returned pointer
//! is not guaranteed to be valid thereafter.
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: No
//!
//! \deprecated Deprecated in TensorRT 10.0. Superseded by IPluginRegistry::getAllCreators(int32_t* const).
//!
TRT_DEPRECATED virtual IPluginCreator* const* getPluginCreatorList(int32_t* const numCreators) const noexcept = 0;
//!
//! \brief Return plugin creator based on plugin name, version, and
//! namespace associated with plugin during network creation.
//!
//! \warning The strings pluginName, pluginVersion, and pluginNamespace must be 1024 bytes or less including the
//! NULL terminator and must be NULL terminated.
//!
//! \warning Returns nullptr if a plugin creator with matching name, version, and namespace is found, but is not a
//! descendent of IPluginCreator
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: Yes
//!
//! \deprecated Deprecated in TensorRT 10.0. Superseded by IPluginRegistry::getCreator(AsciiChar const* const,
//! AsciiChar const* const, AsciiChar const* const).
//!
TRT_DEPRECATED virtual IPluginCreator* getPluginCreator(AsciiChar const* const pluginName,
AsciiChar const* const pluginVersion, AsciiChar const* const pluginNamespace = "") noexcept = 0;
// @cond SuppressDoxyWarnings
IPluginRegistry() = default;
IPluginRegistry(IPluginRegistry const&) = delete;
@@ -121,7 +69,7 @@ public:
// @endcond
protected:
virtual ~IPluginRegistry() noexcept = default;
virtual ~IPluginRegistry() noexcept = 0;
public:
//!
@@ -159,25 +107,6 @@ public:
//!
virtual IErrorRecorder* getErrorRecorder() const noexcept = 0;
//!
//! \brief Deregister a previously registered plugin creator implementing IPluginCreator.
//!
//! Since there may be a desire to limit the number of plugins,
//! this function provides a mechanism for removing plugin creators registered in TensorRT.
//! The plugin creator that is specified by \p creator is removed from TensorRT and no longer tracked.
//!
//! \return True if the plugin creator was deregistered, false if it was not found in the registry or otherwise
//! could not be deregistered.
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: Yes
//!
//! \deprecated Deprecated in TensorRT 10.0. Superseded by
//! IPluginRegistry::deregisterCreator(IPluginCreatorInterface const&).
//!
TRT_DEPRECATED virtual bool deregisterCreator(IPluginCreator const& creator) noexcept = 0;
//!
//! \brief Return whether the parent registry will be searched if a plugin is not found in this registry
//! default: true
@@ -317,6 +246,8 @@ public:
virtual IPluginCreatorInterface* const* getAllCreatorsRecursive(int32_t* const numCreators) noexcept = 0;
};
inline IPluginRegistry::~IPluginRegistry() noexcept = default;
} // namespace nvinfer1
#endif /* NV_INFER_RUNTIME_COMMON_H */
include/NvInferRuntimePlugin.h
+7 -62
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -500,57 +500,6 @@ public:
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
= 0;
//!
//! \brief Return true if the output tensor is broadcast across a batch.
//!
//! \param outputIndex The index of the output tensor, which will be in the valid range between 0 and
//! nbOutputs()-1.
//! \param inputIsBroadcasted A boolean array of length nbInputs. The i-th element will be true if and only if
//! the tensor for the ith input is broadcast across a batch.
//! \param nbInputs The number of inputs. Will be a non-negative integer.
//!
//! The values in inputIsBroadcasted refer to broadcasting at the semantic level,
//! i.e. are unaffected by whether method canBroadcastInputAcrossBatch requests
//! physical replication of the values.
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: Yes, this method is required to be thread-safe and may be called from multiple threads
//! when building networks on multiple devices sharing the same plugin.
//!
//! \deprecated Deprecated in TensorRT 10.0. Implicit batch support is removed in TensorRT 10.0.
//!
TRT_DEPRECATED virtual bool isOutputBroadcastAcrossBatch(
int32_t outputIndex, bool const* inputIsBroadcasted, int32_t nbInputs) const noexcept
= 0;
//!
//! \brief Return true if the plugin can use an input tensor that is broadcast across batch without replication.
//!
//! \param inputIndex Index of input that could be broadcast. Will be in the valid range between 0 and
//! nbInputs - 1 where nbInputs is the maximum number of input tensors supported by this plugin.
//!
//! \return true if the index is in the valid range and the plugin is able to broadcast a single copy of this
//! input tensor across the batch. False otherwise.
//!
//! For each input whose tensor is semantically broadcast across a batch,
//! TensorRT calls this method before calling configurePlugin.
//! If canBroadcastInputAcrossBatch returns true, TensorRT will not replicate the input tensor;
//! i.e., there will be a single copy that the plugin must share across the batch.
//! If it returns false, TensorRT will replicate the input tensor
//! so that it appears like a non-broadcasted tensor.
//!
//! This method is called only for inputs that can be broadcast.
//!
//! \usage
//! - Allowed context for the API call
//! - Thread-safe: Yes, this method is required to be thread-safe and may be called from multiple threads
//! when building networks on multiple devices sharing the same plugin.
//!
//! \deprecated Deprecated in TensorRT 10.0. Implicit batch support is removed in TensorRT 10.0.
//!
TRT_DEPRECATED virtual bool canBroadcastInputAcrossBatch(int32_t inputIndex) const noexcept = 0;
//!
//! \brief Configure the layer with input and output data types.
//!
@@ -573,10 +522,8 @@ public:
//! The dimensions passed here do not include the outermost batch size (i.e. for 2D image networks, they will be
//! 3-dimensional CHW dimensions). When inputIsBroadcast or outputIsBroadcast is true, the outermost batch size for
//! that input or output must be treated as if it is one.
//! Index 'i' of inputIsBroadcast is true only if the input is semantically broadcast across the batch and
//! calling canBroadcastInputAcrossBatch with argument 'i' returns true.
//! Index 'i' of outputIsBroadcast is true only if calling isOutputBroadcastAcrossBatch with argument 'i'
//! returns true.
//! Index 'i' of inputIsBroadcast is true only if the input is semantically broadcast across the batch.
//! Index 'i' of outputIsBroadcast is true only if the output is semantically broadcast across the batch.
//!
//! \warning for the floatFormat field, the values PluginFormat::kCHW4, PluginFormat::kCHW16, and
//! PluginFormat::kCHW32 will not be passed in, this is to keep backward compatibility with TensorRT 5.x series. Use
@@ -600,10 +547,8 @@ public:
//! \brief Attach the plugin object to an execution context and grant the plugin the access to some context
//! resources.
//!
//! \param cudnn The cuDNN context handle of the execution context. Will be a valid cuDNN context handle, or
//! nullptr if TacticSource::kCUDNN is disabled.
//! \param cublas The cuBLAS context handle of the execution context. Will be a valid cuBLAS context handle, or
//! nullptr if TacticSource::kCUBLAS is disabled.
//! \param cudnn The cuDNN context handle of the execution context. Always nullptr.
//! \param cublas The cuBLAS context handle of the execution context. Always nullptr.
//! \param allocator The allocator used by the execution context
//!
//! This function is called automatically for each plugin when a new execution context is created. If the context
@@ -611,9 +556,9 @@ public:
//! new resources are assigned to the context.
//!
//! If the plugin needs per-context resource, it can be allocated here.
//! The plugin can also get context-owned cuDNN and cuBLAS context here.
//!
//! \note The TacticSource::kCUDNN and TacticSource::kCUBLAS flag is disabled by default.
//! \note The cuDNN and cuBLAS handles are always nullptr. Plugins that need cuDNN or cuBLAS
//! should create their own handles.
//! The allocator pointer is unique to each building or execution context instance having overlapping lifetimes.
//! It can be used as a key to manage resources across plugin instances sharing the same context.
//! Plugins attached to different contexts will have different handles as their execution will not overlap.
include/NvInferVersion.h
+4 -4
View File
@@ -23,10 +23,10 @@
#ifndef NV_INFER_VERSION_H
#define NV_INFER_VERSION_H
#define TRT_MAJOR_ENTERPRISE 10
#define TRT_MINOR_ENTERPRISE 16
#define TRT_PATCH_ENTERPRISE 1
#define TRT_BUILD_ENTERPRISE 11
#define TRT_MAJOR_ENTERPRISE 11
#define TRT_MINOR_ENTERPRISE 0
#define TRT_PATCH_ENTERPRISE 0
#define TRT_BUILD_ENTERPRISE 114
#define NV_TENSORRT_MAJOR TRT_MAJOR_ENTERPRISE //!< TensorRT major version.
#define NV_TENSORRT_MINOR TRT_MINOR_ENTERPRISE //!< TensorRT minor version.
#define NV_TENSORRT_PATCH TRT_PATCH_ENTERPRISE //!< TensorRT patch version.
include/NvOnnxConfig.h
+4 -2
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -46,7 +46,7 @@ namespace nvonnxparser
class IOnnxConfig
{
public:
virtual ~IOnnxConfig() noexcept = default;
virtual ~IOnnxConfig() noexcept = 0;
//!
//! \typedef Verbosity
//!
@@ -191,6 +191,8 @@ public:
}; // class IOnnxConfig
inline IOnnxConfig::~IOnnxConfig() noexcept = default;
TENSORRTAPI IOnnxConfig* createONNXConfig();
} // namespace nvonnxparser
include/NvOnnxParser.h
+5 -55
View File
@@ -20,8 +20,6 @@
#include "NvInfer.h"
#include <stddef.h>
#include <string>
#include <vector>
//!
//! \file NvOnnxParser.h
@@ -36,22 +34,6 @@
static constexpr int32_t NV_ONNX_PARSER_VERSION
= ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH);
//!
//! \typedef SubGraph_t
//!
//! \brief The data structure containing the parsing capability of
//! a set of nodes in an ONNX graph.
//!
typedef std::pair<std::vector<size_t>, bool> SubGraph_t;
//!
//! \typedef SubGraphCollection_t
//!
//! \brief The data structure containing all SubGraph_t partitioned
//! out of an ONNX graph.
//!
typedef std::vector<SubGraph_t> SubGraphCollection_t;
//!
//! \namespace nvonnxparser
//!
@@ -134,7 +116,7 @@ enum class OnnxParserFlag : int32_t
template <>
constexpr int32_t EnumMax<OnnxParserFlag>() noexcept
{
return 5;
return 4;
}
//!
@@ -235,39 +217,7 @@ public:
//! \return true if the model was parsed successfully
//!
//!
virtual bool parseFromFile(const char* onnxModelFile, int verbosity) noexcept = 0;
//!
//! [DEPRECATED] Deprecated in TensorRT 10.1. See supportsModelV2.
//!
//! \brief Check whether TensorRT supports a particular ONNX model.
//! If the function returns True, one can proceed to engine building
//! without having to call \p parse or \p parseFromFile.
//!
//! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serialized_onnx_model_size Size of the serialized ONNX model
//! in bytes
//! \param sub_graph_collection Container to hold supported subgraphs
//! \param model_path Absolute path to the model file for loading external weights if required
//! \return true if the model is supported
//!
TRT_DEPRECATED virtual bool supportsModel(void const* serialized_onnx_model, size_t serialized_onnx_model_size,
SubGraphCollection_t& sub_graph_collection, const char* model_path = nullptr) noexcept = 0;
//!
//! [DEPRECATED] Deprecated in TensorRT 10.13. See loadInitializer().
//!
//!\brief Parse a serialized ONNX model into the TensorRT network
//! with consideration of user provided weights
//!
//! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serialized_onnx_model_size Size of the serialized ONNX model
//! in bytes
//! \return true if the model was parsed successfully
//! \see getNbErrors() getError()
//!
TRT_DEPRECATED virtual bool parseWithWeightDescriptors(
void const* serialized_onnx_model, size_t serialized_onnx_model_size) noexcept = 0;
virtual bool parseFromFile(char const* onnxModelFile, int verbosity) noexcept = 0;
//!
//!\brief Returns whether the specified operator may be supported by the
@@ -318,7 +268,7 @@ public:
//! \param[out] nbPluginLibs Returns the number of plugin libraries in the array, or -1 if there was an error.
//! \return Array of `nbPluginLibs` C-strings describing plugin library paths on the filesystem if nbPluginLibs > 0,
//! or nullptr otherwise. This array is owned by the IParser, and the pointers in the array are only valid until
//! the next call to parse(), supportsModel(), parseFromFile(), or parseWithWeightDescriptors().
//! the next call to parse() or parseFromFile().
//!
virtual char const* const* getUsedVCPluginLibraries(int64_t& nbPluginLibs) const noexcept = 0;
@@ -432,8 +382,8 @@ public:
virtual int64_t* getSubgraphNodes(int64_t const index, int64_t& subgraphLength) noexcept = 0;
//!
//! \brief Load a serialized ONNX model into the parser. Unlike the parse(), parseFromFile(), or
//! parseWithWeightDescriptors() functions, this function does not immediately convert the model into a TensorRT
//! \brief Load a serialized ONNX model into the parser. Unlike the parse() or parseFromFile()
//! functions, this function does not immediately convert the model into a TensorRT
//! INetworkDefinition. Using this function allows users to provide their own initializers for the ONNX model
//! through the loadInitializer() function.
//!
include/impl/NvInferPythonPlugin.h
+14 -5
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -73,9 +73,11 @@ public:
virtual int32_t getNbSymExprs() const noexcept = 0;
virtual bool setNbSymExprs(int32_t count) noexcept = 0;
virtual ~ISymExprsImpl() noexcept = default;
virtual ~ISymExprsImpl() noexcept = 0;
};
inline ISymExprsImpl::~ISymExprsImpl() noexcept = default;
//! \class ISymExprs
//! \brief Allows for a sequence of symbolic expressions to be communicated to the TensorRT backend
//! \note Clients must not implement this class.
@@ -112,9 +114,11 @@ public:
protected:
ISymExprsImpl* mImpl{nullptr};
virtual ~ISymExprs() noexcept = default;
virtual ~ISymExprs() noexcept = 0;
};
inline ISymExprs::~ISymExprs() noexcept = default;
//! \enum QuickPluginCreationRequest
//! \brief Communicates preference when a quickly deployable plugin is to be added to the network
enum class QuickPluginCreationRequest : int32_t
@@ -156,9 +160,11 @@ public:
virtual ISymExpr* getSharedMem() noexcept = 0;
virtual bool setSharedMem(ISymExpr* sharedMem) noexcept = 0;
virtual ~IKernelLaunchParamsImpl() noexcept = default;
virtual ~IKernelLaunchParamsImpl() noexcept = 0;
};
inline IKernelLaunchParamsImpl::~IKernelLaunchParamsImpl() noexcept = default;
//! \class IKernelLaunchParams
//! \brief Allows for kernel launch parameters to be communicated to the TensorRT backend
//! \note Clients must not implement this class.
@@ -258,9 +264,12 @@ public:
protected:
IKernelLaunchParamsImpl* mImpl{nullptr};
virtual ~IKernelLaunchParams() noexcept = default;
virtual ~IKernelLaunchParams() noexcept = 0;
};
inline IKernelLaunchParams::~IKernelLaunchParams() noexcept = default;
namespace v_1_0
{

Some files were not shown because too many files have changed in this diff Show More