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

TensorRT Open Source Release/5.1

Browse Source
This commit is contained in:
Rajeev Rao
2019-06-17 10:13:09 -07:00
committed by Rajeev Rao
parent e80b96d795
commit 0d36bbb297
225 changed files with 50777 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.clang-format
+90
View File
@@ -0,0 +1,90 @@
---
Language: Cpp
# BasedOnStyle: WebKit
AccessModifierOffset: -4
AlignAfterOpenBracket: DontAlign
AlignConsecutiveAssignments: false
AlignConsecutiveDeclarations: false
AlignEscapedNewlinesLeft: false
AlignOperands: false
AlignTrailingComments: true
AllowAllParametersOfDeclarationOnNextLine: true
AllowShortBlocksOnASingleLine: false
AllowShortCaseLabelsOnASingleLine: true
AllowShortFunctionsOnASingleLine: Empty
AllowShortIfStatementsOnASingleLine: false
AllowShortLoopsOnASingleLine: false
AlwaysBreakAfterDefinitionReturnType: None
AlwaysBreakAfterReturnType: None
AlwaysBreakBeforeMultilineStrings: true
AlwaysBreakTemplateDeclarations: true
BinPackArguments: true
BinPackParameters: true
BraceWrapping:
AfterClass: true
AfterControlStatement: true
AfterEnum: true
AfterFunction: true
AfterNamespace: true
AfterObjCDeclaration: true
AfterStruct: true
AfterUnion: true
BeforeCatch: true
BeforeElse: true
IndentBraces: false
BreakBeforeBinaryOperators: All
BreakBeforeBraces: Custom
BreakBeforeTernaryOperators: true
BreakConstructorInitializersBeforeComma: true
ColumnLimit: 120
CommentPragmas: '^ IWYU pragma:'
ConstructorInitializerAllOnOneLineOrOnePerLine: false
ConstructorInitializerIndentWidth: 4
ContinuationIndentWidth: 4
Cpp11BracedListStyle: true
DerivePointerAlignment: false
DisableFormat: false
ExperimentalAutoDetectBinPacking: false
ForEachMacros: [ foreach, Q_FOREACH, BOOST_FOREACH ]
IncludeCategories:
- Regex: '^"(llvm|llvm-c|clang|clang-c)/'
Priority: 2
- Regex: '^(<|"(gtest|isl|json)/)'
Priority: 3
- Regex: '.*'
Priority: 1
IndentCaseLabels: false
IndentWidth: 4
IndentWrappedFunctionNames: false
KeepEmptyLinesAtTheStartOfBlocks: true
MacroBlockBegin: ''
MacroBlockEnd: ''
MaxEmptyLinesToKeep: 1
NamespaceIndentation: None
ObjCBlockIndentWidth: 4
ObjCSpaceAfterProperty: true
ObjCSpaceBeforeProtocolList: true
PenaltyBreakBeforeFirstCallParameter: 19
PenaltyBreakComment: 300
PenaltyBreakFirstLessLess: 120
PenaltyBreakString: 1000
PenaltyExcessCharacter: 1000000
PenaltyReturnTypeOnItsOwnLine: 60
PointerAlignment: Left
ReflowComments: true
SortIncludes: true
SpaceAfterCStyleCast: true
SpaceBeforeAssignmentOperators: true
SpaceBeforeParens: ControlStatements
SpaceInEmptyParentheses: false
SpacesBeforeTrailingComments: 1
SpacesInAngles: false
SpacesInContainerLiterals: true
SpacesInCStyleCastParentheses: false
SpacesInParentheses: false
SpacesInSquareBrackets: false
Standard: Cpp11
TabWidth: 4
UseTab: Never
...
.gitignore
+1
View File
@@ -0,0 +1 @@
build/
.gitmodules
+12
View File
@@ -0,0 +1,12 @@
[submodule "parsers/onnx"]
path = parsers/onnx
url = https://github.com/onnx/onnx-tensorrt.git
branch = 5.1
[submodule "third_party/protobuf"]
path = third_party/protobuf
url = https://github.com/protocolbuffers/protobuf.git
branch = 3.8.x
[submodule "third_party/cub"]
path = third_party/cub
url = https://github.com/NVlabs/cub.git
branch = 1.8.0
CMakeLists.txt
+201
View File
@@ -0,0 +1,201 @@
#
# Copyright (c) 2019, 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.
#
cmake_minimum_required(VERSION 3.13 FATAL_ERROR)
include(cmake/modules/set_ifndef.cmake)
include(cmake/modules/find_library_create_target.cmake)
set_ifndef(TRT_LIB_DIR ${CMAKE_BINARY_DIR})
set_ifndef(TRT_BIN_DIR ${CMAKE_BINARY_DIR})
file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/include/NvInfer.h" VERSION_STRINGS REGEX "#define NV_TENSORRT_.*")
foreach(TYPE MAJOR MINOR PATCH BUILD)
string(REGEX MATCH "NV_TENSORRT_${TYPE} [0-9]" TRT_TYPE_STRING ${VERSION_STRINGS})
string(REGEX MATCH "[0-9]" TRT_${TYPE} ${TRT_TYPE_STRING})
endforeach(TYPE)
foreach(TYPE MAJOR MINOR PATCH)
string(REGEX MATCH "NV_TENSORRT_SONAME_${TYPE} [0-9]" TRT_TYPE_STRING ${VERSION_STRINGS})
string(REGEX MATCH "[0-9]" TRT_SO_${TYPE} ${TRT_TYPE_STRING})
endforeach(TYPE)
set(TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}.${TRT_BUILD}" CACHE STRING "TRT project version")
set(TRT_SOVERSION "${TRT_SO_MAJOR}.${TRT_SO_MINOR}.${TRT_SO_PATCH}" CACHE STRING "TRT library so version")
message("Building for TensorRT version: ${TRT_VERSION}, library version: ${TRT_SOVERSION}")
set(FIND_CUDA "")
if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++)
set(FIND_CUDA "CUDA")
endif()
message("CHECK for ${FIND_CUDA}")
project(TensorRT
LANGUAGES CXX ${FIND_CUDA}
VERSION ${TRT_VERSION}
DESCRIPTION "TensorRT is a C++ library that facilitates high performance inference on NVIDIA GPUs and deep learning accelerators."
HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT")
if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
set(CMAKE_INSTALL_PREFIX ${TRT_LIB_DIR}/../ CACHE PATH "TensorRT installation" FORCE)
endif(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
option(BUILD_PLUGINS "Build TensorRT plugin" ON)
option(BUILD_PARSERS "Build TensorRT parsers" ON)
option(BUILD_SAMPLES "Build TensorRT samples" ON)
option(NVPARTNER "Build partner repos from source" OFF)
option(NVINTERNAL "Build in NVIDIA internal source tree" OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
############################# CROSS COMPILATION SETTINGS ##################################
set_ifndef(TRT_PLATFORM_ID "x86_64")
message(STATUS "Targeting TRT Platform: ${TRT_PLATFORM_ID}")
############################################################################################
set(TRT_DEBUG_POSTFIX _debug CACHE STRING "suffix for debug builds")
if (CMAKE_BUILD_TYPE STREQUAL "Debug")
message("Building in debug mode ${DEBUG_POSTFIX}")
endif()
set(CMAKE_CXX_FLAGS "-Wno-deprecated-declarations ${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss")
################################### DEPENDENCIES ##########################################
set(DEFAULT_CUDA_VERSION 10.1)
set(DEFAULT_CUDNN_VERSION 7.5)
set(DEFAULT_PROTOBUF_VERSION 3.0.0)
set(DEFAULT_PROTOBUF_INTERNAL_VERSION 10.0)
set(DEFAULT_CUB_VERSION 1.8.0)
# Dependency Version Resolution
set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION})
message(STATUS "CUDA version set to ${CUDA_VERSION}")
set_ifndef(CUDNN_VERSION ${DEFAULT_CUDNN_VERSION})
message(STATUS "cuDNN version set to ${CUDNN_VERSION}")
if (NVINTERNAL)
#TODO: Change this to set_ifndef(PROTOBUF_INTERNAL_VERSION ${DEFAULT_PROTOBUF_INTERNAL_VERSION}) once onnx-tensorrts build system is fixed
set_ifndef(PROTOBUF_VERSION ${DEFAULT_PROTOBUF_VERSION})
message(STATUS "Protobuf version set to ${PROTOBUF_INTERNAL_VERSION}")
set_ifndef(CUB_VERSION ${DEFAULT_CUB_VERSION})
message(STATUS "CUB version set to ${CUB_VERSION}")
#TODO: Remove this once CMake is fully intergrated in the P4 build system
set_ifndef(NVINTERNAL_SUFFIX "V2")
else()
set_ifndef(PROTOBUF_VERSION ${DEFAULT_PROTOBUF_VERSION})
message(STATUS "Protobuf version set to ${PROTOBUF_VERSION}")
endif()
find_package(Threads REQUIRED)
include(third_party/zlib.cmake)
include(third_party/protobuf.cmake)
if (NVINTERNAL)
########################################### DEPENDENCIES FOR BUILDING IN NVIDIA's TREE ############################################
set(EXTERNALS ${PROJECT_SOURCE_DIR}/../externals)
set(CUB_ROOT_DIR ${EXTERNALS}/cub/${CUB_VERSION} CACHE STRING "directory of CUB installation")
set(Protobuf_DIR ${EXTERNALS}/protobuf/${TRT_PLATFORM_ID} CACHE STRING "directory of PROTOBUF installation")
## This needs to be fixed to work with externals
if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
find_package(CUDA REQUIRED)
endif()
# Set this for ONNX Parser
set(CUDNN_ROOT_DIR ${EXTERNALS}/cudnn/${TRT_PLATFORM_ID}/${CUDNN_VERSION}/cuda-${CUDA_VERSION})
include_directories(
${CUDNN_ROOT_DIR}/include
${CUDA_TOOLKIT_ROOT_DIR}/include
/usr/local/cuda-${CUDA_VERSION}/include
)
#Check externals before using system
find_library(CUDNN_LIB cudnn HINTS
${CUDNN_ROOT_DIR}/lib64 NO_DEFAULT_PATH)
find_library(CUDNN_LIB cudnn HINTS
${CUDNN_ROOT_DIR}/lib64)
find_library(CUBLAS_LIB cublas HINTS
${CUDA_TOOLKIT_ROOT_DIR}/lib NO_DEFAULT_PATH)
find_library(CUBLAS_LIB cublas HINTS
${CUDA_TOOLKIT_ROOT_DIR}/lib)
if(BUILD_PARSERS)
#TODO: Change this to configure_protobuf_internal(${PROTOBUF_INTERNAL_VERSION}) once onnx-tensorrts build system is fixed
configure_protobuf(${PROTOBUF_VERSION})
endif()
########################################### DEPENDENCIES FOR BUILDING IN NVIDIA's TREE ############################################
else()
########################################### DEPENDENCIES FOR BUILDING OUTSIDE OF NVIDIA ############################################
if(NOT CUB_ROOT_DIR)
set(CUB_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/cub CACHE STRING "directory of CUB installation")
endif()
find_package(CUDA ${CUDA_VERSION} REQUIRED)
include_directories(
${CUDA_INCLUDE_DIRS}
)
find_library(CUDNN_LIB cudnn HINTS
${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDNN_ROOT_DIR}/lib64)
find_library(CUBLAS_LIB cublas HINTS
${CUDA_TOOLKIT_ROOT_DIR}/lib64)
if(BUILD_PARSERS)
configure_protobuf(${PROTOBUF_VERSION})
endif()
########################################### DEPENDENCIES FOR BUILDING OUTSIDE OF NVIDIA ############################################
endif()
find_library_create_target(nvinfer nvinfer SHARED ${TRT_LIB_DIR})
if (NOT (NVINTERNAL OR NVPARTNER))
find_library_create_target(nvuffparser nvparsers SHARED ${TRT_LIB_DIR})
endif()
find_library(CUDART_LIB cudart HINTS ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
find_library(RT_LIB rt)
############################################################################################
# TensorRT
if(BUILD_PLUGINS)
add_subdirectory(plugin${NVINTERNAL_SUFFIX})
else()
find_library_create_target(nvinfer_plugin nvinfer_plugin SHARED ${TRT_BIN_DIR} ${TRT_LIB_DIR})
endif()
if(BUILD_PARSERS)
add_subdirectory(parsers${NVINTERNAL_SUFFIX})
else()
if(NVPARTNER OR NVINTERNAL)
find_library_create_target(nvuffparser nvparsers SHARED ${TRT_BIN_DIR} ${TRT_LIB_DIR})
endif()
find_library_create_target(nvcaffeparser nvparsers SHARED ${TRT_BIN_DIR} ${TRT_LIB_DIR})
find_library_create_target(nvonnxparser nvonnxparser SHARED ${TRT_BIN_DIR} ${TRT_LIB_DIR})
endif()
if(BUILD_SAMPLES)
add_subdirectory(samples${NVINTERNAL_SUFFIX})
endif()
CODING-GUIDELINES.md
+436
View File
@@ -0,0 +1,436 @@
## TensorRT C++ Coding Guidelines
The TensorRT C++ Coding Guidelines are derived from several sources, primarily:
- [AUTOSAR C++ 2014](https://www.autosar.org/fileadmin/user_upload/standards/adaptive/17-03/AUTOSAR_RS_CPP14Guidelines.pdf)
- [MISRA C++ 2008](https://www.misra.org.uk/Activities/MISRAC/tabid/171/Default.aspx)
- [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html)
------
#### Namespaces
1. *MISRA C++: 2008 Rule 7-3-1*
Global namespace shall only contain main, namespace declarations and extern "C" declarations. Use explicit or anon namespaces for everything else.
2. Closing braces of namespaces should have a comment saying the namespace it closes:
```cpp
namespace foo
{
...
} // namespace foo
```
#### Constants
1. Prefer `const` or `constexpr` variables over `#defines` whenever possible, as the latter are not visible to the compiler.
2. *MISRA C++: 2008 Rule 7-1-1 and 7-1-2*
A variable that is not modified after its initialization should be declared as `const`.
3. For naming of constants, see the Naming section of this document.
#### Literals
1. Except `0` (only used in comparison for checking signness/existence/emptiness) and `nullptr`, `true`, `false`, all other literals should only be used for variable initialization.
Example:
```cpp
if (nbInputs == 2U){/*...*/}
```
Should be changed to:
```cpp
constexpr size_t kNbInputsWBias = 2U;
if (nbInputs == kNbInputsWBias) {/*...*/}
```
#### Brace Notation
1. Use the [Allman indentation](https://en.wikipedia.org/wiki/Indent_style#Allman_style) style.
2. Put the semicolon for an empty `for` or `while` loop in a new line.
3. *AUTOSAR C++14 Rule 6.6.3*, *MISRA C++: 2008 6-3-1*
The statement forming the body of a `switch`, `while`, `do .. while` or `for` statement shall be a compound statement. (use brace-delimited statements)
4. *AUTOSAR C++14 Rule 6.6.4*, *MISRA C++: 2008 Rule 6-4-1*
`If` and `else` should always be followed by brace-delimited statements, even if empty or a single statement.
#### Naming
1. Filenames
* Camel case with first letter lowercase: `thisIsASubDir` and `thisIsAFilename.cpp`
* *NOTE*: All files involved in the compilation of a compilation target (.exe/.so) must have filenames that are case-insensitive unique.
2. Types
* All types (including, but not limited to, class names) are [camel case](https://en.wikipedia.org/wiki/Camel_case) with uppercase first letter. Example: `FooBarClass`
3. Local variables, methods and namespaces
* Camel case with first letter lowercase. Example: `localFooBar`
4. Non-magic-number global variables that are non-static and not defined in anonymous namespace
* Camel case prefixed by a lower case 'g'. Example: `gDontUseGlobalFoos`
5. Non-magic-number global variables that are static or defined in an anonymous namespace
* Camel case prefixed by a lower case 's'. Example: `sMutableStaticGlobal`
6. Locally visible static variable
* Camel case with lowercase prefix ''s" as the first letter of the name. Example: `static std::once_flag sCaskInitOnce;`
7. Public, private and protected class member variables
* Camelcase prefixed with an 'm': `mNbFooValues`.
* Public member variables do not require the 'm' prefix but it is highly encouraged to use the prefix when needed to improve code clarity, especially in cases where the class is a base class in an inheritance chain.
8. Constants
* Enumerations, global constants, static constants at class-scope and function-scope magic-number/literal constants are uppercase snakecase with prefix 'k':
```cpp
const int kDIGIT_NUM = 10;
```
> *NOTE*: Function-scope constants that are not magic numbers or literals are named like non-constant variables:
```cpp
const bool pass = a && b;
```
9. Macros
* See [Constants](CODING-GUIDELINES.md#constants), which are preferred over `#define`.
* If you must use macros, however, follow uppercase snakecase: `FOO_VERSION`
Notes:
* In general we don't use [hungarian notation](https://en.wikipedia.org/wiki/Hungarian_notation), except for 'apps hungarian' in some cases such as 'nb' in a variable name to indicate count: `mNbTensorDescriptors`
* If a constructor's parameter name `foo` conflicts with a public member name `foo`, add a trailing underscore to the parameter name: `foo_`.
* *MISRA C++: 2008 Rule 2-13-4*
Literal suffixes should be upper case. For example, use `1234L` instead of `1234l`.
#### Tabs vs Spaces
1. Use only spaces. Do not use tabs.
2. Indent 4 spaces at a time. This is enforced automatically if you format your code using our clang-format config.
#### Formatting
1. Use the [LLVM clang-format](https://clang.llvm.org/docs/ClangFormat.html) tool for formatting your changes prior to submitting the PR.
2. Use a maximum of 120 characters per line. The auto formatting tool will wrap longer lines.
3. Exceptions to formatting violations must be justified on a per-case basis. Bypassing the formatting rules is discouraged, but can be achieved for exceptions as follows:
```cpp
// clang-format off
// .. Unformatted code ..
// clang-format on
```
#### Pointers and Memory Allocation
1. *AUTOSAR C++ 2014: 18-5-2/3*
Use smart pointers for allocating objects on the heap.
2. When picking a smart pointer, prefer `unique_ptr` for single resource ownership and `shared_ptr` for shared resource ownership. Use `weak_ptr` only in exceptional cases.
3. Do not use smart pointers that have been deprecated in C++11.
#### Comments
1. C++ comments are required. C comments are not allowed except for special cases (inline).
2. C++ style for single-line comments. `// This is a single line comment`
3. In function calls where parameters are not obvious from inspection, it can be helpful to use an inline C comment to document the parameter for readers:
```cpp
doSomeOperation(/* checkForErrors = */ false);
```
4. If the comment is a full sentence, it should be capitalized i.e. start with capital letter and punctuated properly.
5. Follow [Doxygen rules](http://www.doxygen.nl/manual/docblocks.html) for documenting new class interfaces and function prototypes.
* For C++-style single-line comments use `//!`.
* For class members, use `//!<`.
```cpp
//! This is a Doxygen comment
//! in C++ style
struct Foo
{
int x; //!< This is a Doxygen comment for members
}
```
#### Disabling Code
1. Use `#if` / `#endif` to disable code, preferably with a mnemonic condition like this:
```cpp
#if DEBUG_CONVOLUTION_INSTRUMENTATION
// ...code to be disabled...
#endif
```
```cpp
// Alternative: use a macro which evaluates to a noop in release code.
#if DEBUG_CONVOLUTION_INSTRUMENTATION
# define DEBUG_CONV_CODE(x) x
#else
# define DEBUG_CONV_CODE(x)
#endif
```
2. *MISRA C++: 2008 Rule 0-1-9*, *AutoSAR C++ 2014: 6-0-1*
Dead code is forbidden in safety-critical software - you may not use compile-time expressions and DCE to disable code. However, this technique can be useful elsewhere (e.g. tools, tests) to help prevent bitrot.
```cpp
// Not allowed in safety-critical code.
const bool gDisabledFeature = false;
void foo()
{
if (gDisabledFeature)
{
doSomething();
}
}
```
3. *MISRA C++: 2008 Rule 2-7-2 and 2-7-3*
Do NOT use comments to disable code. Use comments to explain code, not hide it.
#### Exceptions
1. Exceptions must not be thrown across library boundaries.
#### Casts
1. Use the least forceful cast necessary, or no cast if possible, to help the compiler diagnose unintended consequences.
2. Casting a pointer to a `void*` should be implicit (except if removing `const`).
3. *MISRA C++: 2008 Rule 5-2-5*
Casting should not remove any `const` or `volatile` qualification from the type of a pointer or reference.
4. *MISRA C++: 2008 Rule 5-2-4*
Do not use C-style casts (other than void casts) and functional notation casts (other than explicit constructor calls).
6. Casting from a `void*` to a `T*` should be done with `static_cast`, not `reinterpret_cast`, since the latter is more forceful.
7. Use `reinterpret_cast` as a last resort, where `const_cast` and `static_cast` won't work.
8. Avoid `dynamic_cast`.
#### Expressions
1. *MISRA C++: 2008 Rule 6-2-1*
Do not use assignment operator in subexpressions.
```cpp
// Not compliant
x = y = z;
// Not compliant
if (x = y)
{
// ...
}
```
#### Ternary operator
1. *AUTOSAR C++ 2014: 7-1-1*
Ternary operator should not be used as a sub-expression. Ternary operator expressions should be encapsulated with braces. Example:
```cpp
const auto var = (condition0 ? a : (condition1 ? b : c));
```
should be changed to:
```cpp
const auto d = (condition1 ? b : c);
const auto var = (condition0 ? a : d);
```
#### Statements
1. When practical, a `switch` statement controlled by an `enum` should have a case for each enum value and not have a default clause so that we get a compile-time error if a new enum value is added.
2. *MISRA C++:2008 Rules 6-4-3, 6-4-4, and 6-4-5*
Switch statements should be well structured. An informal guideline is to treat switch statements as structured multi-way branches and not "glorified gotos" such as:
```cpp
// Not compliant
switch (x) case 4: if (y) case 5: return 0; else default: return 1;
```
3. The "well structured" requirement prohibits fall-though except from one case label to another. Each case clause must be terminated in a break or throw. If a case clause has multiple statements, the braces are optional. The following example illustrates these requirements:
```cpp
switch (x)
{
case 0: // Fall-through allowed from case 0: to case 1: since case 0 is empty.
case 1:
a();
b();
break;
case 2:
case 4:
{ // With optional braces
c();
d();
break;
}
case 5:
c();
throw 42; // Terminating with throw is okay
default:
throw 42;
}
```
4. *MISRA C++:2008 Rule 6-4-3*
Ending a case clause with return is not allowed.
5. If a switch clause is a compound statement, put the break inside the braces.
```cpp
switch (x)
{
case 0:
case 1:
{
y();
z();
break;
}
...other cases...
}
```
#### Functions
1. Avoid declaring large functions as `inline`, absent a quantifiable benefit. Remember that functions defined in class declarations are implicitly inline.
2. Rather than using the `static` keyword to mark a function as having internal linkage, prefer to use anonymous namespaces instead.
3. *MISRA C++:2008 Rule 0-1-10*
Every defined function must be called at least once. That is, do not have unused methods.
4. *MISRA C++:2008 Rule 8-4-2*
Parameter names should be consistent across function definition and corresponding function declarations.
#### Forward declarations and extern variables
1. *MISRA C++: 2008 Rule 3-2-3*
For safety critical code, a type, object or function that is used in multiple translation units shall be declared in one and only one file.
* This means we cannot forward declare incomplete types in files where they are needed. Instead, we should put forward declarations in header files, and include these header files as needed.
#### Structures and Classes
1. *MISRA C++: 2008 Rule 14-7-1*
All class templates, function templates, class template member functions and class template static members shall be instantiated at least once. This prevents use of uninitialized variables.
2. *MISRA C++: 2008 Rule 11-01*
If class is not a *Plain Old Data Structure*, then its data members should be private.
#### Preprocessor Directives
1. *MISRA C++: 2008 Rule 16-0-2*
`#define` and `#undef` of macros should be done only at global namespace.
2. Avoid the use of `#ifdef` and `#ifndef` directives (except in the case of header include guards). Prefer to use `#if defined(...)` or `#if !defined(...)` instead. The latter syntax is more consistent with C syntax, and allows you to use more complicated preprocessor conditionals, e.g.:
```cpp
#if defined(FOO) || defined(BAR)
void foo();
#endif // defined(FOO) || defined(BAR)
```
3. When nesting preprocessor directives, use indentation after the hash mark (#). For example:
```cpp
#if defined(FOO)
# if FOO == 0
# define BAR 0
# elif FOO == 1
# define BAR 5
# else
# error "invalid FOO value"
# endif
#endif
```
4. Do not use `#pragma` once as include guard.
5. Use a preprocessor guard. It's standard-conforming and modern compilers are smart enough to open the file only once.
* The guard name must have prefix `TRT_` followed by the filename, all in caps. For a header file named `FooBarHello.h`, name the symbol as `TRT_FOO_BAR_HELLO_H`.
* Only use the file name to create the symbol. Unlike the Google C++ guideline, we do not include the directory names in the symbol. This is because we ensure all filenames are unique in the compilation unit.
* Do not use prefix with underscore. Such symbols are reserved in C++ standard for compilers or implementation.
* Do not use trailing underscore for the symbol. We differ in this from Google C++ guideline, which uses trailing underscore: `TRT_FOO_BAR_HELLO_H_`
```cpp
#ifndef TRT_FOO_BAR_HELLO_H
#define TRT_FOO_BAR_HELLO_H
// ...
#endif // TRT_FOO_BAR_HELLO_H
```
6. *AUTOSAR C++ 2014: 7-1-6*
Use `using` instead of `typedef`.
#### Signed vs Unsigned Integers
1. Use signed integers instead of unsigned, except for the cases below.
* The integer is a bitmap - use an unsigned type, since sign extension could lead to surprises.
* The integer is being used with an external library that expects an unsigned integer. A common example is a loop that compares against `std::vector::size()`, such as:
```cpp
for (size_t i = 0; i < mTensors.size(); ++i) // preferred style
```
* Using only signed integers for the above would lead to prolixity and perhaps unsafe narrowing:
```cpp
for (int i = 0; i < static_cast<int>(mTensors.size()); ++i)
```
#### Special Considerations for API
1. The API consists, with very few exceptions, of methodless structs and pure virtual interface classes.
2. API class methods should be either virtual or inline.
3. The API does not use integral types with platform-dependent sizes, other than `int`, `unsigned`, and `bool`. `size_t` should be used only for sizes of memory buffers.
4. The API does not use any aggregate types (e.g. `std::string`) which may be compiled differently with different compilers and libraries.
5. The API minimizes dependencies on system headers - currently only `<cstddef>` and `<cstdint>`.
6. Memory ownership may not be transferred across API boundaries - any memory allocated inside a library must be freed inside the library.
7. The API should be C++03.
8. New methods should be added at the end of interfaces so as to preserve v-table compatibility (compilers don't guarantee this, but de facto it works.)
9. Avoid optional arguments to functions, since they can make it difficult to extend interfaces.
10. Do not throw exceptions across library boundaries.
11. Document all APIs with doxygen.
#### Common Pitfalls
1. C headers should not be used directly.
- Example: Use `<cstdint>` instead of `<stdint.h>`
2. Do not use C library functions, whenever possible.
* Use brace initialization or `std::fill_n()` instead of `memset()`. This is especially important when dealing with non-[POD types](http://en.cppreference.com/w/cpp/concept/PODType). In the example below, using `memset()` will corrupt the vtable of `Foo:`
```cpp
struct Foo {
virtual int getX() { return x; }
int x;
};
...
// Bad: use memset() to initialize Foo
{
Foo foo;
memset(&foo, 0, sizeof(foo)); // Destroys hiddien virtual-function-table pointer!
}
// Good: use brace initialization to initialize Foo
{
Foo foo = {};
}
```
2. When specifying pointers to `const` data, the pointer itself may be `const`, in some usecases.
```cpp
char const * const errStr = getErrorStr(status);
```
----
## Appendix
#### Abbreviation Words and Compound Words as Part of Names
* Abbreviation words, which are usually fully-capitalized in literature, are treated as normal words without special capitalization, e.g. `gpuAllocator`, where GPU is converted to `gpu` before constructing the camel case name.
* Compound words, which are usually used in full in literature, e.g. `runtime`, can be abbreviated into fully capitalized letters, e.g. `RT` in NvInferRT.h.
#### Terminology
* *CUDA code* is code that must be compiled with a CUDA compiler. Typically, it includes:
* Declaration or definition of global or static variables with one of the following CUDA keywords: `__device__`, `__managed__` and `__constant__`.
* Declaration or definition of device functions decorated with `__device__`.
* Declaration or definition of kernels decorated with `__global__`.
* Kernel launching with <<<...>>> syntax.
> NOTE:
* Definition of kernel function pointer type aliases is not device code, e.g. `typedef __global__ void(*KernelFunc)(void* /*arg*/);`.
* Definition of pointers to kernel functions is not device code, either, e.g. `__global__ void(*KernelFunc)(void* /*arg*/) = getKernelFunc(parameters);` .
* Kernel launching with the CUDA runtime/driver API's, e.g. `cuLaunch` and `cudaLaunch`, is not CUDA code.
----
## NVIDIA Copyright
1. All TensorRT Open Source Software code should contain an NVIDIA copyright header that includes the current year. The following block of text should be prepended to the top of all OSS files. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
```cpp
/*
* Copyright (c) 2019, 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.
*/
```
CONTRIBUTING.md
+141
View File
@@ -0,0 +1,141 @@
## TensorRT OSS Contribution Rules
#### Issue Tracking
* All enhancement, bugfix, or change requests must begin with the creation of a [TensorRT Issue Request](https://github.com/nvidia/TensorRT/issues).
* The issue request must be reviewed by TensorRT engineers and approved prior to code review.
#### Coding Guidelines
- All source code contributions must strictly adhere to the [TensorRT Coding Guidelines](CODING-GUIDELINES.md).
- In addition, please follow the existing conventions in the relevant file, submodule, module, and project when you add new code or when you extend/fix existing functionality.
- To maintain consistency in code formatting and style, you should also run `clang-format` on the modified sources with the provided configuration file. This applies TensorRT code formatting rules to:
- class, function/method, and variable/field naming
- comment style
- indentation
- line length
- Format git changes:
```bash
# Commit ID is optional - if unspecified, run format on staged changes.
git-clang-format --style file [commit ID/reference]
```
- Format individual source files:
```bash
# -style=file : Obtain the formatting rules from .clang-format
# -i : In-place modification of the processed file
clang-format -style=file -i -fallback-style=none <file(s) to process>
```
- Format entire codebase (for project maintainers only):
```bash
find samples plugin -iname *.h -o -iname *.c -o -iname *.cpp -o -iname *.hpp \
| xargs clang-format -style=file -i -fallback-style=none
```
- Avoid introducing unnecessary complexity into existing code so that maintainability and readability are preserved.
- Try to keep pull requests (PRs) as concise as possible:
- Avoid committing commented-out code.
- Wherever possible, each PR should address a single concern. If there are several otherwise-unrelated things that should be fixed to reach a desired endpoint, our recommendation is to open several PRs and indicate the dependencies in the description. The more complex the changes are in a single PR, the more time it will take to review those changes.
- Write commit titles using imperative mood and [these rules](https://chris.beams.io/posts/git-commit/), and reference the Issue number corresponding to the PR. Following is the recommended format for commit texts:
```
#<Issue Number> - <Commit Title>
<Commit Body>
```
- Ensure that the build log is clean, meaning no warnings or errors should be present.
- Ensure that all `sample_*` tests pass prior to submitting your code.
- All OSS components must contain accompanying documentation (READMEs) describing the functionality, dependencies, and known issues.
- See `README.md` for existing samples and plugins for reference.
- All OSS components must have an accompanying test.
- If introducing a new component, such as a plugin, provide a test sample to verify the functionality.
- To add or disable functionality:
- Add a CMake option with a default value that matches the existing behavior.
- Where entire files can be included/excluded based on the value of this option, selectively include/exclude the relevant files from compilation by modifying `CMakeLists.txt` rather than using `#if` guards around the entire body of each file.
- Where the functionality involves minor changes to existing files, use `#if` guards.
- Make sure that you can contribute your work to open source (no license and/or patent conflict is introduced by your code). You will need to [`sign`](#signing-your-work) your commit.
- Thanks in advance for your patience as we review your contributions; we do appreciate them!
#### Pull Requests
Developer workflow for code contributions is as follows:
1. Developers must first [fork](https://help.github.com/en/articles/fork-a-repo) the [upstream](https://github.com/nvidia/TensorRT) TensorRT OSS repository.
2. Git clone the forked repository and push changes to the personal fork.
```bash
git clone https://github.com/YOUR_USERNAME/YOUR_FORK.git TensorRT
# Checkout the targeted branch and commit changes
# Push the commits to a branch on the fork (remote).
git push -u origin <local-branch>:<remote-branch>
```
3. Once the code changes are staged on the fork and ready for review, a [Pull Request](https://help.github.com/en/articles/about-pull-requests) (PR) can be [requested](https://help.github.com/en/articles/creating-a-pull-request) to merge the changes from a branch of the fork into a selected branch of upstream.
* Exercise caution when selecting the source and target branches for the PR.
Note that versioned releases of TensorRT OSS are posted to `release/` branches of the upstream repo.
* Creation of a PR creation kicks off the code review process.
* Atleast one TensorRT engineer will be assigned for the review.
* While under review, mark your PRs as work-in-progress by prefixing the PR title with [WIP].
4. Since there is no CI/CD process in place yet, the PR will be accepted and the corresponding issue closed only after adequate testing has been completed, manually, by the developer and/or TensorRT engineer reviewing the code.
#### Signing Your Work
* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
* Any contribution which contains commits that are not Signed-Off will not be accepted.
* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
```bash
$ git commit -s -m "Add cool feature."
```
This will append the following to your commit message:
```
Signed-off-by: Your Name <your@email.com>
```
* Full text of the DCO:
```
Developer Certificate of Origin
Version 1.1
Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
1 Letterman Drive
Suite D4700
San Francisco, CA, 94129
Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
```
```
Developer's Certificate of Origin 1.1
By making a contribution to this project, I certify that:
(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
(d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
```
LICENSE
+267
View File
@@ -0,0 +1,267 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
Copyright 2019 NVIDIA Corporation
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.
PORTIONS LICENSED AS FOLLOWS
> samples/common/windows/getopt.c
Copyright (c) 2002 Todd C. Miller <Todd.Miller@courtesan.com>
Permission to use, copy, modify, and distribute this software for any
purpose with or without fee is hereby granted, provided that the above
copyright notice and this permission notice appear in all copies.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
Sponsored in part by the Defense Advanced Research Projects
Agency (DARPA) and Air Force Research Laboratory, Air Force
Materiel Command, USAF, under agreement number F39502-99-1-0512.
Copyright (c) 2000 The NetBSD Foundation, Inc.
All rights reserved.
This code is derived from software contributed to The NetBSD Foundation
by Dieter Baron and Thomas Klausner.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
- Copyright (c) 2002 Todd C. Miller <Todd.Miller@courtesan.com>
- Copyright (c) 2000 The NetBSD Foundation, Inc.
> samples/common/half.h
Copyright (c) 2012-2017 Christian Rau <rauy@users.sourceforge.net>
Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
README.md
+212 -2
View File
@@ -1,3 +1,213 @@
# NVIDIA TensorRT Open Source Software
[![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)
Coming Soon
# TensorRT Open Source Software
This repository contains the Open Source Software (OSS) components of NVIDIA TensorRT. Included are the sources for TensorRT plugins and parsers (Caffe and ONNX), as well as sample applications demonstrating usage and capabilities of the TensorRT platform.
## Prerequisites
To build the TensorRT OSS components, ensure you meet the following package requirements:
**System Packages**
* [CUDA](https://developer.nvidia.com/cuda-toolkit)
* Recommended versions:
* [cuda-10.1](https://developer.nvidia.com/cuda-10.1-download-archive-base) + cuDNN-7.5
* [cuda-10.0](https://developer.nvidia.com/cuda-10.0-download-archive) + cuDNN-7.5
* [cuda-9.0](https://developer.nvidia.com/cuda-90-download-archive) + cuDNN 7.3
* [GNU Make](https://ftp.gnu.org/gnu/make/) >= v4.1
* [CMake](https://github.com/Kitware/CMake/releases) >= v3.13
* [Python](<https://www.python.org/downloads/>)
* Recommended versions:
* [Python2](https://www.python.org/downloads/release/python-2715/) >= v2.7.15
* [Python3](https://www.python.org/downloads/release/python-365/) >= v3.6.5
* [PIP](https://pypi.org/project/pip/#history) >= v19.0
* Essential libraries and utilities
* [Git](https://git-scm.com/downloads), [pkg-config](https://www.freedesktop.org/wiki/Software/pkg-config/), [Wget](https://www.gnu.org/software/wget/faq.html#download), [Zlib](https://zlib.net/)
**Optional Packages**
* Containerized builds
* [Docker](https://docs.docker.com/install/) >= 1.12
* [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker) >= 2.0
* Code formatting tools
* [Clang-format](https://clang.llvm.org/docs/ClangFormat.html)
* [Git-clang-format](https://github.com/llvm-mirror/clang/blob/master/tools/clang-format/git-clang-format)
**TensorRT Release**
* [TensorRT](https://developer.nvidia.com/nvidia-tensorrt-5x-download) v5.1.5
NOTE: Along with the TensorRT OSS components, the following source packages will also be downloaded, and they are not required to be installed on the system.
- [ONNX-TensorRT](https://github.com/onnx/onnx-tensorrt) v5.1
- [CUB](http://nvlabs.github.io/cub/) v1.8.0
- [Protobuf](https://github.com/protocolbuffers/protobuf.git) v3.8.x
## Downloading The TensorRT Components
1. #### Download TensorRT OSS sources.
```bash
git clone -b release/5.1 https://github.com/nvidia/TensorRT TensorRT
cd TensorRT
git submodule update --init --recursive
export TRT_SOURCE=`pwd`
```
2. #### Download the TensorRT binary release.
To build the TensorRT OSS, obtain the corresponding TensorRT 5.1.5 binary release from [NVidia Developer Zone](https://developer.nvidia.com/nvidia-tensorrt-5x-download). For a list of key features, known and fixed issues, see the [TensorRT 5.1.5 Release Notes](https://docs.nvidia.com/deeplearning/sdk/tensorrt-release-notes/tensorrt-5.html#rel_5-1-5).
**Example: Ubuntu 18.04 with cuda-10.1**
Download and extract the *TensorRT 5.1.5.0 GA for Ubuntu 18.04 and CUDA 10.1 tar package*
```bash
cd ~/Downloads
# Download TensorRT-5.1.5.0.Ubuntu-18.04.2.x86_64-gnu.cuda-10.1.cudnn7.5.tar.gz
tar -xvzf TensorRT-5.1.5.0.Ubuntu-18.04.2.x86_64-gnu.cuda-10.1.cudnn7.5.tar.gz
export TRT_RELEASE=`pwd`/TensorRT-5.1.5.0
```
**Example: CentOS/RedHat 7 with cuda-9.0**
Download and extract the *TensorRT 5.1.5.0 GA for CentOS/RedHat 7 and CUDA 9.0 tar package*
```bash
cd ~/Downloads
# Download TensorRT-5.1.5.0.Red-Hat.x86_64-gnu.cuda-9.0.cudnn7.5.tar.gz
tar -xvzf TensorRT-5.1.5.0.Red-Hat.x86_64-gnu.cuda-9.0.cudnn7.5.tar.gz
export TRT_RELEASE=~/Downloads/TensorRT-5.1.5.0
```
## Setting Up The Build Environment
* Install the *System Packages* list of components in the *Prerequisites* section.
* Alternatively, use the build containers as described below:
1. #### Generate the TensorRT build container.
The docker container can be built using the included Dockerfile. The build container is configured with the environment and packages required for building TensorRT OSS.
**Example: Ubuntu 18.04 with cuda-10.1**
```bash
docker build -f docker/ubuntu-18.04.Dockerfile --build-arg CUDA_VERSION=10.1 --tag=tensorrt .
```
**Example: CentOS/RedHat 7 with cuda-9.0**
```bash
docker build -f docker/centos-7.Dockerfile --build-arg CUDA_VERSION=9.0 --tag=tensorrt .
```
2. #### Launch the TensorRT build container.
```bash
docker run -v $TRT_RELEASE:/tensorrt -v $TRT_SOURCE:/workspace/TensorRT -it tensorrt:latest
```
> NOTE: To run TensorRT/CUDA programs within the build container, install [nvidia-docker](#prerequisites). Replace the `docker run` command with `nvidia-docker run` or `docker run --runtime=nvidia`.
## Building The TensorRT OSS Components
* Generate Makefiles and build.
```bash
cd $TRT_SOURCE
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_BIN_DIR=`pwd`/out
make -j$(nproc)
```
> NOTE:
> 1. The default CUDA version used by CMake is 10.1. To override this, for example to 9.0, append `-DCUDA_VERSION=9.0` to the cmake command.
> 2. If linking against the plugin and parser libraries obtained from TensorRT release (default behavior) is causing compatibility issues with TensorRT OSS, try building the OSS components separately in the following dependency order:
> ```bash
> # 1. Build Plugins
> cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_BIN_DIR=`pwd`/out \
> -DBUILD_PLUGINS=ON -DBUILD_PARSERS=OFF -DBUILD_SAMPLES=OFF
> make -j$(nproc)
> # 2. Build Parsers
> cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_BIN_DIR=`pwd`/out \
> -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=ON -DBUILD_SAMPLES=OFF
> make -j$(nproc)
> # 3. Build Samples
> cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_BIN_DIR=`pwd`/out \
> -DBUILD_PLUGINS=OFF -DBUILD_PARSERS=OFF -DBUILD_SAMPLES=ON
> make -j$(nproc)
> ```
The required CMake arguments are:
- `TRT_LIB_DIR`: Path to the TensorRT installation directory containing libraries.
- `TRT_BIN_DIR`: Output directory where generated build artifacts will be copied.
The following CMake build parameters are *optional*:
- `CMAKE_BUILD_TYPE`: Specify if binaries generated are for release or debug (contain debug symbols). Values consists of [`Release`] | `Debug`
- `CUDA_VERISON`: The version of CUDA to target, for example [`10.1`].
- `CUDNN_VERSION`: The version of cuDNN to target, for example [`7.5`].
- `PROTOBUF_VERSION`: The version of Protobuf to use, for example [`3.8.x`]. Note: Changing this will not configure CMake to use a system version of Protobuf, it will configure CMake to download and try building that version.
- `CMAKE_TOOLCHAIN_FILE`: The path to a toolchain file for cross compilation.
- `BUILD_PARSERS`: Specify if the parsers should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find precompiled versions of the parser libraries to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_PLUGINS`: Specify if the plugins should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find a precompiled version of the plugin library to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_SAMPLES`: Specify if the samples should be built, for example [`ON`] | `OFF`.
Other build options with limited applicability:
- `NVPARTNER`: For use by NVIDIA partners with exclusive source access. Values consists of [`OFF`] | `ON`.
- `CUB_VERSION`: The version of CUB to use, for example [`1.8.0`].
- `GPU_ARCHS`: GPU (SM) architectures to target. By default we generate CUDA code for the latest SM version. If lower SM versions are desired, they can be specified here as a comma separated list. Table of compute capabilities of NVIDIA GPUs can be found [here](https://developer.nvidia.com/cuda-gpus). Examples:
- Titan V: `-DGPU_ARCHS="70"`
- Tesla V100: `-DGPU_ARCHS="70"`
- GeForce RTX 2080: `-DGPU_ARCHS="75"`
- Tesla T4: `-DGPU_ARCHS="75"`
## Install the TensorRT OSS Components [Optional]
* Copy the build artifacts into the TensorRT installation directory, updating the installation.
* TensorRT installation directory is determined as `$TRT_LIB_DIR/..`
* Installation might require superuser privileges depending on the path and permissions of files being replaced.
```bash
sudo make install
```
## Useful Resources
#### TensorRT
* [TensorRT Homepage](https://developer.nvidia.com/tensorrt)
* [TensorRT Developer Guide](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html)
* [TensorRT Sample Support Guide](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html)
* [TensorRT Discussion Forums](https://devtalk.nvidia.com/default/board/304/tensorrt/)
## Known Issues
#### TensorRT 5.1.5
* FP16/INT8 modes have been disabled in SampleSSD (Caffe version). Please see the [SampleSSD README](samples/opensource/sampleSSD/README.md#known-issues) for details.
* Additionally, see the TensorRT [Release Notes](https://docs.nvidia.com/deeplearning/sdk/tensorrt-release-notes/tensorrt-5.html#rel_5-1-5).
VERSION
+1
View File
@@ -0,0 +1 @@
5.1.5.0
cmake/modules/find_library_create_target.cmake
+29
View File
@@ -0,0 +1,29 @@
#
# Copyright (c) 2019, 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.
#
macro(find_library_create_target target_name lib libtype hints)
message(STATUS "========================= Importing and creating target ${target_name} ==========================")
message(STATUS "Looking for library ${lib}")
if (CMAKE_BUILD_TYPE STREQUAL "Debug")
find_library(${lib}_LIB_PATH ${lib}${TRT_DEBUG_POSTFIX} HINTS ${hints} NO_DEFAULT_PATH)
endif()
find_library(${lib}_LIB_PATH ${lib} HINTS ${hints} NO_DEFAULT_PATH)
find_library(${lib}_LIB_PATH ${lib})
message(STATUS "Library that was found ${${lib}_LIB_PATH}")
add_library(${target_name} ${libtype} IMPORTED)
set_property(TARGET ${target_name} PROPERTY IMPORTED_LOCATION ${${lib}_LIB_PATH})
message(STATUS "==========================================================================================")
endmacro()
cmake/modules/set_ifndef.cmake
+20
View File
@@ -0,0 +1,20 @@
#
# Copyright (c) 2019, 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.
#
function (set_ifndef variable value)
if(NOT DEFINED ${variable})
set(${variable} ${value} PARENT_SCOPE)
endif()
endfunction()
cmake/toolchains/cmake_aarch64-android.toolchain
+47
View File
@@ -0,0 +1,47 @@
#
# Copyright (c) 2019, 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.
#
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} -lcublas -lcudart -lnvToolsExt -lculibos -lcudadevrt -llog)
set(DISABLE_SWIG TRUE)
set(TRT_PLATFORM_ID "aarch64-android")
cmake/toolchains/cmake_aarch64.toolchain
+57
View File
@@ -0,0 +1,57 @@
#
# Copyright (c) 2019, 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.
#
set(CMAKE_SYSTEM_NAME Linux)
set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(TRT_PLATFORM_ID "aarch64")
set(CUDA_PLATFORM_ID "aarch64-linux")
set(CMAKE_C_COMPILER $ENV{AARCH64_CC})
set(CMAKE_CXX_COMPILER $ENV{AARCH64_CC})
set(CMAKE_C_FLAGS "$ENV{AARCH64_CFLAGS}" CACHE STRING "" FORCE)
set(CMAKE_CXX_FLAGS "$ENV{AARCH64_CFLAGS}" CACHE STRING "" FORCE)
set(CMAKE_C_COMPILER_TARGET aarch64)
set(CMAKE_CXX_COMPILER_TARGET aarch64)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
if(NVINTERNAL)
set(EXT_PATH ${PROJECT_SOURCE_DIR}/../externals)
set(CUDA_ROOT ${EXT_PATH}/cuda-${CUDA_VERSION}-${TRT_PLATFORM_ID}/${CUDA_PLATFORM_ID})
else()
set(CUDA_ROOT /usr/local/cuda-${CUDA_VERSION}/targets/${CUDA_PLATFORM_ID})
endif()
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)
if(DEFINED ENV{VULCAN} AND NOT $ENV{VULCAN} STREQUAL "")
message("cmake_aarch64.toolchain using VULCAN mode")
set(CUDA_LIBS ${DEVLIBPATHS} -L${CUDA_ROOT}/lib64)
else()
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
endif()
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm)
set(DISABLE_SWIG TRUE)
cmake/toolchains/cmake_ppc64le.toolchain
+39
View File
@@ -0,0 +1,39 @@
#
# Copyright (c) 2019, 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.
#
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_C_COMPILER_TARGET ppc64le)
set(CMAKE_CXX_COMPILER_TARGET ppc64le)
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)
if(DEFINED ENV{VULCAN} AND NOT $ENV{VULCAN} STREQUAL "")
set(DISABLE_SWIG TRUE)
endif()
set(TRT_PLATFORM_ID "ppc64le")
cmake/toolchains/cmake_qnx.toolchain
+68
View File
@@ -0,0 +1,68 @@
#
# Copyright (c) 2019, 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.
#
set(CMAKE_SYSTEM_NAME qnx)
set(CMAKE_SYSTEM_PROCESSOR aarch64)
if(DEFINED ENV{QNX_BASE})
set(QNX_BASE $ENV{QNX_BASE})
message(STATUS "Found QNX_BASE = ${QNX_BASE}")
elseif(DEFINED ENV{TOOLS_BASE})
set(QNX_BASE $ENV{TOOLS_BASE}/embedded/qnx/qnx700-ga3)
message(STATUS "Found QNX_BASE = ${QNX_BASE}")
else()
message(FATAL_ERROR "QNX_BASE was not found")
endif()
set(ENV{QNX_HOST} ${QNX_BASE}/host/linux/x86_64)
set(ENV{QNX_TARGET} ${QNX_BASE}/target/qnx7)
set(QNX_HOST $ENV{QNX_HOST})
set(QNX_TARGET $ENV{QNX_TARGET})
message(STATUS "QNX_HOST = ${QNX_HOST}")
message(STATUS "QNX_TARGET = ${QNX_TARGET}")
set(CMAKE_C_COMPILER ${QNX_HOST}/usr/bin/aarch64-unknown-nto-qnx7.0.0-gcc)
set(CMAKE_CXX_COMPILER ${QNX_HOST}/usr/bin/aarch64-unknown-nto-qnx7.0.0-g++)
set(CMAKE_C_COMPILER_TARGET aarch64)
set(CMAKE_CXX_COMPILER_TARGET aarch64)
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" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
#...VULCAN (eris) does not use externals, more over libnvToolsExt.so is not even part of .vulcan
#...So, to resolve for all libs we need to use lib64 (vs normal lib in externals).
#...Moreover, for libnvToolsExt.so we still need to add externals to the PATH
if(DEFINED ENV{VULCAN} AND NOT $ENV{VULCAN} STREQUAL "")
set(CUDA_LIBS ${DEVLIBPATHS} -L${CUDA_ROOT}/lib64 -L${EXT_PATH}/cuda-10.0-qnx/lib)
else()
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
endif()
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart)
#...Disable swig
set(DISABLE_SWIG TRUE)
set(TRT_PLATFORM_ID "aarch64-qnx")
cmake/toolchains/cmake_x64_win.toolchain
+48
View File
@@ -0,0 +1,48 @@
#
# Copyright (c) 2019, 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.
#
set(CMAKE_SYSTEM_NAME WindowsStore)
set(CMAKE_SYSTEM_VERSION 10.0)
set(CMAKE_C_COMPILER ${CC})
set(CMAKE_CXX_COMPILER ${CC})
if(DEFINED CUDA_TOOLKIT)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT})
endif()
set(CMAKE_CUDA_COMPILER ${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc.exe)
set(CMAKE_CUDA_COMPILER_ID "NVIDIA")
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
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)
set(ADDITIONAL_PLATFORM_INCL_FLAGS "-I${MSVC_COMPILER_DIR}/include -I${MSVC_COMPILER_DIR}/../ucrt/include")
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${ADDITIONAL_PLATFORM_LIB_FLAGS} "-LIBPATH:${NV_TOOLS}/ddk/wddmv2/dev/rs4/17130/Lib/10.0.17130.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 cublas.lib")
set(TRT_PLATFORM_ID "win10")
cmake/toolchains/cmake_x86_64.toolchain
+33
View File
@@ -0,0 +1,33 @@
#
# Copyright (c) 2019, 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.
#
set(CMAKE_SYSTEM_NAME Linux)
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})
endif()
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
if(DEFINED ENV{VULCAN} AND NOT $ENV{VULCAN} STREQUAL "")
set(DISABLE_SWIG TRUE)
endif()
set(TRT_PLATFORM_ID "x86_64")
docker/centos-7.Dockerfile
+47
View File
@@ -0,0 +1,47 @@
# Copyright (c) 2019, 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.
ARG CUDA_VERSION=10.1
FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-cudnn7-devel-centos7
LABEL maintainer="NVIDIA CORPORATION"
# Install requried libraries
RUN yum -y install \
libcurl4-openssl-dev \
wget \
zlib-devel \
git \
pkg-config \
python3 \
python3-pip \
make
RUN cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip
# Install Cmake
RUN cd /tmp &&\
wget https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.sh &&\
chmod +x cmake-3.14.4-Linux-x86_64.sh &&\
./cmake-3.14.4-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license &&\
rm ./cmake-3.14.4-Linux-x86_64.sh
# Set environment and working directory
ENV TRT_RELEASE /tensorrt
ENV TRT_SOURCE /workspace/TensorRT
WORKDIR /workspace
RUN ["/bin/bash"]
docker/ubuntu-16.04.Dockerfile
+46
View File
@@ -0,0 +1,46 @@
# Copyright (c) 2019, 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.
ARG CUDA_VERSION=10.1
FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-cudnn7-devel-ubuntu16.04
LABEL maintainer="NVIDIA CORPORATION"
# Install requried libraries
RUN apt-get update && apt-get install -y --no-install-recommends \
libcurl4-openssl-dev \
wget \
zlib1g-dev \
git \
pkg-config \
python3 \
python3-pip
RUN cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip
# Install Cmake
RUN cd /tmp &&\
wget https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.sh &&\
chmod +x cmake-3.14.4-Linux-x86_64.sh &&\
./cmake-3.14.4-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license &&\
rm ./cmake-3.14.4-Linux-x86_64.sh
# Set environment and working directory
ENV TRT_RELEASE /tensorrt
ENV TRT_SOURCE /workspace/TensorRT
WORKDIR /workspace
RUN ["/bin/bash"]
docker/ubuntu-18.04.Dockerfile
+46
View File
@@ -0,0 +1,46 @@
# Copyright (c) 2019, 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.
ARG CUDA_VERSION=10.1
FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-cudnn7-devel-ubuntu18.04
LABEL maintainer="NVIDIA CORPORATION"
# Install requried libraries
RUN apt-get update && apt-get install -y --no-install-recommends \
libcurl4-openssl-dev \
wget \
zlib1g-dev \
git \
pkg-config \
python3 \
python3-pip
RUN cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip
# Install Cmake
RUN cd /tmp &&\
wget https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.sh &&\
chmod +x cmake-3.14.4-Linux-x86_64.sh &&\
./cmake-3.14.4-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license &&\
rm ./cmake-3.14.4-Linux-x86_64.sh
# Set environment and working directory
ENV TRT_RELEASE /tensorrt
ENV TRT_SOURCE /workspace/TensorRT
WORKDIR /workspace
RUN ["/bin/bash"]
include/NvCaffeParser.h
+277
View File
@@ -0,0 +1,277 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_CAFFE_PARSER_H
#define NV_CAFFE_PARSER_H
#include "NvInfer.h"
namespace ditcaffe
{
class NetParameter;
}
namespace nvcaffeparser1
{
//!
//! \class IBlobNameToTensor
//!
//! \brief Object used to store and query Tensors after they have been extracted from a Caffe model using the ICaffeParser.
//!
//! \note The lifetime of IBlobNameToTensor is the same as the lifetime of its parent ICaffeParser.
//!
//! \see nvcaffeparser1::ICaffeParser
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IBlobNameToTensor
{
public:
//! \brief Given a blob name, returns a pointer to a ITensor object.
//!
//! \param name Caffe blob name for which the user wants the corresponding ITensor.
//!
//! \return ITensor* corresponding to the queried name. If no such ITensor exists, then nullptr is returned.
//!
virtual nvinfer1::ITensor* find(const char* name) const = 0;
protected:
virtual ~IBlobNameToTensor() {}
};
//!
//! \class IBinaryProtoBlob
//!
//! \brief Object used to store and query data extracted from a binaryproto file using the ICaffeParser.
//!
//! \see nvcaffeparser1::ICaffeParser
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IBinaryProtoBlob
{
public:
virtual const void* getData() = 0;
virtual nvinfer1::DimsNCHW getDimensions() = 0;
virtual nvinfer1::DataType getDataType() = 0;
virtual void destroy() = 0;
protected:
virtual ~IBinaryProtoBlob() {}
};
//!
//! \class IPluginFactory
//!
//! \brief Plugin factory used to configure plugins.
//!
class IPluginFactory
{
public:
//!
//! \brief A user implemented function that determines if a layer configuration is provided by an IPlugin.
//!
//! \param layerName Name of the layer which the user wishes to validate.
//!
virtual bool isPlugin(const char* layerName) = 0;
//!
//! \brief Creates a plugin.
//!
//! \param layerName Name of layer associated with the plugin.
//! \param weights Weights used for the layer.
//! \param nbWeights Number of weights.
//!
virtual nvinfer1::IPlugin* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights) = 0;
};
//!
//! \class IPluginFactoryExt
//!
//! \brief Plugin factory used to configure plugins with added support for TRT versioning.
//!
class IPluginFactoryExt : public IPluginFactory
{
public:
virtual int getVersion() const
{
return NV_TENSORRT_VERSION;
}
//!
//! \brief A user implemented function that determines if a layer configuration is provided by an IPluginExt.
//!
//! \param layerName Name of the layer which the user wishes to validate.
//!
virtual bool isPluginExt(const char* layerName) = 0;
};
//!
//! \class IPluginFactoryV2
//!
//! \brief Plugin factory used to configure plugins.
//!
class IPluginFactoryV2
{
public:
//!
//! \brief A user implemented function that determines if a layer configuration is provided by an IPluginV2.
//!
//! \param layerName Name of the layer which the user wishes to validate.
//!
virtual bool isPluginV2(const char* layerName) = 0;
//!
//! \brief Creates a plugin.
//!
//! \param layerName Name of layer associated with the plugin.
//! \param weights Weights used for the layer.
//! \param nbWeights Number of weights.
//! \param libNamespace Library Namespace associated with the plugin object
//!
virtual nvinfer1::IPluginV2* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights, const char* libNamespace = "") = 0;
};
//!
//! \class ICaffeParser
//!
//! \brief Class used for parsing Caffe models.
//!
//! Allows users to export models trained using Caffe to TRT.
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class ICaffeParser
{
public:
//!
//! \brief Parse a prototxt file and a binaryproto Caffe model to extract
//! network configuration and weights associated with the network, respectively.
//!
//! \param deploy The plain text, prototxt file used to define the network configuration.
//! \param model The binaryproto Caffe model that contains the weights associated with the network.
//! \param network Network in which the CaffeParser will fill the layers.
//! \param weightType The type to which the weights will transformed.
//!
//! \return A pointer to an IBlobNameToTensor object that contains the extracted data.
//!
//! \see nvcaffeparser1::IBlobNameToTensor
//!
virtual const IBlobNameToTensor* parse(const char* deploy,
const char* model,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightType)
= 0;
//!
//! \brief Parse a deploy prototxt a binaryproto Caffe model from memory buffers to extract
//! network configuration and weights associated with the network, respectively.
//!
//! \param deployBuffer The plain text deploy prototxt used to define the network configuration.
//! \param deployLength The length of the deploy buffer.
//! \param modelBuffer The binaryproto Caffe memory buffer that contains the weights associated with the network.
//! \param modelLength The length of the model buffer.
//! \param network Network in which the CaffeParser will fill the layers.
//! \param weightType The type to which the weights will transformed.
//!
//! \return A pointer to an IBlobNameToTensor object that contains the extracted data.
//!
//! \see nvcaffeparser1::IBlobNameToTensor
//!
virtual const IBlobNameToTensor* parseBuffers(const char* deployBuffer,
std::size_t deployLength,
const char* modelBuffer,
std::size_t modelLength,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightType) = 0;
//!
//! \brief Parse and extract data stored in binaryproto file.
//!
//! The binaryproto file contains data stored in a binary blob. parseBinaryProto() converts it
//! to an IBinaryProtoBlob object which gives the user access to the data and meta-data about data.
//!
//! \param fileName Path to file containing binary proto.
//!
//! \return A pointer to an IBinaryProtoBlob object that contains the extracted data.
//!
//! \see nvcaffeparser1::IBinaryProtoBlob
//!
virtual IBinaryProtoBlob* parseBinaryProto(const char* fileName) = 0;
//!
//! \brief Set buffer size for the parsing and storage of the learned model.
//!
//! \param size The size of the buffer specified as the number of bytes.
//!
//! \note Default size is 2^30 bytes.
//!
virtual void setProtobufBufferSize(size_t size) = 0;
//!
//! \brief Set the IPluginFactory used to create the user defined plugins.
//!
//! \param factory Pointer to an instance of the user implmentation of IPluginFactory.
//!
virtual void setPluginFactory(IPluginFactory* factory) = 0;
//!
//! \brief Set the IPluginFactoryExt used to create the user defined pluginExts.
//!
//! \param factory Pointer to an instance of the user implmentation of IPluginFactoryExt.
//!
virtual void setPluginFactoryExt(IPluginFactoryExt* factory) = 0;
//!
//! \brief Destroy this ICaffeParser object.
//!
virtual void destroy() = 0;
//!
//! \brief Set the IPluginFactoryV2 used to create the user defined pluginV2 objects.
//!
//! \param factory Pointer to an instance of the user implmentation of IPluginFactoryV2.
//!
virtual void setPluginFactoryV2(IPluginFactoryV2* factory) = 0;
//!
//! \brief Set the namespace used to lookup and create plugins in the network.
//!
virtual void setPluginNamespace(const char* libNamespace) = 0;
protected:
virtual ~ICaffeParser() {}
};
//!
//! \brief Creates a ICaffeParser object.
//!
//! \return A pointer to the ICaffeParser object is returned.
//!
//! \see nvcaffeparser1::ICaffeParser
//!
TENSORRTAPI ICaffeParser* createCaffeParser();
//!
//! \brief Shuts down protocol buffers library.
//!
//! \note No part of the protocol buffers library can be used after this function is called.
//!
TENSORRTAPI void shutdownProtobufLibrary();
}
extern "C" TENSORRTAPI void* createNvCaffeParser_INTERNAL();
#endif
include/NvInfer.h
+5909
View File
File diff suppressed because it is too large Load Diff
include/NvInferPlugin.h
+486
View File
@@ -0,0 +1,486 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_INFER_PLUGIN_H
#define NV_INFER_PLUGIN_H
#include "NvInfer.h"
#include "NvInferPluginUtils.h"
//!
//! \file NvInferPlugin.h
//!
//! This is the API for the Nvidia provided TensorRT plugins.
//!
namespace nvinfer1
{
//!
//! \enum PluginType
//!
//! \brief The type values for the various plugins.
//!
//! \see INvPlugin::getPluginType()
//!
enum class PluginType : int
{
kFASTERRCNN = 0, //!< FasterRCNN fused plugin (RPN + ROI pooling).
kNORMALIZE = 1, //!< Normalize plugin.
kPERMUTE = 2, //!< Permute plugin.
kPRIORBOX = 3, //!< PriorBox plugin.
kSSDDETECTIONOUTPUT = 4, //!< SSD DetectionOutput plugin.
kCONCAT = 5, //!< Concat plugin.
kPRELU = 6, //!< YOLO PReLU Plugin.
kYOLOREORG = 7, //!< YOLO Reorg Plugin.
kYOLOREGION = 8, //!< YOLO Region Plugin.
kANCHORGENERATOR = 9, //!< SSD Grid Anchor Generator.
};
//!< Maximum number of elements in PluginType enum. \see PluginType
template <>
inline int EnumMax<PluginType>()
{
return 10;
}
namespace plugin
{
//!
//! \class INvPlugin
//!
//! \brief Common interface for the Nvidia created plugins.
//!
//! This class provides a common subset of functionality that is used
//! to provide distinguish the Nvidia created plugins. Each plugin provides a
//! function to validate the parameter options and create the plugin
//! object.
//!
class INvPlugin : public IPlugin
{
public:
//!
//! \brief Get the parameter plugin ID.
//!
//! \return The ID of the plugin.
//!
virtual PluginType getPluginType() const = 0;
//!
//! \brief Get the name of the plugin from the ID
//!
//! \return The name of the plugin specified by \p id. Return nullptr if invalid ID is specified.
//!
//! The valid \p id values are ranged [0, numPlugins()).
//!
virtual const char* getName() const = 0;
//!
//! \brief Destroy the plugin.
//!
//! The valid \p id values are ranged [0, numPlugins()).
//!
virtual void destroy() = 0;
protected:
~INvPlugin() {}
}; // INvPlugin
//!
//! \param featureStride Feature stride.
//! \param preNmsTop Number of proposals to keep before applying NMS.
//! \param nmsMaxOut Number of remaining proposals after applying NMS.
//! \param iouThreshold IoU threshold.
//! \param minBoxSize Minimum allowed bounding box size before scaling.
//! \param spatialScale Spatial scale between the input image and the last feature map.
//! \param pooling Spatial dimensions of pooled ROIs.
//! \param anchorRatios Aspect ratios for generating anchor windows.
//! \param anchorScales Scales for generating anchor windows.
//! \brief Create a plugin layer that fuses the RPN and ROI pooling using user-defined parameters.
//!
//! \return Returns a FasterRCNN fused RPN+ROI pooling plugin. Returns nullptr on invalid inputs.
//!
//! \see INvPlugin
//! \deprecated. This plugin is superseded by createRPNROIPlugin()
//!
TENSORRTAPI INvPlugin* createFasterRCNNPlugin(int featureStride, int preNmsTop,
int nmsMaxOut, float iouThreshold, float minBoxSize,
float spatialScale, DimsHW pooling,
Weights anchorRatios, Weights anchorScales);
TENSORRTAPI INvPlugin* createFasterRCNNPlugin(const void* data, size_t length);
//!
//! \brief The Normalize plugin layer normalizes the input to have L2 norm of 1 with scale learnable.
//! \param scales Scale weights that are applied to the output tensor.
//! \param acrossSpatial Whether to compute the norm over adjacent channels (acrossSpatial is true) or nearby spatial locations (within channel in which case acrossSpatial is false).
//! \param channelShared Whether the scale weight(s) is shared across channels.
//! \param eps Epsilon for not diviiding by zero.
//! \deprecated. This plugin is superseded by createNormalizePlugin()
//!
TENSORRTAPI INvPlugin* createSSDNormalizePlugin(const Weights* scales, bool acrossSpatial, bool channelShared, float eps);
TENSORRTAPI INvPlugin* createSSDNormalizePlugin(const void* data, size_t length);
//!
//! \brief The Permute plugin layer permutes the input tensor by changing the memory order of the data.
//! Quadruple defines a structure that contains an array of 4 integers. They can represent the permute orders or the strides in each dimension.
//!
typedef struct
{
int data[4];
} Quadruple;
//!
//! \param permuteOrder The new orders that are used to permute the data.
//! \deprecated. Please use the TensorRT Shuffle layer for Permute operation
//!
TENSORRTAPI INvPlugin* createSSDPermutePlugin(Quadruple permuteOrder);
TENSORRTAPI INvPlugin* createSSDPermutePlugin(const void* data, size_t length);
//!
//! \brief The PriorBox plugin layer generates the prior boxes of designated sizes and aspect ratios across all dimensions (H x W).
//! PriorBoxParameters defines a set of parameters for creating the PriorBox plugin layer.
//! It contains:
//! \param minSize Minimum box size in pixels. Can not be nullptr.
//! \param maxSize Maximum box size in pixels. Can be nullptr.
//! \param aspectRatios Aspect ratios of the boxes. Can be nullptr.
//! \param numMinSize Number of elements in minSize. Must be larger than 0.
//! \param numMaxSize Number of elements in maxSize. Can be 0 or same as numMinSize.
//! \param numAspectRatios Number of elements in aspectRatios. Can be 0.
//! \param flip If true, will flip each aspect ratio. For example, if there is aspect ratio "r", the aspect ratio "1.0/r" will be generated as well.
//! \param clip If true, will clip the prior so that it is within [0,1].
//! \param variance Variance for adjusting the prior boxes.
//! \param imgH Image height. If 0, then the H dimension of the data tensor will be used.
//! \param imgW Image width. If 0, then the W dimension of the data tensor will be used.
//! \param stepH Step in H. If 0, then (float)imgH/h will be used where h is the H dimension of the 1st input tensor.
//! \param stepW Step in W. If 0, then (float)imgW/w will be used where w is the W dimension of the 1st input tensor.
//! \param offset Offset to the top left corner of each cell.
//!
struct PriorBoxParameters
{
float *minSize, *maxSize, *aspectRatios;
int numMinSize, numMaxSize, numAspectRatios;
bool flip;
bool clip;
float variance[4];
int imgH, imgW;
float stepH, stepW;
float offset;
};
//!
//! \brief The Anchor Generator plugin layer generates the prior boxes of designated sizes and aspect ratios across all dimensions (H x W).
//! GridAnchorParameters defines a set of parameters for creating the plugin layer for all feature maps.
//! It contains:
//! \param minScale Scale of anchors corresponding to finest resolution.
//! \param maxScale Scale of anchors corresponding to coarsest resolution.
//! \param aspectRatios List of aspect ratios to place on each grid point.
//! \param numAspectRatios Number of elements in aspectRatios.
//! \param H Height of feature map to generate anchors for.
//! \param W Width of feature map to generate anchors for.
//! \param variance Variance for adjusting the prior boxes.
//!
struct GridAnchorParameters
{
float minSize, maxSize;
float* aspectRatios;
int numAspectRatios, H, W;
float variance[4];
};
//!
//! \param param Set of parameters for creating the PriorBox plugin layer.
//! \deprecated. This plugin is superseded by createPriorBoxPlugin()
//!
TENSORRTAPI INvPlugin* createSSDPriorBoxPlugin(PriorBoxParameters param);
TENSORRTAPI INvPlugin* createSSDPriorBoxPlugin(const void* data, size_t length);
//!
//! \brief The Grid Anchor Generator plugin layer generates the prior boxes of
//! designated sizes and aspect ratios across all dimensions (H x W) for all feature maps.
//! GridAnchorParameters defines a set of parameters for creating the GridAnchorGenerator plugin layer.
//! \deprecated. This plugin is superseded by createAnchorGeneratorPlugin()
//!
TENSORRTAPI INvPlugin* createSSDAnchorGeneratorPlugin(GridAnchorParameters* param, int numLayers);
TENSORRTAPI INvPlugin* createSSDAnchorGeneratorPlugin(const void* data, size_t length);
//!
//! \enum CodeTypeSSD
//! \brief The type of encoding used for decoding the bounding boxes and loc_data.
//!
enum class CodeTypeSSD : int
{
CORNER = 0, //!< Use box corners.
CENTER_SIZE = 1, //!< Use box centers and size.
CORNER_SIZE = 2, //!< Use box centers and size.
TF_CENTER = 3 //!< Use box centers and size but flip x and y coordinates.
};
//!
//! \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 bouding boxes.
//! DetectionOutputParameters defines a set of parameters for creating the DetectionOutput plugin layer.
//! It contains:
//! \param shareLocation If true, bounding box are shared among different classes.
//! \param varianceEncodedInTarget If true, variance is encoded in target. Otherwise we need to adjust the predicted offset accordingly.
//! \param backgroundLabelId Background label ID. If there is no background class, set it as -1.
//! \param numClasses Number of classes to be predicted.
//! \param topK Number of boxes per image with top confidence scores that are fed into the NMS algorithm.
//! \param keepTopK Number of total bounding boxes to be kept per image after NMS step.
//! \param confidenceThreshold Only consider detections whose confidences are larger than a threshold.
//! \param nmsThreshold Threshold to be used in NMS.
//! \param codeType Type of coding method for bbox.
//! \param inputOrder Specifies the order of inputs {loc_data, conf_data, priorbox_data}.
//! \param confSigmoid Set to true to calculate sigmoid of confidence scores.
//! \param isNormalized Set to true if bounding box data is normalized by the network.
//!
struct DetectionOutputParameters
{
bool shareLocation, varianceEncodedInTarget;
int backgroundLabelId, numClasses, topK, keepTopK;
float confidenceThreshold, nmsThreshold;
CodeTypeSSD codeType;
int inputOrder[3];
bool confSigmoid;
bool isNormalized;
};
//!
//! \param param Set of parameters for creating the DetectionOutput plugin layer.
//! \deprecated. This plugin is superseded by createNMSPlugin()
//!
TENSORRTAPI INvPlugin* createSSDDetectionOutputPlugin(DetectionOutputParameters param);
TENSORRTAPI INvPlugin* createSSDDetectionOutputPlugin(const void* data, size_t length);
//!
//! \brief The Concat plugin layer basically performs the concatention for 4D tensors. Unlike the Concatenation layer in early version of TensorRT,
//! it allows the user to specify the axis along which to concatenate. The axis can be 1 (across channel), 2 (across H), or 3 (across W).
//! More particularly, this Concat plugin layer also implements the "ignoring the batch dimension" switch. If turned on, all the input tensors will be treated as if their batch sizes were 1.
//! \param concatAxis Axis along which to concatenate. Can't be the "N" dimension.
//! \param ignoreBatch If true, all the input tensors will be treated as if their batch sizes were 1.
//! \deprecated. This plugin is superseded by native TensorRT concatenation layer
//!
TENSORRTAPI INvPlugin* createConcatPlugin(int concatAxis, bool ignoreBatch);
TENSORRTAPI INvPlugin* createConcatPlugin(const void* data, size_t length);
//!
//! \brief The PReLu plugin layer performs leaky ReLU for 4D tensors. Give an input value x, the PReLU layer computes the output as x if x > 0
//! and negative_slope //! x if x <= 0.
//! \param negSlope Negative_slope value.
//! \deprecated. This plugin is superseded by createLReLUPlugin()
//!
TENSORRTAPI INvPlugin* createPReLUPlugin(float negSlope);
TENSORRTAPI INvPlugin* createPReLUPlugin(const void* data, size_t length);
//!
//! \brief The Reorg plugin layer maps the 512x26x26 feature map onto a 2048x13x13 feature map, so that it can be concatenated with the feature maps at 13x13 resolution.
//! \param stride Strides in H and W.
//! \deprecated. This plugin is superseded by createReorgPlugin()
//!
TENSORRTAPI INvPlugin* createYOLOReorgPlugin(int stride);
TENSORRTAPI INvPlugin* createYOLOReorgPlugin(const void* data, size_t length);
//!
//! \brief The Region plugin layer performs region proposal calculation: generate 5 bounding boxes per cell (for yolo9000, generate 3 bounding boxes per cell).
//! For each box, calculating its probablities of objects detections from 80 pre-defined classifications (yolo9000 has 9416 pre-defined classifications,
//! and these 9416 items are organized as work-tree structure).
//! RegionParameters defines a set of parameters for creating the Region plugin layer.
//! \param num Number of predicted bounding box for each grid cell.
//! \param coords Number of coordinates for a bounding box.
//! \param classes Number of classfications to be predicted.
//! \param softmaxTree When performing yolo9000, softmaxTree is helping to do softmax on confidence scores, for element to get the precise classfication through word-tree structured classfication definition.
//! \deprecated. This plugin is superseded by createRegionPlugin()
//!
typedef struct
{
int* leaf;
int n;
int* parent;
int* child;
int* group;
char** name;
int groups;
int* groupSize;
int* groupOffset;
} softmaxTree; // softmax tree
struct RegionParameters
{
int num;
int coords;
int classes;
softmaxTree* smTree;
};
TENSORRTAPI INvPlugin* createYOLORegionPlugin(RegionParameters params);
TENSORRTAPI INvPlugin* createYOLORegionPlugin(const void* data, size_t length);
//!
//! \brief The NMSParameters are used by the BatchedNMSPlugin for performing
//! the non_max_suppression operation over boxes for object detection networks.
//! \param 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.
//! \param backgroundLabelId Label ID for the background class. If there is no background class, set it as -1
//! \param numClasses Number of classes in the network.
//! \param topK Number of bounding boxes to be fed into the NMS step.
//! \param keepTopK Number of total bounding boxes to be kept per image after NMS step.
//! Should be less than or equal to the topK value.
//! \param scoreThreshold Scalar threshold for score (low scoring boxes are removed).
//! \param iouThreshold scalar threshold for IOU (new boxes that have high IOU overlap
//! with previously selected boxes are removed).
//! \param isNormalized Set to false, if the box coordinates are not
//! normalized, i.e. not in the range [0,1]. Defaults to true.
//!
struct NMSParameters
{
bool shareLocation;
int backgroundLabelId, numClasses, topK, keepTopK;
float scoreThreshold, iouThreshold;
bool isNormalized;
};
} // end plugin namespace
} // end nvinfer1 namespace
extern "C"
{
//!
//! \brief Create a plugin layer that fuses the RPN and ROI pooling using user-defined parameters.
//! Registered plugin type "RPROI_TRT". Registered plugin version "1".
//! \param featureStride Feature stride.
//! \param preNmsTop Number of proposals to keep before applying NMS.
//! \param nmsMaxOut Number of remaining proposals after applying NMS.
//! \param iouThreshold IoU threshold.
//! \param minBoxSize Minimum allowed bounding box size before scaling.
//! \param spatialScale Spatial scale between the input image and the last feature map.
//! \param pooling Spatial dimensions of pooled ROIs.
//! \param anchorRatios Aspect ratios for generating anchor windows.
//! \param anchorScales Scales for generating anchor windows.
//!
//! \return Returns a FasterRCNN fused RPN+ROI pooling plugin. Returns nullptr on invalid inputs.
//!
TENSORRTAPI nvinfer1::IPluginV2* createRPNROIPlugin(int featureStride, int preNmsTop,
int nmsMaxOut, float iouThreshold, float minBoxSize,
float spatialScale, nvinfer1::DimsHW pooling,
nvinfer1::Weights anchorRatios, nvinfer1::Weights anchorScales);
//!
//! \brief The Normalize plugin layer normalizes the input to have L2 norm of 1 with scale learnable.
//! Registered plugin type "Normalize_TRT". Registered plugin version "1".
//! \param scales Scale weights that are applied to the output tensor.
//! \param acrossSpatial Whether to compute the norm over adjacent channels (acrossSpatial is true) or nearby spatial locations (within channel in which case acrossSpatial is false).
//! \param channelShared Whether the scale weight(s) is shared across channels.
//! \param eps Epsilon for not diviiding by zero.
//!
TENSORRTAPI nvinfer1::IPluginV2* createNormalizePlugin(const nvinfer1::Weights* scales, bool acrossSpatial, bool channelShared, float eps);
//!
//! \brief The PriorBox plugin layer generates the prior boxes of designated sizes and aspect ratios across all dimensions (H x W).
//! PriorBoxParameters defines a set of parameters for creating the PriorBox plugin layer.
//! Registered plugin type "PriorBox_TRT". Registered plugin version "1".
//!
TENSORRTAPI nvinfer1::IPluginV2* createPriorBoxPlugin(nvinfer1::plugin::PriorBoxParameters param);
//!
//! \brief The Grid Anchor Generator plugin layer generates the prior boxes of
//! designated sizes and aspect ratios across all dimensions (H x W) for all feature maps.
//! GridAnchorParameters defines a set of parameters for creating the GridAnchorGenerator plugin layer.
//! Registered plugin type "GridAnchor_TRT". Registered plugin version "1".
//!
TENSORRTAPI nvinfer1::IPluginV2* createAnchorGeneratorPlugin(nvinfer1::plugin::GridAnchorParameters* param, int numLayers);
//!
//! \brief The DetectionOutput plugin layer generates the detection output based on location and confidence predictions by doing non maximum suppression.
//! DetectionOutputParameters defines a set of parameters for creating the DetectionOutput plugin layer.
//! Registered plugin type "NMS_TRT". Registered plugin version "1".
//!
TENSORRTAPI nvinfer1::IPluginV2* createNMSPlugin(nvinfer1::plugin::DetectionOutputParameters param);
//!
//! \brief The LReLu plugin layer performs leaky ReLU for 4D tensors. Give an input value x, the PReLU layer computes the output as x if x > 0 and negative_slope //! x if x <= 0.
//! Registered plugin type "LReLU_TRT". Registered plugin version "1".
//! \param negSlope Negative_slope value.
//!
TENSORRTAPI nvinfer1::IPluginV2* createLReLUPlugin(float negSlope);
//!
//! \brief The Reorg plugin reshapes input of shape CxHxW into a (C*stride*stride)x(H/stride)x(W/stride) shape, used in YOLOv2.
//! It does that by taking 1 x stride x stride slices from tensor and flattening them into (stridexstride) x 1 x 1 shape.
//! Registered plugin type "Reorg_TRT". Registered plugin version "1".
//! \param stride Strides in H and W, it should divide both H and W. Also stride * stride should be less than or equal to C.
//!
TENSORRTAPI nvinfer1::IPluginV2* createReorgPlugin(int stride);
//!
//! \brief The Region plugin layer performs region proposal calculation: generate 5 bounding boxes per cell (for yolo9000, generate 3 bounding boxes per cell).
//! For each box, calculating its probablities of objects detections from 80 pre-defined classifications (yolo9000 has 9416 pre-defined classifications,
//! and these 9416 items are organized as work-tree structure).
//! RegionParameters defines a set of parameters for creating the Region plugin layer.
//! Registered plugin type "Region_TRT". Registered plugin version "1".
//!
TENSORRTAPI nvinfer1::IPluginV2* createRegionPlugin(nvinfer1::plugin::RegionParameters params);
//!
//! \brief The Clip Plugin performs a clip operation on the input tensor. It
//! clips the tensor values to a specified min and max. Any value less than clipMin are set to clipMin.
//! Any values greater than clipMax are set to clipMax. For example, this plugin can be used
//! to perform a Relu6 operation by specifying clipMin=0.0 and clipMax=6.0
//! Registered plugin type "Clip_TRT". Registered plugin version "1".
//! \param layerName The name of the TensorRT layer.
//! \param clipMin The minimum value to clip to.
//! \param clipMax The maximum value to clip to.
//!
TENSORRTAPI nvinfer1::IPluginV2* createClipPlugin(const char* layerName, float clipMin, float clipMax);
//!
//! \brief The BatchedNMS Plugin performs non_max_suppression on the input boxes, per batch, across all classes.
//! It greedily selects a subset of bounding boxes in descending order of
//! score. Prunes away boxes that have a high intersection-over-union (IOU)
//! overlap with previously selected boxes. Bounding boxes are supplied as [y1, x1, y2, x2],
//! where (y1, x1) and (y2, x2) are the coordinates of any
//! diagonal pair of box corners and the coordinates can be provided as normalized
//! (i.e., lying in the interval [0, 1]) or absolute.
//! The plugin expects two inputs.
//! Input0 is expected to be 4-D float boxes tensor of shape [batch_size, num_boxes,
//! q, 4], where q can be either 1 (if shareLocation is true) or num_classes.
//! Input1 is expected to be a 3-D float scores tensor of shape [batch_size, num_boxes, num_classes]
//! representing a single score corresponding to each box.
//! The plugin returns four outputs.
//! num_detections : A [batch_size] int32 tensor indicating the number of valid
//! detections per batch item. Can be less than keepTopK. Only the top num_detections[i] entries in
//! nmsed_boxes[i], nmsed_scores[i] and nmsed_classes[i] are valid.
//! nmsed_boxes : A [batch_size, max_detections, 4] float32 tensor containing
//! the co-ordinates of non-max suppressed boxes.
//! nmsed_scores : A [batch_size, max_detections] float32 tensor containing the
//! scores for the boxes.
//! nmsed_classes : A [batch_size, max_detections] float32 tensor containing the
//! classes for the boxes.
//!
//! Registered plugin type "BatchedNMS_TRT". Registered plugin version "1".
//!
TENSORRTAPI nvinfer1::IPluginV2* createBatchedNMSPlugin(nvinfer1::plugin::NMSParameters param);
//!
//! \brief Initialize and register all the existing TensorRT plugins to the Plugin Registry with an optional namespace.
//! The plugin library author should ensure that this function name is unique to the library.
//! This function should be called once before accessing the Plugin Registry.
//! \param logger Logger object to print plugin registration information
//! \param libNamespace Namespace used to register all the plugins in this library
//!
TENSORRTAPI bool initLibNvInferPlugins(void* logger, const char* libNamespace);
} // extern "C"
#endif // NV_INFER_PLUGIN_H
include/NvInferPluginUtils.h
+61
View File
@@ -0,0 +1,61 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_INFER_PLUGIN_UTILS_H
#define NV_INFER_PLUGIN_UTILS_H
//!
//! \file NvPluginUtils.h
//!
//! This is the API for the Nvidia provided TensorRT plugin utilities.
//! It lists all the parameters utilized by the TensorRT plugins.
//!
namespace nvinfer1
{
namespace plugin
{
//!
//! \brief RPROIParams is used to create the RPROIPlugin instance.
//! It contains:
//! \param poolingH Height of the output in pixels after ROI pooling on feature map.
//! \param poolingW Width of the output in pixels after ROI pooling on feature map.
//! \param featureStride Feature stride; ratio of input image size to feature map size. Assuming that max pooling layers in neural network use square filters.
//! \param preNmsTop Number of proposals to keep before applying NMS.
//! \param nmsMaxOut Number of remaining proposals after applying NMS.
//! \param anchorsRatioCount Number of anchor box ratios.
//! \param anchorsScaleCount Number of anchor box scales.
//! \param iouThreshold IoU (Intersection over Union) threshold used for the NMS step.
//! \param minBoxSize Minimum allowed bounding box size before scaling, used for anchor box calculation.
//! \param spatialScale Spatial scale between the input image and the last feature map.
//!
struct RPROIParams
{
int poolingH;
int poolingW;
int featureStride;
int preNmsTop;
int nmsMaxOut;
int anchorsRatioCount;
int anchorsScaleCount;
float iouThreshold;
float minBoxSize;
float spatialScale;
};
} // end plugin namespace
} // end nvinfer1 namespace
#endif
include/NvOnnxConfig.h
+188
View File
@@ -0,0 +1,188 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_ONNX_CONFIG_H
#define NV_ONNX_CONFIG_H
#include "NvInfer.h"
namespace nvonnxparser
{
//!
//! \mainpage
//!
//! This is the API documentation for the Configuration Manager for Open Neural Network Exchange (ONNX) Parser for Nvidia TensorRT Inference Engine.
//! It provides information on individual functions, classes
//! and methods. Use the index on the left to navigate the documentation.
//!
//! Please see the accompanying user guide and samples for higher-level information and general advice on using ONNX Parser and TensorRT.
//!
//!
//! \file NvOnnxConfig.h
//!
//! This is the API file for the Configuration Manager for ONNX Parser for Nvidia TensorRT.
//!
//!
//! \class IOnnxConfig
//! \brief Configuration Manager Class.
//!
class IOnnxConfig
{
protected:
virtual ~IOnnxConfig() {}
public:
//!
//! \typedef Verbosity
//! \brief Defines Verbosity level.
//!
typedef int Verbosity;
//!
//! \brief Set the Model Data Type.
//!
//! Sets the Model DataType, one of the following: float -d 32 (default), half precision -d 16, and int8 -d 8 data types.
//!
//! \see getModelDtype()
//!
virtual void setModelDtype(const nvinfer1::DataType) = 0;
//!
//! \brief Get the Model Data Type.
//!
//! \return DataType nvinfer1::DataType
//!
//! \see setModelDtype() and #DataType
//!
virtual nvinfer1::DataType getModelDtype() const = 0;
//!
//! \brief Get the Model FileName.
//!
//! \return Return the Model Filename, as a pointer to a NULL-terminated character sequence.
//!
//! \see setModelFileName()
//!
virtual const char* getModelFileName() const = 0;
//!
//! \brief Set the Model File Name.
//!
//! The Model File name contains the Network Description in ONNX pb format.
//!
//! This method copies the name string.
//!
//! \param onnxFilename The name.
//!
//! \see getModelFileName()
//!
virtual void setModelFileName(const char* onnxFilename) = 0;
//!
//! \brief Get the Verbosity Level.
//!
//! \return The Verbosity Level.
//!
//! \see addVerbosity(), reduceVerbosity()
//!
virtual Verbosity getVerbosityLevel() const = 0;
//!
//! \brief Increase the Verbosity Level.
//!
//! \return The Verbosity Level.
//!
//! \see addVerbosity(), reduceVerbosity(), setVerbosity(Verbosity)
//!
virtual void addVerbosity() = 0; //!< Increase verbosity Level.
virtual void reduceVerbosity() = 0; //!< Decrease verbosity Level.
virtual void setVerbosityLevel(Verbosity) = 0; //!< Set to specific verbosity Level.
//!
//! \brief Returns the File Name of the Network Description as a Text File.
//!
//! \return Return the name of the file containing the network description converted to a plain text, used for debugging purposes.
//!
//! \see setTextFilename()
//!
virtual const char* getTextFileName() const = 0;
//!
//! \brief Set the File Name of the Network Description as a Text File.
//!
//! This API allows setting a file name for the network description in plain text, equivalent of the ONNX protobuf.
//!
//! This method copies the name string.
//!
//! \param textFileName Name of the file.
//!
//! \see getTextFilename()
//!
virtual void setTextFileName(const char* textFileName) = 0;
//!
//! \brief Get the File Name of the Network Description as a Text File, including the weights.
//!
//! \return Return the name of the file containing the network description converted to a plain text, used for debugging purposes.
//!
//! \see setFullTextFilename()
//!
virtual const char* getFullTextFileName() const = 0;
//!
//! \brief Set the File Name of the Network Description as a Text File, including the weights.
//!
//! This API allows setting a file name for the network description in plain text, equivalent of the ONNX protobuf.
//!
//! This method copies the name string.
//!
//! \param fullTextFileName Name of the file.
//!
//! \see getFullTextFilename()
//!
virtual void setFullTextFileName(const char* fullTextFileName) = 0;
//!
//! \brief Get whether the layer information will be printed.
//!
//! \return Returns whether the layer information will be printed.
//!
//! \see setPrintLayerInfo()
//!
virtual bool getPrintLayerInfo() const = 0;
//!
//! \brief Set whether the layer information will be printed.
//!
//! \see getPrintLayerInfo()
//!
virtual void setPrintLayerInfo(bool) = 0;
//!
//! \brief Destroy IOnnxConfig object.
//!
virtual void destroy() = 0;
}; // class IOnnxConfig
TENSORRTAPI IOnnxConfig* createONNXConfig();
} // namespace nvonnxparser
#endif
include/NvOnnxParser.h
+225
View File
@@ -0,0 +1,225 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_ONNX_PARSER_H
#define NV_ONNX_PARSER_H
#include "NvInfer.h"
#define NV_ONNX_PARSER_MAJOR 0
#define NV_ONNX_PARSER_MINOR 1
#define NV_ONNX_PARSER_PATCH 0
static const int NV_ONNX_PARSER_VERSION = ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH);
class onnxTensorDescriptorV1;
namespace nvonnxparser
{
template <typename T>
inline int EnumMax();
/** \enum ErrorCode
*
* \brief the type of parser error
*/
enum class ErrorCode : int
{
kSUCCESS = 0,
kINTERNAL_ERROR = 1,
kMEM_ALLOC_FAILED = 2,
kMODEL_DESERIALIZE_FAILED = 3,
kINVALID_VALUE = 4,
kINVALID_GRAPH = 5,
kINVALID_NODE = 6,
kUNSUPPORTED_GRAPH = 7,
kUNSUPPORTED_NODE = 8
};
template <>
inline int EnumMax<ErrorCode>()
{
return 9;
}
/** \class IParserError
*
* \brief an object containing information about an error
*/
class IParserError
{
public:
/** \brief the error code
*/
virtual ErrorCode code() const = 0;
/** \brief description of the error
*/
virtual const char* desc() const = 0;
/** \brief source file in which the error occurred
*/
virtual const char* file() const = 0;
/** \brief source line at which the error occurred
*/
virtual int line() const = 0;
/** \brief source function in which the error occurred
*/
virtual const char* func() const = 0;
/** \brief index of the ONNX model node in which the error occurred
*/
virtual int node() const = 0;
protected:
virtual ~IParserError() {}
};
/** \class IParser
*
* \brief an object for parsing ONNX models into a TensorRT network definition
*/
class IParser
{
public:
/** \brief Parse a serialized ONNX model into the TensorRT network.
* This method has very limited diagnostic. If parsing the serialized model
* fails for any reason (e.g. unsupported IR version, unsupported opset, etc.)
* it the user responsibility to intercept and report the error.
* To obtain a better diagnostic, use the parseFromFile method below.
*
* \param serialized_onnx_model Pointer to the serialized ONNX model
* \param serialized_onnx_model_size Size of the serialized ONNX model
* in bytes
* \return true if the model was parsed successfully
* \see getNbErrors() getError()
*/
virtual bool parse(void const* serialized_onnx_model,
size_t serialized_onnx_model_size)
= 0;
/** \brief Parse an onnx model file, can be a binary protobuf or a text onnx model
* calls parse method inside.
*
* \param File name
* \param Verbosity Level
*
* \return true if the model was parsed successfully
*
*/
virtual bool parseFromFile(const char* onnxModelFile, int verbosity) = 0;
/** \brief Check whether TensorRT supports a particular ONNX model
*
* \param serialized_onnx_model Pointer to the serialized ONNX model
* \param serialized_onnx_model_size Size of the serialized ONNX model
* in bytes
* \return true if the model is supported
*/
virtual bool supportsModel(void const* serialized_onnx_model,
size_t serialized_onnx_model_size)
= 0;
/** \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
* \param serialized_onnx_model_size Size of the serialized ONNX model
* in bytes
* \param weight_count number of user provided weights
* \param weight_descriptors pointer to user provided weight array
* \return true if the model was parsed successfully
* \see getNbErrors() getError()
*/
virtual bool parseWithWeightDescriptors(
void const* serialized_onnx_model, size_t serialized_onnx_model_size,
uint32_t weight_count,
onnxTensorDescriptorV1 const* weight_descriptors)
= 0;
/** \brief Returns whether the specified operator may be supported by the
* parser.
*
* Note that a result of true does not guarantee that the operator will be
* supported in all cases (i.e., this function may return false-positives).
*
* \param op_name The name of the ONNX operator to check for support
*/
virtual bool supportsOperator(const char* op_name) const = 0;
/** \brief destroy this object
*/
virtual void destroy() = 0;
/** \brief Get the number of errors that occurred during prior calls to
* \p parse
*
* \see getError() clearErrors() IParserError
*/
virtual int getNbErrors() const = 0;
/** \brief Get an error that occurred during prior calls to \p parse
*
* \see getNbErrors() clearErrors() IParserError
*/
virtual IParserError const* getError(int index) const = 0;
/** \brief Clear errors from prior calls to \p parse
*
* \see getNbErrors() getError() IParserError
*/
virtual void clearErrors() = 0;
protected:
virtual ~IParser() {}
};
} // namespace nvonnxparser
extern "C" TENSORRTAPI void* createNvOnnxParser_INTERNAL(void* network, void* logger, int version);
extern "C" TENSORRTAPI int getNvOnnxParserVersion();
namespace nvonnxparser
{
#ifdef SWIG
inline IParser* createParser(nvinfer1::INetworkDefinition* network,
nvinfer1::ILogger* logger)
{
return static_cast<IParser*>(
createNvOnnxParser_INTERNAL(network, logger, NV_ONNX_PARSER_VERSION));
}
#endif // SWIG
namespace
{
/** \brief Create a new parser object
*
* \param network The network definition that the parser will write to
* \param logger The logger to use
* \return a new parser object or NULL if an error occurred
* \see IParser
*/
#ifdef _MSC_VER
TENSORRTAPI IParser* createParser(nvinfer1::INetworkDefinition& network,
nvinfer1::ILogger& logger)
#else
inline IParser* createParser(nvinfer1::INetworkDefinition& network,
nvinfer1::ILogger& logger)
#endif
{
return static_cast<IParser*>(
createNvOnnxParser_INTERNAL(&network, &logger, NV_ONNX_PARSER_VERSION));
}
} // namespace
} // namespace nvonnxparser
#endif // NV_ONNX_PARSER_H
include/NvOnnxParserRuntime.h
+79
View File
@@ -0,0 +1,79 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_ONNX_PARSER_RUNTIME_H
#define NV_ONNX_PARSER_RUNTIME_H
#include "NvOnnxParser.h"
namespace nvonnxparser
{
/** \class IPluginFactory
*
* \brief a destroyable plugin factory object
*/
class IPluginFactory : public nvinfer1::IPluginFactory
{
public:
/** \brief destroy this object
*/
virtual void destroy() = 0;
protected:
virtual ~IPluginFactory() {}
};
} // namespace nvonnxparser
extern "C" TENSORRTAPI void* createNvOnnxParserPluginFactory_INTERNAL(void* logger, int version);
namespace nvonnxparser
{
#ifdef SWIG
inline IPluginFactory* createPluginFactory(nvinfer1::ILogger* logger)
{
return static_cast<IPluginFactory*>(
createNvOnnxParserPluginFactory_INTERNAL(logger, NV_ONNX_PARSER_VERSION));
}
#endif // SWIG
namespace
{
/** \brief Create a new plugin factory for deserializing engines built using
* the ONNX parser.
*
* This plugin factory handles deserialization of the plugins that are built
* into the ONNX parser. Engines built using the ONNX parser must use this
* plugin factory during deserialization.
*
* \param logger The logger to use
*
* \return a new plugin factory object or NULL if an error occurred
* \see IPluginFactory
*/
inline IPluginFactory* createPluginFactory(nvinfer1::ILogger& logger)
{
return static_cast<IPluginFactory*>(
createNvOnnxParserPluginFactory_INTERNAL(&logger, NV_ONNX_PARSER_VERSION));
}
} // namespace
} // namespace nvonnxparser
#endif // NV_ONNX_PARSER_RUNTIME_H
include/NvUffParser.h
+242
View File
@@ -0,0 +1,242 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_UFF_PARSER_H
#define NV_UFF_PARSER_H
#include "NvInfer.h"
//Current supported Universal Framework Format (UFF) version for the parser.
#define UFF_REQUIRED_VERSION_MAJOR 0
#define UFF_REQUIRED_VERSION_MINOR 6
#define UFF_REQUIRED_VERSION_PATCH 3
namespace nvuffparser
{
//!
//! \enum UffInputOrder
//! \brief The different possible supported input order.
//!
enum class UffInputOrder : int
{
kNCHW = 0, //!< NCHW order.
kNHWC = 1, //!< NHWC order.
kNC = 2 //!< NC order.
};
//!
//! \enum FieldType
//! \brief The possible field types for custom layer.
//!
enum class FieldType : int
{
kFLOAT = 0, //!< FP32 field type.
kINT32 = 1, //!< INT32 field type.
kCHAR = 2, //!< char field type. String for length>1.
kDIMS = 4, //!< nvinfer1::Dims field type.
kDATATYPE = 5, //!< nvinfer1::DataType field type.
kUNKNOWN = 6
};
//!
//! \class FieldMap
//!
//! \brief An array of field params used as a layer parameter for plugin layers.
//!
//! The node fields are passed by the parser to the API through the plugin
//! constructor. The implementation of the plugin should parse the contents of
//! the fieldMap as part of the plugin constructor
//!
class TENSORRTAPI FieldMap
{
public:
const char* name;
const void* data;
FieldType type = FieldType::kUNKNOWN;
int length = 1;
FieldMap(const char* name, const void* data, const FieldType type, int length = 1);
};
struct FieldCollection
{
int nbFields;
const FieldMap* fields;
};
//!
//! \class IPluginFactory
//!
//! \brief Plugin factory used to configure plugins.
//!
class IPluginFactory
{
public:
//!
//! \brief A user implemented function that determines if a layer configuration is provided by an IPlugin.
//!
//! \param layerName Name of the layer which the user wishes to validate.
//!
virtual bool isPlugin(const char* layerName) = 0;
//!
//! \brief Creates a plugin.
//!
//! \param layerName Name of layer associated with the plugin.
//! \param weights Weights used for the layer.
//! \param nbWeights Number of weights.
//! \param fc A collection of FieldMaps used as layer parameters for different plugin layers.
//!
//! \see FieldCollection
//!
virtual nvinfer1::IPlugin* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights,
const FieldCollection fc) = 0;
};
//!
//! \class IPluginFactoryExt
//!
//! \brief Plugin factory used to configure plugins with added support for TRT versioning.
//!
class IPluginFactoryExt : public IPluginFactory
{
public:
virtual int getVersion() const
{
return NV_TENSORRT_VERSION;
}
//!
//! \brief A user implemented function that determines if a layer configuration is provided by an IPluginExt.
//!
//! \param layerName Name of the layer which the user wishes to validate.
//!
virtual bool isPluginExt(const char* layerName) = 0;
};
//!
//! \class IUffParser
//!
//! \brief Class used for parsing models described using the UFF format.
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IUffParser
{
public:
//!
//! \brief Register an input name of a UFF network with the associated Dimensions.
//!
//! \param inputName Input name.
//! \param inputDims Input dimensions.
//! \param inputOrder Input order on which the framework input was originally.
//!
virtual bool registerInput(const char* inputName, nvinfer1::Dims inputDims, UffInputOrder inputOrder) = 0;
//!
//! \brief Register an output name of a UFF network.
//!
//! \param outputName Output name.
//!
virtual bool registerOutput(const char* outputName) = 0;
//!
//! \brief Parse a UFF file.
//!
//! \param file File name of the UFF file.
//! \param network Network in which the UFFParser will fill the layers.
//! \param weightsType The type on which the weights will transformed in.
//!
virtual bool parse(const char* file,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightsType=nvinfer1::DataType::kFLOAT) = 0;
//!
//! \brief Parse a UFF buffer, useful if the file already live in memory.
//!
//! \param buffer Buffer of the UFF file.
//! \param size Size of buffer of the UFF file.
//! \param network Network in which the UFFParser will fill the layers.
//! \param weightsType The type on which the weights will transformed in.
//!
virtual bool parseBuffer(const char* buffer, std::size_t size,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightsType=nvinfer1::DataType::kFLOAT) = 0;
virtual void destroy() = 0;
//!
//! \brief Return Version Major of the UFF.
//!
virtual int getUffRequiredVersionMajor() = 0;
//!
//! \brief Return Version Minor of the UFF.
//!
virtual int getUffRequiredVersionMinor() = 0;
//!
//! \brief Return Patch Version of the UFF.
//!
virtual int getUffRequiredVersionPatch() = 0;
//!
//! \brief Set the IPluginFactory used to create the user defined plugins.
//!
//! \param factory Pointer to an instance of the user implmentation of IPluginFactory.
//!
virtual void setPluginFactory(IPluginFactory* factory) = 0;
//!
//! \brief Set the IPluginFactoryExt used to create the user defined pluginExts.
//!
//! \param factory Pointer to an instance of the user implmentation of IPluginFactoryExt.
//!
virtual void setPluginFactoryExt(IPluginFactoryExt* factory) = 0;
//!
//! \brief Set the namespace used to lookup and create plugins in the network.
//!
virtual void setPluginNamespace(const char* libNamespace) = 0;
protected:
virtual ~IUffParser() {}
};
//!
//! \brief Creates a IUffParser object.
//!
//! \return A pointer to the IUffParser object is returned.
//!
//! \see nvuffparser::IUffParser
//!
TENSORRTAPI IUffParser* createUffParser();
//!
//! \brief Shuts down protocol buffers library.
//!
//! \note No part of the protocol buffers library can be used after this function is called.
//!
TENSORRTAPI void shutdownProtobufLibrary(void);
}
extern "C" TENSORRTAPI void* createNvUffParser_INTERNAL();
#endif /* !NV_UFF_PARSER_H */
include/NvUtils.h
+129
View File
@@ -0,0 +1,129 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef NV_UTILS_H
#define NV_UTILS_H
#include "NvInfer.h"
namespace nvinfer1
{
namespace utils
{
//!
//! \param input The input weights to reshape.
//! \param shape The shape of the weights.
//! \param shapeOrder The order of the dimensions to process for the output.
//! \param data The location where the output data is placed.
//! \param nbDims The number of dimensions to process.
//!
//! \brief Reformat the input weights of the given shape based on the new
//! order of dimensions.
//!
//! Take the weights specified by \p input with the dimensions specified by
//! \p shape and re-order the weights based on the new dimensions specified
//! by \p shapeOrder. The size of each dimension and the input data is not
//! modified. The output volume pointed to by \p data must be the same as
//! he \p input volume.
//!
//! Example usage:
//! float *out = new float[N*C*H*W];
//! Weights input{DataType::kFLOAT, {0 ... N*C*H*W-1}, N*C*H*W size};
//! int order[4]{1, 0, 3, 2};
//! int shape[4]{C, N, W, H};
//! reshapeWeights(input, shape, order, out, 4);
//! Weights reshaped{input.type, out, input.count};
//!
//! Input Matrix{3, 2, 3, 2}:
//! { 0 1}, { 2 3}, { 4 5} <-- {0, 0, *, *}
//! { 6 7}, { 8 9}, {10 11} <-- {0, 1, *, *}
//! {12 13}, {14 15}, {16 17} <-- {1, 0, *, *}
//! {18 19}, {20 21}, {22 23} <-- {1, 1, *, *}
//! {24 25}, {26 27}, {28 29} <-- {2, 0, *, *}
//! {30 31}, {32 33}, {34 35} <-- {2, 1, *, *}
//!
//! Output Matrix{2, 3, 2, 3}:
//! { 0 2 4}, { 1 3 5} <-- {0, 0, *, *}
//! {12 14 16}, {13 15 17} <-- {0, 1, *, *}
//! {24 26 28}, {25 27 29} <-- {0, 2, *, *}
//! { 6 8 10}, { 7 9 11} <-- {1, 0, *, *}
//! {18 20 22}, {19 21 23} <-- {1, 1, *, *}
//! {30 32 34}, {31 33 35} <-- {1, 2, *, *}
//!
//! \return True on success, false on failure.
//!
TENSORRTAPI bool reshapeWeights(const Weights &input, const int *shape, const int *shapeOrder, void *data, int nbDims);
//!
//! \param input The input data to re-order.
//! \param order The new order of the data sub-buffers.
//! \param num The number of data sub-buffers to re-order.
//! \param size The size of each data sub-buffer in bytes.
//!
//! \brief Takes an input stream and re-orders \p num chunks of the data
//! given the \p size and \p order.
//!
//! In some frameworks, the ordering of the sub-buffers within a dimension
//! is different than the way that TensorRT expects them.
//! TensorRT expects the gate/bias sub-buffers for LSTM's to be in fico order.
//! TensorFlow however formats the sub-buffers in icfo order.
//! This helper function solves this in a generic fashion.
//!
//! Example usage output of reshapeWeights above:
//! int indir[1]{1, 0}
//! int stride = W*H;
//! for (int x = 0, y = N*C; x < y; ++x)
//! reorderSubBuffers(out + x * stride, indir, H, W);
//!
//! Input Matrix{2, 3, 2, 3}:
//! { 0 2 4}, { 1 3 5} <-- {0, 0, *, *}
//! {12 14 16}, {13 15 17} <-- {0, 1, *, *}
//! {24 26 28}, {25 27 29} <-- {0, 2, *, *}
//! { 6 8 10}, { 7 9 11} <-- {1, 0, *, *}
//! {18 20 22}, {19 21 23} <-- {1, 1, *, *}
//! {30 32 34}, {31 33 35} <-- {1, 2, *, *}
//!
//! Output Matrix{2, 3, 2, 3}:
//! { 1 3 5}, { 0 2 4} <-- {0, 0, *, *}
//! {13 15 17}, {12 14 16} <-- {0, 1, *, *}
//! {25 27 29}, {24 26 28} <-- {0, 2, *, *}
//! { 7 9 11}, { 6 8 10} <-- {1, 0, *, *}
//! {19 21 23}, {18 20 22} <-- {1, 1, *, *}
//! {31 33 35}, {30 32 34} <-- {1, 2, *, *}
//!
//! \return True on success, false on failure.
//!
//! \see reshapeWeights()
//!
TENSORRTAPI bool reorderSubBuffers(void *input, const int *order, int num, int size);
//!
//! \param input The input data to transpose.
//! \param type The type of the data to transpose.
//! \param num The number of data sub-buffers to transpose.
//! \param height The size of the height dimension to transpose.
//! \param width The size of the width dimension to transpose.
//!
//! \brief Transpose \p num sub-buffers of \p height * \p width.
//!
//! \return True on success, false on failure.
//!
TENSORRTAPI bool transposeSubBuffers(void *input, DataType type, int num, int height, int width);
} // utils namespace
} // nvinfer1 namespace
#endif // NV_UTILS_H
parsers/CMakeLists.txt
+42
View File
@@ -0,0 +1,42 @@
#
# Copyright (c) 2019, 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.
#
add_custom_target(parsers DEPENDS
nvcaffeparserlibs
nvonnxparser
if (NVINTERNAL OR NVPARTNER)
nvuffparserlibs
nvparserslibs
endif()
)
add_subdirectory(caffe)
if (NVINTERNAL OR NVPARTNER)
add_subdirectory(uff)
include(uff/NvParsersCMakeLists.txt)
endif()
add_definitions("-D_PROTOBUF_INSTALL_DIR=${Protobuf_INSTALL_DIR}")
add_compile_options("-Dgoogle=google_private")
set(TENSORRT_ROOT ${PROJECT_SOURCE_DIR})
set(TENSORRT_BUILD ${TRT_BIN_DIR} ${TRT_LIB_DIR})
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${TRT_BIN_DIR})
include_directories(
${Protobuf_INCLUDE_DIR}
)
add_subdirectory(onnx)
parsers/caffe/CMakeLists.txt
+142
View File
@@ -0,0 +1,142 @@
#
# Copyright (c) 2019, 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.
#
############################# GENERATE C++ PROTO FILES ###################################
protobuf_generate_cpp(CAFFE_PROTO_SRC CAFFE_PROTO_HDR proto/trtcaffe.proto)
add_custom_target(caffe_proto
DEPENDS
${CAFFE_PROTO_SRC} ${CAFFE_PROTO_HDR}
)
############################## BUILD CAFFE PARSER ########################################
add_custom_target(nvcaffeparserlibs)
set(TARGET_NAME nvcaffeparser)
set(SHARED_TARGET ${TARGET_NAME})
set(STATIC_TARGET ${TARGET_NAME}_static)
################################# DEFINE SOURCES ########################################
include(CaffeParserSources.txt)
#########################################################################################
################################## SHARED LIBRARY #######################################
add_library(${SHARED_TARGET} SHARED
${CAFFE_PARSER_SRCS}
)
add_dependencies(${SHARED_TARGET} caffe_proto)
target_include_directories(${SHARED_TARGET}
PUBLIC ${PROJECT_SOURCE_DIR}/include
PRIVATE .
PRIVATE caffeParser
PRIVATE caffeParser/opParsers
PRIVATE caffeWeightFactory
PRIVATE ../common
PRIVATE ${Protobuf_INCLUDE_DIR}
PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/proto
)
set_target_properties(${SHARED_TARGET}
PROPERTIES
CXX_STANDARD 11
CXX_STANDARD_REQUIRED YES
CXX_EXTENSIONS NO
ARCHIVE_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
LIBRARY_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
RUNTIME_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
)
target_link_libraries(${SHARED_TARGET}
${Protobuf_LIBRARY}
)
# modify google namespace to avoid namespace collision.
set(GOOGLE google_private)
target_compile_definitions(${SHARED_TARGET}
PRIVATE
"-Dgoogle=${GOOGLE}"
"-DGOOGLE_PROTOBUF_ARCH_64_BIT"
)
set_target_properties(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-Wl,--exclude-libs,ALL")
set_target_properties(${SHARED_TARGET} PROPERTIES DEBUG_POSTFIX ${TRT_DEBUG_POSTFIX})
set_target_properties(${SHARED_TARGET} PROPERTIES VERSION ${TRT_VERSION} SOVERSION ${TRT_SOVERSION} )
set_property(TARGET ${SHARED_TARGET} PROPERTY CUDA_STANDARD 11)
################################## STATIC LIBRARY #######################################
add_library(${STATIC_TARGET} STATIC
${CAFFE_PARSER_SRCS}
)
add_dependencies(${STATIC_TARGET} caffe_proto)
target_include_directories(${STATIC_TARGET}
PUBLIC ${PROJECT_SOURCE_DIR}/include
PRIVATE .
PRIVATE caffeParser
PRIVATE caffeParser/opParsers
PRIVATE caffeWeightFactory
PRIVATE ../common
PRIVATE ${Protobuf_INCLUDE_DIR}
PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/proto
)
set_target_properties(${STATIC_TARGET}
PROPERTIES
CXX_STANDARD 11
CXX_STANDARD_REQUIRED YES
CXX_EXTENSIONS NO
ARCHIVE_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
LIBRARY_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
RUNTIME_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
)
target_link_libraries(${STATIC_TARGET}
${Protobuf_LIBRARY}
)
# modify google namespace to avoid namespace collision.
set(GOOGLE google_private)
target_compile_definitions(${STATIC_TARGET}
PRIVATE
"-Dgoogle=${GOOGLE}"
"-DGOOGLE_PROTOBUF_ARCH_64_BIT"
)
set_target_properties(${STATIC_TARGET} PROPERTIES LINK_FLAGS "-Wl,--exclude-libs,ALL")
set_target_properties(${STATIC_TARGET} PROPERTIES DEBUG_POSTFIX ${TRT_DEBUG_POSTFIX})
set_target_properties(${STATIC_TARGET} PROPERTIES VERSION ${TRT_VERSION} SOVERSION ${TRT_SOVERSION} )
set_property(TARGET ${STATIC_TARGET} PROPERTY CUDA_STANDARD 11)
#########################################################################################
add_dependencies(nvcaffeparserlibs ${SHARED_TARGET} ${STATIC_TARGET})
################################### INSTALLATION ########################################
install(TARGETS ${TARGET_NAME}
RUNTIME DESTINATION bin
LIBRARY DESTINATION lib
ARCHIVE DESTINATION lib
)
parsers/caffe/CaffeParserSources.txt
+45
View File
@@ -0,0 +1,45 @@
#
# Copyright (c) 2019, 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.
#
set(CAFFE_PARSER_SRCS
${CAFFE_PROTO_SRC}
caffeParser/opParsers/opParsers.h
caffeParser/opParsers/parseAbsVal.cpp
caffeParser/opParsers/parseBatchNorm.cpp
caffeParser/opParsers/parseBNLL.cpp
caffeParser/opParsers/parseClip.cpp
caffeParser/opParsers/parseConcat.cpp
caffeParser/opParsers/parseConv.cpp
caffeParser/opParsers/parseCrop.cpp
caffeParser/opParsers/parseDeconv.cpp
caffeParser/opParsers/parseEltwise.cpp
caffeParser/opParsers/parseELU.cpp
caffeParser/opParsers/parseInnerProduct.cpp
caffeParser/opParsers/parseLRN.cpp
caffeParser/opParsers/parsePermute.cpp
caffeParser/opParsers/parsePooling.cpp
caffeParser/opParsers/parsePower.cpp
caffeParser/opParsers/parseReduction.cpp
caffeParser/opParsers/parseReLU.cpp
caffeParser/opParsers/parseReshape.cpp
caffeParser/opParsers/parseScale.cpp
caffeParser/opParsers/parseSigmoid.cpp
caffeParser/opParsers/parseSoftMax.cpp
caffeParser/opParsers/parseTanH.cpp
caffeWeightFactory/caffeWeightFactory.cpp
caffeParser/caffeParser.cpp
NvCaffeParser.cpp
)
parsers/caffe/NvCaffeParser.cpp
+35
View File
@@ -0,0 +1,35 @@
/*
* Copyright (c) 2019, 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.
*/
#include "NvCaffeParser.h"
#include "caffeParser.h"
using namespace nvcaffeparser1;
void nvcaffeparser1::shutdownProtobufLibrary()
{
google::protobuf::ShutdownProtobufLibrary();
}
extern "C" void* createNvCaffeParser_INTERNAL()
{
return nvcaffeparser1::createCaffeParser();
}
ICaffeParser* nvcaffeparser1::createCaffeParser()
{
return new CaffeParser;
}
parsers/caffe/binaryProtoBlob.h
+66
View File
@@ -0,0 +1,66 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_BINARY_PROTO_BLOB_H
#define TRT_CAFFE_PARSER_BINARY_PROTO_BLOB_H
#include <stdlib.h>
#include "NvCaffeParser.h"
#include "NvInfer.h"
namespace nvcaffeparser1
{
class BinaryProtoBlob : public IBinaryProtoBlob
{
public:
BinaryProtoBlob(void* memory, nvinfer1::DataType type, nvinfer1::DimsNCHW dimensions)
: mMemory(memory)
, mDataType(type)
, mDimensions(dimensions)
{
}
nvinfer1::DimsNCHW getDimensions() override
{
return mDimensions;
}
nvinfer1::DataType getDataType() override
{
return mDataType;
}
const void* getData() override
{
return mMemory;
}
void destroy() override
{
delete this;
}
~BinaryProtoBlob() override
{
free(mMemory);
}
void* mMemory;
nvinfer1::DataType mDataType;
nvinfer1::DimsNCHW mDimensions;
};
} // namespace nvcaffeparser1
#endif // TRT_CAFFE_PARSER_BINARY_PROTO_BLOB_H
parsers/caffe/blobNameToTensor.h
+71
View File
@@ -0,0 +1,71 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_BLOB_NAME_TO_TENSOR_H
#define TRT_CAFFE_PARSER_BLOB_NAME_TO_TENSOR_H
#include <map>
#include <string>
#include "NvCaffeParser.h"
#include "NvInfer.h"
namespace nvcaffeparser1
{
class BlobNameToTensor : public IBlobNameToTensor
{
public:
void add(const std::string& name, nvinfer1::ITensor* tensor)
{
mMap[name] = tensor;
}
nvinfer1::ITensor* find(const char* name) const override
{
auto p = mMap.find(name);
if (p == mMap.end())
{
return nullptr;
}
return p->second;
}
nvinfer1::ITensor*& operator[](const std::string& name)
{
return mMap[name];
}
void setTensorNames()
{
for (auto& p : mMap)
{
p.second->setName(p.first.c_str());
}
}
~BlobNameToTensor() override = default;
bool isOK()
{
return !mError;
}
private:
std::map<std::string, nvinfer1::ITensor*> mMap;
bool mError{false};
};
} // namespace nvcaffeparser1
#endif // TRT_CAFFE_PARSER_BLOB_NAME_TO_TENSOR_H
parsers/caffe/caffeMacros.h
+46
View File
@@ -0,0 +1,46 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_MACROS_H
#define TRT_CAFFE_PARSER_MACROS_H
#ifdef _MSC_VER
#define FN_NAME __FUNCTION__
#else
#define FN_NAME __func__
#endif
#define CHECK_NULL(ptr) \
if ((ptr) == nullptr) \
{ \
std::cout << "Error: input " << #ptr << " is NULL in " << FN_NAME << std::endl; \
return; \
}
#define CHECK_NULL_RET_NULL(ptr) \
if ((ptr) == nullptr) \
{ \
std::cout << "Error: input " << #ptr << " is NULL in " << FN_NAME << std::endl; \
return nullptr; \
}
#define CHECK_NULL_RET_VAL(ptr, val) \
if ((ptr) == nullptr) \
{ \
std::cout << "Error: input " << #ptr << " is NULL in " << FN_NAME << std::endl; \
return val; \
}
#include "parserUtils.h"
#define RETURN_AND_LOG_ERROR(ret, message) RETURN_AND_LOG_ERROR_IMPL(ret, message, "CaffeParser: ")
#endif // TRT_CAFFE_PARSER_MACROS_H
parsers/caffe/caffeParser/caffeParser.cpp
+686
View File
@@ -0,0 +1,686 @@
/*
* Copyright (c) 2019, 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.
*/
#include <iostream>
#include "caffeMacros.h"
#include "caffeParser.h"
#include "opParsers.h"
#include "parserUtils.h"
#include "readProto.h"
#include "binaryProtoBlob.h"
#include "google/protobuf/text_format.h"
#include "half.h"
#include "NvInferPlugin.h"
using namespace nvinfer1;
using namespace nvcaffeparser1;
CaffeParser::~CaffeParser()
{
for (auto v : mTmpAllocs)
{
free(v);
}
for (auto p : mNewPlugins)
{
if (p)
{
p->destroy();
}
}
delete mBlobNameToTensor;
}
std::vector<nvinfer1::PluginField> CaffeParser::parseNormalizeParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
std::vector<nvinfer1::PluginField> f;
const trtcaffe::NormalizeParameter& p = msg.norm_param();
int* acrossSpatial = allocMemory<int32_t>();
*acrossSpatial = p.across_spatial() ? 1 : 0;
f.emplace_back("acrossSpatial", acrossSpatial, PluginFieldType::kINT32, 1);
int* channelShared = allocMemory<int32_t>();
*channelShared = p.channel_shared() ? 1 : 0;
f.emplace_back("channelShared", channelShared, PluginFieldType::kINT32, 1);
auto* eps = allocMemory<float>();
*eps = p.eps();
f.emplace_back("eps", eps, PluginFieldType::kFLOAT32, 1);
std::vector<Weights> w;
// If .caffemodel is not provided, need to randomize the weight
if (!weightFactory.isInitialized())
{
int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
w.emplace_back(weightFactory.allocateWeights(C, std::normal_distribution<float>(0.0F, 1.0F)));
}
else
{
// Use the provided weight from .caffemodel
w = weightFactory.getAllWeights(msg.name());
}
for (auto weight : w)
{
f.emplace_back("weights", weight.values, PluginFieldType::kFLOAT32, weight.count);
}
int* nbWeights = allocMemory<int32_t>();
*nbWeights = w.size();
f.emplace_back("nbWeights", nbWeights, PluginFieldType::kINT32, 1);
return f;
}
std::vector<nvinfer1::PluginField> CaffeParser::parsePriorBoxParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
std::vector<nvinfer1::PluginField> f;
const trtcaffe::PriorBoxParameter& p = msg.prior_box_param();
int minSizeSize = p.min_size_size();
auto* minSize = allocMemory<float>(minSizeSize);
for (int i = 0; i < minSizeSize; ++i)
{
minSize[i] = p.min_size(i);
}
f.emplace_back("minSize", minSize, PluginFieldType::kFLOAT32, minSizeSize);
int maxSizeSize = p.max_size_size();
auto* maxSize = allocMemory<float>(maxSizeSize);
for (int i = 0; i < maxSizeSize; ++i)
{
maxSize[i] = p.max_size(i);
}
f.emplace_back("maxSize", maxSize, PluginFieldType::kFLOAT32, maxSizeSize);
int aspectRatiosSize = p.aspect_ratio_size();
auto* aspectRatios = allocMemory<float>(aspectRatiosSize);
for (int i = 0; i < aspectRatiosSize; ++i)
{
aspectRatios[i] = p.aspect_ratio(i);
}
f.emplace_back("aspectRatios", aspectRatios, PluginFieldType::kFLOAT32, aspectRatiosSize);
int varianceSize = p.variance_size();
auto* variance = allocMemory<float>(varianceSize);
for (int i = 0; i < varianceSize; ++i)
{
variance[i] = p.variance(i);
}
f.emplace_back("variance", variance, PluginFieldType::kFLOAT32, varianceSize);
int* flip = allocMemory<int32_t>();
*flip = p.flip() ? 1 : 0;
f.emplace_back("flip", flip, PluginFieldType::kINT32, 1);
int* clip = allocMemory<int32_t>();
*clip = p.clip() ? 1 : 0;
f.emplace_back("clip", clip, PluginFieldType::kINT32, 1);
int* imgH = allocMemory<int32_t>();
*imgH = p.has_img_h() ? p.img_h() : p.img_size();
f.emplace_back("imgH", imgH, PluginFieldType::kINT32, 1);
int* imgW = allocMemory<int32_t>();
*imgW = p.has_img_w() ? p.img_w() : p.img_size();
f.emplace_back("imgW", imgW, PluginFieldType::kINT32, 1);
auto* stepH = allocMemory<float>();
*stepH = p.has_step_h() ? p.step_h() : p.step();
f.emplace_back("stepH", stepH, PluginFieldType::kFLOAT32, 1);
auto* stepW = allocMemory<float>();
*stepW = p.has_step_w() ? p.step_w() : p.step();
f.emplace_back("stepW", stepW, PluginFieldType::kFLOAT32, 1);
auto* offset = allocMemory<float>();
*offset = p.offset();
f.emplace_back("offset", offset, PluginFieldType::kFLOAT32, 1);
return f;
}
std::vector<nvinfer1::PluginField> CaffeParser::parseDetectionOutputParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
std::vector<nvinfer1::PluginField> f;
const trtcaffe::DetectionOutputParameter& p = msg.detection_output_param();
const trtcaffe::NonMaximumSuppressionParameter& nmsp = p.nms_param();
int* shareLocation = allocMemory<int32_t>();
*shareLocation = p.share_location() ? 1 : 0;
f.emplace_back("shareLocation", shareLocation, PluginFieldType::kINT32, 1);
int* varianceEncodedInTarget = allocMemory<int32_t>();
*varianceEncodedInTarget = p.variance_encoded_in_target() ? 1 : 0;
f.emplace_back("varianceEncodedInTarget", varianceEncodedInTarget, PluginFieldType::kINT32, 1);
int* backgroundLabelId = allocMemory<int32_t>();
*backgroundLabelId = p.background_label_id();
f.emplace_back("backgroundLabelId", backgroundLabelId, PluginFieldType::kINT32, 1);
int* numClasses = allocMemory<int32_t>();
*numClasses = p.num_classes();
f.emplace_back("numClasses", numClasses, PluginFieldType::kINT32, 1);
//nms
int* topK = allocMemory<int32_t>();
*topK = nmsp.top_k();
f.emplace_back("topK", topK, PluginFieldType::kINT32, 1);
int* keepTopK = allocMemory<int32_t>();
*keepTopK = p.keep_top_k();
f.emplace_back("keepTopK", keepTopK, PluginFieldType::kINT32, 1);
auto* confidenceThreshold = allocMemory<float>();
*confidenceThreshold = p.confidence_threshold();
f.emplace_back("confidenceThreshold", confidenceThreshold, PluginFieldType::kFLOAT32, 1);
//nms
auto* nmsThreshold = allocMemory<float>();
*nmsThreshold = nmsp.nms_threshold();
f.emplace_back("nmsThreshold", nmsThreshold, PluginFieldType::kFLOAT32, 1);
// input order = {0, 1, 2} in Caffe
int* inputOrder = allocMemory<int32_t>(3);
inputOrder[0] = 0;
inputOrder[1] = 1;
inputOrder[2] = 2;
f.emplace_back("inputOrder", inputOrder, PluginFieldType::kINT32, 3);
// confSigmoid = false for Caffe
int* confSigmoid = allocMemory<int32_t>();
*confSigmoid = 0;
f.emplace_back("confSigmoid", confSigmoid, PluginFieldType::kINT32, 1);
// isNormalized = true for Caffe
int* isNormalized = allocMemory<int32_t>();
*isNormalized = 1;
f.emplace_back("isNormalized", isNormalized, PluginFieldType::kINT32, 1);
// codeTypeSSD : from NvInferPlugin.h
// CORNER = 0, CENTER_SIZE = 1, CORNER_SIZE = 2, TF_CENTER = 3
int* codeType = allocMemory<int32_t>();
switch (p.code_type())
{
case trtcaffe::PriorBoxParameter::CORNER_SIZE:
*codeType = static_cast<int>(plugin::CodeTypeSSD::CORNER_SIZE);
break;
case trtcaffe::PriorBoxParameter::CENTER_SIZE:
*codeType = static_cast<int>(plugin::CodeTypeSSD::CENTER_SIZE);
break;
case trtcaffe::PriorBoxParameter::CORNER: // CORNER is default
default:
*codeType = static_cast<int>(plugin::CodeTypeSSD::CORNER);
break;
}
f.emplace_back("codeType", codeType, PluginFieldType::kINT32, 1);
return f;
}
std::vector<nvinfer1::PluginField> CaffeParser::parseLReLUParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
std::vector<nvinfer1::PluginField> f;
const trtcaffe::ReLUParameter& p = msg.relu_param();
auto* negSlope = allocMemory<float>();
*negSlope = p.negative_slope();
f.emplace_back("negSlope", negSlope, PluginFieldType::kFLOAT32, 1);
return f;
}
std::vector<nvinfer1::PluginField> CaffeParser::parseRPROIParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
std::vector<nvinfer1::PluginField> f;
const trtcaffe::ROIPoolingParameter& p1 = msg.roi_pooling_param();
const trtcaffe::RegionProposalParameter& p2 = msg.region_proposal_param();
// Memory allocations for plugin field variables
int* poolingH = allocMemory<int32_t>();
int* poolingW = allocMemory<int32_t>();
auto* spatialScale = allocMemory<float>();
int* preNmsTop = allocMemory<int32_t>();
int* nmsMaxOut = allocMemory<int32_t>();
auto* iouThreshold = allocMemory<float>();
auto* minBoxSize = allocMemory<float>();
int* featureStride = allocMemory<int32_t>();
int* anchorsRatioCount = allocMemory<int32_t>();
int* anchorsScaleCount = allocMemory<int32_t>();
int anchorsRatiosSize = p2.anchor_ratio_size();
auto* anchorsRatios = allocMemory<float>(anchorsRatiosSize);
int anchorsScalesSize = p2.anchor_scale_size();
auto* anchorsScales = allocMemory<float>(anchorsScalesSize);
// Intialize the plugin fields with values from the prototxt
*poolingH = p1.pooled_h();
f.emplace_back("poolingH", poolingH, PluginFieldType::kINT32, 1);
*poolingW = p1.pooled_w();
f.emplace_back("poolingW", poolingW, PluginFieldType::kINT32, 1);
*spatialScale = p1.spatial_scale();
f.emplace_back("spatialScale", spatialScale, PluginFieldType::kFLOAT32, 1);
*preNmsTop = p2.prenms_top();
f.emplace_back("preNmsTop", preNmsTop, PluginFieldType::kINT32, 1);
*nmsMaxOut = p2.nms_max_out();
f.emplace_back("nmsMaxOut", nmsMaxOut, PluginFieldType::kINT32, 1);
*iouThreshold = p2.iou_threshold();
f.emplace_back("iouThreshold", iouThreshold, PluginFieldType::kFLOAT32, 1);
*minBoxSize = p2.min_box_size();
f.emplace_back("minBoxSize", minBoxSize, PluginFieldType::kFLOAT32, 1);
*featureStride = p2.feature_stride();
f.emplace_back("featureStride", featureStride, PluginFieldType::kINT32, 1);
*anchorsRatioCount = p2.anchor_ratio_count();
f.emplace_back("anchorsRatioCount", anchorsRatioCount, PluginFieldType::kINT32, 1);
*anchorsScaleCount = p2.anchor_scale_count();
f.emplace_back("anchorsScaleCount", anchorsScaleCount, PluginFieldType::kINT32, 1);
for (int i = 0; i < anchorsRatiosSize; ++i) {
anchorsRatios[i] = p2.anchor_ratio(i);
}
f.emplace_back("anchorsRatios", anchorsRatios, PluginFieldType::kFLOAT32, anchorsRatiosSize);
for (int i = 0; i < anchorsScalesSize; ++i) {
anchorsScales[i] = p2.anchor_scale(i);
}
f.emplace_back("anchorsScales", anchorsScales, PluginFieldType::kFLOAT32, anchorsScalesSize);
return f;
}
const IBlobNameToTensor* CaffeParser::parseBuffers(const char* deployBuffer,
std::size_t deployLength,
const char* modelBuffer,
std::size_t modelLength,
INetworkDefinition& network,
DataType weightType)
{
mDeploy = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
google::protobuf::io::ArrayInputStream deployStream(deployBuffer, deployLength);
if (!google::protobuf::TextFormat::Parse(&deployStream, mDeploy.get()))
{
RETURN_AND_LOG_ERROR(nullptr, "Could not parse deploy file");
}
if (modelBuffer)
{
mModel = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
google::protobuf::io::ArrayInputStream modelStream(modelBuffer, modelLength);
google::protobuf::io::CodedInputStream codedModelStream(&modelStream);
codedModelStream.SetTotalBytesLimit(modelLength, -1);
if (!mModel->ParseFromCodedStream(&codedModelStream))
{
RETURN_AND_LOG_ERROR(nullptr, "Could not parse model file");
}
}
return parse(network, weightType, modelBuffer != nullptr);
}
const IBlobNameToTensor* CaffeParser::parse(const char* deployFile,
const char* modelFile,
INetworkDefinition& network,
DataType weightType)
{
CHECK_NULL_RET_NULL(deployFile)
// this is used to deal with dropout layers which have different input and output
mModel = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
if (modelFile && !readBinaryProto(mModel.get(), modelFile, mProtobufBufferSize))
{
RETURN_AND_LOG_ERROR(nullptr, "Could not parse model file");
}
mDeploy = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
if (!readTextProto(mDeploy.get(), deployFile))
{
RETURN_AND_LOG_ERROR(nullptr, "Could not parse deploy file");
}
return parse(network, weightType, modelFile != nullptr);
}
const IBlobNameToTensor* CaffeParser::parse(INetworkDefinition& network,
DataType weightType,
bool hasModel)
{
bool ok = true;
CaffeWeightFactory weights(*mModel.get(), weightType, mTmpAllocs, hasModel);
mBlobNameToTensor = new (BlobNameToTensor);
// Get list of all available plugin creators
int numCreators = 0;
nvinfer1::IPluginCreator* const* tmpList = getPluginRegistry()->getPluginCreatorList(&numCreators);
for (int k = 0; k < numCreators; ++k)
{
if (!tmpList[k])
{
std::cout << "Plugin Creator for plugin " << k << " is a nullptr." << std::endl;
continue;
}
std::string pluginName = tmpList[k]->getPluginName();
mPluginRegistry[pluginName] = tmpList[k];
}
for (int i = 0; i < mDeploy->input_size(); i++)
{
DimsCHW dims;
if (mDeploy->input_shape_size())
{
dims = DimsCHW{(int) mDeploy->input_shape().Get(i).dim().Get(1), (int) mDeploy->input_shape().Get(i).dim().Get(2), (int) mDeploy->input_shape().Get(i).dim().Get(3)};
}
else
{ // deprecated, but still used in a lot of networks
dims = DimsCHW{(int) mDeploy->input_dim().Get(i * 4 + 1), (int) mDeploy->input_dim().Get(i * 4 + 2), (int) mDeploy->input_dim().Get(i * 4 + 3)};
}
ITensor* tensor = network.addInput(mDeploy->input().Get(i).c_str(), DataType::kFLOAT, dims);
(*mBlobNameToTensor)[mDeploy->input().Get(i)] = tensor;
}
for (int i = 0; i < mDeploy->layer_size() && ok; i++)
{
const trtcaffe::LayerParameter& layerMsg = mDeploy->layer(i);
if (layerMsg.has_phase() && layerMsg.phase() == trtcaffe::TEST)
{
continue;
}
// If there is a inplace operation and the operation is
// modifying the input, emit an error as
for (int j = 0; ok && j < layerMsg.top_size(); ++j)
{
for (int k = 0; ok && k < layerMsg.bottom_size(); ++k)
{
if (layerMsg.top().Get(j) == layerMsg.bottom().Get(k))
{
auto iter = mBlobNameToTensor->find(layerMsg.top().Get(j).c_str());
if (iter != nullptr && iter->isNetworkInput())
{
ok = false;
std::cout << "TensorRT does not support in-place operations on input tensors in a prototxt file." << std::endl;
}
}
}
}
// If there is a pluginFactory provided, use layer name matching to handle the plugin construction
if (mPluginFactory && mPluginFactory->isPlugin(layerMsg.name().c_str()))
{
std::vector<Weights> w = weights.getAllWeights(layerMsg.name());
IPlugin* plugin = mPluginFactory->createPlugin(layerMsg.name().c_str(), w.empty() ? nullptr : &w[0], w.size());
std::vector<ITensor*> inputs;
for (int i = 0, n = layerMsg.bottom_size(); i < n; i++)
{
inputs.push_back((*mBlobNameToTensor)[layerMsg.bottom(i)]);
}
bool isExt = mPluginFactoryIsExt && static_cast<IPluginFactoryExt*>(mPluginFactory)->isPluginExt(layerMsg.name().c_str());
ILayer* layer = isExt ? network.addPluginExt(&inputs[0], int(inputs.size()), *static_cast<IPluginExt*>(plugin))
: network.addPlugin(&inputs[0], int(inputs.size()), *plugin);
layer->setName(layerMsg.name().c_str());
if (plugin->getNbOutputs() != layerMsg.top_size())
{
std::cout << "Plugin layer output count is not equal to caffe output count" << std::endl;
ok = false;
}
for (int i = 0, n = std::min(layer->getNbOutputs(), layerMsg.top_size()); i < n; i++)
{
(*mBlobNameToTensor)[layerMsg.top(i)] = layer->getOutput(i);
}
if (layer == nullptr)
{
std::cout << "error parsing layer type " << layerMsg.type() << " index " << i << std::endl;
ok = false;
}
continue;
}
if (getInferLibVersion() >= 5000)
{
if (mPluginFactoryV2 && mPluginFactoryV2->isPluginV2(layerMsg.name().c_str()))
{
if (mPluginFactory)
{
RETURN_AND_LOG_ERROR(nullptr, "Both IPluginFactory and IPluginFactoryV2 are set. If using TensorRT 5.0 or later, switch to IPluginFactoryV2");
}
std::vector<Weights> w = weights.getAllWeights(layerMsg.name());
nvinfer1::IPluginV2* plugin = mPluginFactoryV2->createPlugin(layerMsg.name().c_str(), w.empty() ? nullptr : &w[0], w.size(), mPluginNamespace.c_str());
std::vector<ITensor*> inputs;
for (int i = 0, n = layerMsg.bottom_size(); i < n; i++)
{
inputs.push_back((*mBlobNameToTensor)[layerMsg.bottom(i)]);
}
ILayer* layer = network.addPluginV2(&inputs[0], int(inputs.size()), *plugin);
layer->setName(layerMsg.name().c_str());
if (plugin->getNbOutputs() != layerMsg.top_size())
{
std::cout << "Plugin layer output count is not equal to caffe output count" << std::endl;
ok = false;
}
for (int i = 0, n = std::min(layer->getNbOutputs(), layerMsg.top_size()); i < n; i++)
{
(*mBlobNameToTensor)[layerMsg.top(i)] = layer->getOutput(i);
}
if (layer == nullptr)
{
std::cout << "error parsing layer type " << layerMsg.type() << " index " << i << std::endl;
ok = false;
}
continue;
}
// Use the TRT5 plugin creator method to check for built-in plugin support
std::string pluginName;
nvinfer1::PluginFieldCollection fc;
std::vector<nvinfer1::PluginField> f;
if (layerMsg.type() == "Normalize")
{
pluginName = "Normalize_TRT";
f = parseNormalizeParam(layerMsg, weights, *mBlobNameToTensor);
}
else if (layerMsg.type() == "PriorBox")
{
pluginName = "PriorBox_TRT";
f = parsePriorBoxParam(layerMsg, weights, *mBlobNameToTensor);
}
else if (layerMsg.type() == "DetectionOutput")
{
pluginName = "NMS_TRT";
f = parseDetectionOutputParam(layerMsg, weights, *mBlobNameToTensor);
}
else if (layerMsg.type() == "RPROI")
{
pluginName = "RPROI_TRT";
f = parseRPROIParam(layerMsg, weights, *mBlobNameToTensor);
}
if (mPluginRegistry.find(pluginName) != mPluginRegistry.end())
{
// Set fc
fc.nbFields = f.size();
fc.fields = f.empty() ? nullptr : f.data();
nvinfer1::IPluginV2* pluginV2 = mPluginRegistry.at(pluginName)->createPlugin(layerMsg.name().c_str(), &fc);
assert(pluginV2);
mNewPlugins.push_back(pluginV2);
std::vector<ITensor*> inputs;
for (int i = 0, n = layerMsg.bottom_size(); i < n; i++)
{
inputs.push_back((*mBlobNameToTensor)[layerMsg.bottom(i)]);
}
auto layer = network.addPluginV2(&inputs[0], int(inputs.size()), *pluginV2);
layer->setName(layerMsg.name().c_str());
if (pluginV2->getNbOutputs() != layerMsg.top_size())
{
std::cout << "Plugin layer output count is not equal to caffe output count" << std::endl;
ok = false;
}
for (int i = 0, n = std::min(layer->getNbOutputs(), layerMsg.top_size()); i < n; i++)
{
(*mBlobNameToTensor)[layerMsg.top(i)] = layer->getOutput(i);
}
if (layer == nullptr)
{
std::cout << "error parsing layer type " << layerMsg.type() << " index " << i << std::endl;
ok = false;
}
continue;
}
}
if (layerMsg.type() == "Dropout")
{
(*mBlobNameToTensor)[layerMsg.top().Get(0)] = (*mBlobNameToTensor)[layerMsg.bottom().Get(0)];
continue;
}
if (layerMsg.type() == "Input")
{
const trtcaffe::InputParameter& p = layerMsg.input_param();
for (int i = 0; i < layerMsg.top_size(); i++)
{
const trtcaffe::BlobShape& shape = p.shape().Get(i);
if (shape.dim_size() != 4)
{
RETURN_AND_LOG_ERROR(nullptr, "error parsing input layer, TensorRT only supports 4 dimensional input");
}
else
{
DimsCHW dims{(int) shape.dim().Get(1), (int) shape.dim().Get(2), (int) shape.dim().Get(3)};
ITensor* tensor = network.addInput(layerMsg.top(i).c_str(), DataType::kFLOAT, dims);
(*mBlobNameToTensor)[layerMsg.top().Get(i)] = tensor;
}
}
continue;
}
if (layerMsg.type() == "Flatten")
{
ITensor* tensor = (*mBlobNameToTensor)[layerMsg.bottom().Get(0)];
(*mBlobNameToTensor)[layerMsg.top().Get(0)] = tensor;
std::cout << "Warning: Flatten layer ignored. TensorRT implicitly"
" flattens input to FullyConnected layers, but in other"
" circumstances this will result in undefined behavior."
<< std::endl;
continue;
}
// Use parser table to lookup the corresponding parse function to handle the rest of the layers
auto v = gParseTable.find(layerMsg.type());
if (v == gParseTable.end())
{
std::cout << "could not parse layer type " << layerMsg.type() << std::endl;
ok = false;
}
else
{
ILayer* layer = (*v->second)(network, layerMsg, weights, *static_cast<BlobNameToTensor*>(mBlobNameToTensor));
if (layer == nullptr)
{
std::cout << "error parsing layer type " << layerMsg.type() << " index " << i << std::endl;
ok = false;
}
else
{
layer->setName(layerMsg.name().c_str());
(*mBlobNameToTensor)[layerMsg.top(0)] = layer->getOutput(0);
}
}
}
mBlobNameToTensor->setTensorNames();
return ok && weights.isOK() && mBlobNameToTensor->isOK() ? mBlobNameToTensor : nullptr;
}
IBinaryProtoBlob* CaffeParser::parseBinaryProto(const char* fileName)
{
CHECK_NULL_RET_NULL(fileName)
using namespace google::protobuf::io;
std::ifstream stream(fileName, std::ios::in | std::ios::binary);
if (!stream)
{
RETURN_AND_LOG_ERROR(nullptr, "Could not open file " + std::string{fileName});
}
IstreamInputStream rawInput(&stream);
CodedInputStream codedInput(&rawInput);
codedInput.SetTotalBytesLimit(INT_MAX, -1);
trtcaffe::BlobProto blob;
bool ok = blob.ParseFromCodedStream(&codedInput);
stream.close();
if (!ok)
{
RETURN_AND_LOG_ERROR(nullptr, "parseBinaryProto: Could not parse mean file");
}
DimsNCHW dims{1, 1, 1, 1};
if (blob.has_shape())
{
int size = blob.shape().dim_size(), s[4] = {1, 1, 1, 1};
for (int i = 4 - size; i < 4; i++)
{
assert(blob.shape().dim(i) < INT32_MAX);
s[i] = static_cast<int>(blob.shape().dim(i));
}
dims = DimsNCHW{s[0], s[1], s[2], s[3]};
}
else
{
dims = DimsNCHW{blob.num(), blob.channels(), blob.height(), blob.width()};
}
const int dataSize = dims.n() * dims.c() * dims.h() * dims.w();
assert(dataSize > 0);
const trtcaffe::Type blobProtoDataType = CaffeWeightFactory::getBlobProtoDataType(blob);
const auto blobProtoData = CaffeWeightFactory::getBlobProtoData(blob, blobProtoDataType, mTmpAllocs);
if (dataSize != (int) blobProtoData.second)
{
std::cout << "CaffeParser::parseBinaryProto: blob dimensions don't match data size!!" << std::endl;
return nullptr;
}
const int dataSizeBytes = dataSize * CaffeWeightFactory::sizeOfCaffeType(blobProtoDataType);
void* memory = malloc(dataSizeBytes);
memcpy(memory, blobProtoData.first, dataSizeBytes);
return new BinaryProtoBlob(memory,
blobProtoDataType == trtcaffe::FLOAT ? DataType::kFLOAT : DataType::kHALF, dims);
std::cout << "CaffeParser::parseBinaryProto: couldn't find any data!!" << std::endl;
return nullptr;
}
parsers/caffe/caffeParser/caffeParser.h
+93
View File
@@ -0,0 +1,93 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_CAFFE_PARSER_H
#define TRT_CAFFE_PARSER_CAFFE_PARSER_H
#include <vector>
#include <memory>
#include <unordered_map>
#include <string>
#include "NvCaffeParser.h"
#include "caffeWeightFactory.h"
#include "blobNameToTensor.h"
#include "trtcaffe.pb.h"
namespace nvcaffeparser1
{
class CaffeParser : public ICaffeParser
{
public:
const IBlobNameToTensor* parse(const char* deploy,
const char* model,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightType) override;
const IBlobNameToTensor* parseBuffers(const char* deployBuffer,
size_t deployLength,
const char* modelBuffer,
size_t modelLength,
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightType) override;
void setProtobufBufferSize(size_t size) override { mProtobufBufferSize = size; }
void setPluginFactory(nvcaffeparser1::IPluginFactory* factory) override { mPluginFactory = factory; }
void setPluginFactoryExt(nvcaffeparser1::IPluginFactoryExt* factory) override
{
mPluginFactory = factory;
mPluginFactoryIsExt = true;
}
void setPluginFactoryV2(nvcaffeparser1::IPluginFactoryV2* factory) override { mPluginFactoryV2 = factory; }
void setPluginNamespace(const char* libNamespace) override { mPluginNamespace = libNamespace; }
IBinaryProtoBlob* parseBinaryProto(const char* fileName) override;
void destroy() override { delete this; }
private:
~CaffeParser() override;
std::vector<nvinfer1::PluginField> parseNormalizeParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
std::vector<nvinfer1::PluginField> parsePriorBoxParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
std::vector<nvinfer1::PluginField> parseDetectionOutputParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
std::vector<nvinfer1::PluginField> parseLReLUParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
std::vector<nvinfer1::PluginField> parseRPROIParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
template <typename T>
T* allocMemory(int size = 1)
{
T* tmpMem = static_cast<T*>(malloc(sizeof(T) * size));
mTmpAllocs.push_back(tmpMem);
return tmpMem;
}
const IBlobNameToTensor* parse(nvinfer1::INetworkDefinition& network,
nvinfer1::DataType weightType,
bool hasModel);
private:
std::shared_ptr<trtcaffe::NetParameter> mDeploy;
std::shared_ptr<trtcaffe::NetParameter> mModel;
std::vector<void*> mTmpAllocs;
BlobNameToTensor* mBlobNameToTensor{nullptr};
size_t mProtobufBufferSize{INT_MAX};
nvcaffeparser1::IPluginFactory* mPluginFactory{nullptr};
nvcaffeparser1::IPluginFactoryV2* mPluginFactoryV2{nullptr};
bool mPluginFactoryIsExt{false};
std::vector<nvinfer1::IPluginV2*> mNewPlugins;
std::unordered_map<std::string, nvinfer1::IPluginCreator*> mPluginRegistry;
std::string mPluginNamespace = "";
};
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_CAFFE_PARSER_H
parsers/caffe/caffeParser/opParsers/opParsers.h
+100
View File
@@ -0,0 +1,100 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_OP_PARSERS_H
#define TRT_CAFFE_PARSER_OP_PARSERS_H
#include <unordered_map>
#include <iostream>
#include "caffeMacros.h"
#include "NvInfer.h"
#include "trtcaffe.pb.h"
#include "parserUtils.h"
#include "half.h"
#include "blobNameToTensor.h"
#include "caffeWeightFactory.h"
namespace nvcaffeparser1
{
inline bool checkBlobs(const trtcaffe::LayerParameter& msg, int bottoms, int tops)
{
if (msg.bottom_size() != bottoms)
{
std::cout << msg.name() << ": expected " << bottoms << " bottom blobs, found " << msg.bottom_size() << std::endl;
return false;
}
if (msg.top_size() != tops)
{
std::cout << msg.name() << ": expected " << tops << " tops blobs, found " << msg.top_size() << std::endl;
return false;
}
return true;
}
typedef nvinfer1::ILayer* (*LayerParseFn)(nvinfer1::INetworkDefinition&, const trtcaffe::LayerParameter&, CaffeWeightFactory&, BlobNameToTensor&);
nvinfer1::ILayer* parseAbsVal(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseBatchNormalization(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseBNLL(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseClip(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseConcat(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseConvolution(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseCrop(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseDeconvolution(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseEltwise(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseELU(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseInnerProduct(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseLRN(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parsePermute(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parsePooling(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parsePower(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseReduction(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseReLU(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseReshape(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseScale(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseSigmoid(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseSoftMax(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
nvinfer1::ILayer* parseTanH(nvinfer1::INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors);
static std::unordered_map<std::string, LayerParseFn> gParseTable
{
{"Convolution", parseConvolution},
{"Pooling", parsePooling},
{"InnerProduct", parseInnerProduct},
{"ReLU", parseReLU},
{"Softmax", parseSoftMax},
{"SoftmaxWithLoss", parseSoftMax},
{"LRN", parseLRN},
{"Power", parsePower},
{"Eltwise", parseEltwise},
{"Concat", parseConcat},
{"Deconvolution", parseDeconvolution},
{"Sigmoid", parseSigmoid},
{"TanH", parseTanH},
{"BatchNorm", parseBatchNormalization},
{"Scale", parseScale},
{"Crop", parseCrop},
{"Reduction", parseReduction},
{"Reshape", parseReshape},
{"Permute", parsePermute},
{"ELU", parseELU},
{"BNLL", parseBNLL},
{"Clip", parseClip},
{"AbsVal", parseAbsVal}
};
} // namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_OP_PARSERS_H
parsers/caffe/caffeParser/opParsers/parseAbsVal.cpp
+31
View File
@@ -0,0 +1,31 @@
/*
* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseAbsVal(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
return network.addUnary(*tensors[msg.bottom(0)], UnaryOperation::kABS);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseBNLL.cpp
+31
View File
@@ -0,0 +1,31 @@
/*
* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseBNLL(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
return network.addActivation(*tensors[msg.bottom(0)], ActivationType::kSOFTPLUS);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseBatchNorm.cpp
+177
View File
@@ -0,0 +1,177 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
template <typename T>
inline bool bnConvertWrap(float scaleFactor, const Weights& variance, const Weights& mean,
const Weights& scaleBlob, const Weights& biasBlob,
Weights& shift, Weights& scale, float eps,
bool nvCaffe, CaffeWeightFactory& weightFactory)
{
assert(shift.count == scale.count);
if (nvCaffe)
{
if (scaleBlob.values == nullptr)
{
return false;
}
if (biasBlob.values == nullptr)
{
return false;
}
}
T* shiftv = reinterpret_cast<T*>(malloc(sizeof(T) * shift.count));
if (!shiftv)
{
return false;
}
T* scalev = reinterpret_cast<T*>(malloc(sizeof(T) * scale.count));
if (!scalev)
{
free(shiftv);
return false;
}
shift.values = shiftv;
scale.values = scalev;
weightFactory.getTmpAllocs().push_back(shiftv);
weightFactory.getTmpAllocs().push_back(scalev);
const T* m = reinterpret_cast<const T*>(mean.values);
const T* v = reinterpret_cast<const T*>(variance.values);
for (int i = 0; i < shift.count; i++)
{
scalev[i] = T(1.0f / std::sqrt(float(v[i]) * scaleFactor + eps));
shiftv[i] = T(-(float(m[i]) * scaleFactor * float(scalev[i])));
}
if (nvCaffe)
{
const T* s = reinterpret_cast<const T*>(scaleBlob.values);
const T* b = reinterpret_cast<const T*>(biasBlob.values);
for (int i = 0; i < shift.count; i++)
{
scalev[i] = T(float(scalev[i]) * s[i]);
shiftv[i] = T(float(shiftv[i]) * s[i]) + b[i];
}
}
return true;
}
ILayer* parseBatchNormalization(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::BatchNormParameter& p = msg.batch_norm_param();
bool nvCaffe = weightFactory.getBlobsSize(msg.name()) == 5;
int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
Weights mean{DataType::kFLOAT, nullptr, 0},
variance{DataType::kFLOAT, nullptr, 0},
scaleBlob{DataType::kFLOAT, nullptr, 0},
biasBlob{DataType::kFLOAT, nullptr, 0},
movingAverage{DataType::kFLOAT, nullptr, 0};
// Because of the incompatible nature of the batch normalizations
// between BLVC Caffe and nvCaffe, two different paths have to be
// used.
if (nvCaffe)
{
if (weightFactory.isInitialized())
{
mean = weightFactory(msg.name(), WeightType::kNVMEAN);
variance = weightFactory(msg.name(), WeightType::kNVVARIANCE);
scaleBlob = weightFactory(msg.name(), WeightType::kNVSCALE);
biasBlob = weightFactory(msg.name(), WeightType::kNVBIAS);
}
else
{
mean = weightFactory.allocateWeights(C);
variance = weightFactory.allocateWeights(C, std::uniform_real_distribution<float>(0.9F, 1.1F));
scaleBlob = weightFactory.allocateWeights(C, std::uniform_real_distribution<float>(0.9F, 1.1F));
biasBlob = weightFactory.allocateWeights(C);
}
}
else
{
if (weightFactory.isInitialized())
{
mean = weightFactory(msg.name(), WeightType::kMEAN);
variance = weightFactory(msg.name(), WeightType::kVARIANCE);
movingAverage = weightFactory(msg.name(), WeightType::kMOVING_AVERAGE);
}
else
{
mean = weightFactory.allocateWeights(C);
variance = weightFactory.allocateWeights(C, std::uniform_real_distribution<float>(0.9F, 1.1F));
movingAverage = weightFactory.allocateWeights(1, std::uniform_real_distribution<float>(0.99F, 1.01F));
}
assert(mean.count == variance.count && movingAverage.count == 1);
}
Weights shift{mean.type, nullptr, mean.count};
Weights scale{mean.type, nullptr, mean.count};
Weights power{mean.type, nullptr, 0};
bool success{false};
float scaleFactor{1.0f};
if (!nvCaffe)
{
float average{0.0f};
// Inside weightFactory, the weights are generated based off the type.
if (mean.type == DataType::kFLOAT)
{
average = *(static_cast<const float*>(movingAverage.values));
}
else
{
average = *(static_cast<const float16*>(movingAverage.values));
}
if (average == 0.0f)
{
std::cout << "Batch normalization moving average is zero" << std::endl;
return nullptr;
}
scaleFactor /= average;
}
if (mean.type == DataType::kFLOAT)
{
success = bnConvertWrap<float>(scaleFactor, variance, mean, scaleBlob, biasBlob, shift, scale, p.eps(), nvCaffe, weightFactory);
}
else
{
success = bnConvertWrap<float16>(scaleFactor, variance, mean, scaleBlob, biasBlob, shift, scale, p.eps(), nvCaffe, weightFactory);
}
if (!success)
{
return nullptr;
}
weightFactory.convert(shift);
weightFactory.convert(scale);
weightFactory.convert(power);
return network.addScale(*tensors[msg.bottom(0)], ScaleMode::kCHANNEL, shift, scale, power);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseClip.cpp
+44
View File
@@ -0,0 +1,44 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseClip(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ClipParameter& p = msg.clip_param();
float alpha = std::numeric_limits<float>::lowest(); // lower bound
float beta = std::numeric_limits<float>::max(); // upper bound
if(p.has_min())
{
alpha = p.min();
}
if(p.has_max())
{
beta = p.max();
}
auto layer = network.addActivation(*tensors[msg.bottom(0)], ActivationType::kCLIP);
layer->setAlpha(alpha);
layer->setBeta(beta);
return layer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseConcat.cpp
+50
View File
@@ -0,0 +1,50 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseConcat(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
const trtcaffe::ConcatParameter& p = msg.concat_param();
bool hasAxis = p.has_axis(); // optional parameter
if (hasAxis && p.axis() <= 0)
{
std::cout << "Caffe parser: Concat along batch axis or negative axis is not supported." << std::endl;
return nullptr;
}
std::vector<ITensor*> ptrs;
for (unsigned int i = 0, n = msg.bottom_size(); i < n; i++)
{
ptrs.push_back(tensors[msg.bottom().Get(i)]);
}
auto concat = network.addConcatenation(&ptrs[0], msg.bottom_size());
// If no axis is explicitly provided, do not call setAxis.
// Rely on the default axis setting inside TRT which takes into account NPCHW and higher dimensional input.
if (hasAxis)
{
concat->setAxis(p.axis() - 1);
}
return concat;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseConv.cpp
+67
View File
@@ -0,0 +1,67 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseConvolution(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ConvolutionParameter& p = msg.convolution_param();
int nbOutputs = p.num_output();
int kernelH = p.has_kernel_h() ? p.kernel_h() : p.kernel_size(0);
int kernelW = p.has_kernel_w() ? p.kernel_w() : p.kernel_size_size() > 1 ? p.kernel_size(1) : p.kernel_size(0);
int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
int G = p.has_group() ? p.group() : 1;
auto CbyG = float(C / G * nbOutputs);
float std_dev = 1.0F / sqrtf((kernelW * kernelH * sqrtf(CbyG)));
Weights kernelWeights = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(kernelW * kernelH * CbyG, std::normal_distribution<float>(0.0F, std_dev));
Weights biasWeights = !p.has_bias_term() || p.bias_term() ? (weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kBIAS) : weightFactory.allocateWeights(nbOutputs)) : weightFactory.getNullWeights();
weightFactory.convert(kernelWeights);
weightFactory.convert(biasWeights);
auto layer = network.addConvolution(*tensors[msg.bottom(0)], nbOutputs, DimsHW{kernelH, kernelW}, kernelWeights, biasWeights);
if (layer)
{
int strideH = p.has_stride_h() ? p.stride_h() : p.stride_size() > 0 ? p.stride(0) : 1;
int strideW = p.has_stride_w() ? p.stride_w() : p.stride_size() > 1 ? p.stride(1) : p.stride_size() > 0 ? p.stride(0) : 1;
int padH = p.has_pad_h() ? p.pad_h() : p.pad_size() > 0 ? p.pad(0) : 0;
int padW = p.has_pad_w() ? p.pad_w() : p.pad_size() > 1 ? p.pad(1) : p.pad_size() > 0 ? p.pad(0) : 0;
int dilationH = p.dilation_size() > 0 ? p.dilation(0) : 1;
int dilationW = p.dilation_size() > 1 ? p.dilation(1) : p.dilation_size() > 0 ? p.dilation(0) : 1;
layer->setStride(DimsHW{strideH, strideW});
layer->setPadding(DimsHW{padH, padW});
layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_DOWN);
layer->setDilation(DimsHW{dilationH, dilationW});
layer->setNbGroups(G);
}
return layer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseCrop.cpp
+117
View File
@@ -0,0 +1,117 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseCrop(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
// To crop, elements of the first bottom are selected to fit the dimensions
// of the second, reference bottom. The crop is configured by
// - the crop `axis` to pick the dimensions for cropping
// - the crop `offset` to set the shift for all/each dimension
// to align the cropped bottom with the reference bottom.
// All dimensions up to but excluding `axis` are preserved, while
// the dimensions including and trailing `axis` are cropped.
// If only one `offset` is set, then all dimensions are offset by this amount.
// Otherwise, the number of offsets must equal the number of cropped axes to
// shift the crop in each dimension accordingly.
// Note: standard dimensions are N,C,H,W so the default is a spatial crop,
// and `axis` may be negative to index from the end (e.g., -1 for the last
// axis).
if (!checkBlobs(msg, 2, 1))
{
return nullptr;
}
// ONLY IMPLEMENT SPATIAL CROPPING
// IF CROP LAYER IS NOT SPATIAL CROP, ABORT
const trtcaffe::CropParameter& p = msg.crop_param();
DimsCHW inputDims = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions());
DimsCHW refDims = parserutils::getCHW(tensors[msg.bottom(1)]->getDimensions());
bool hasAxis = p.has_axis(); // optional parameter
int axis = hasAxis ? p.axis() : 2; // default is 2 - spatial crop
axis = (axis < 0) ? 4 + axis : axis; // axis negative number correction
// acceptable axis values: 2, 3, -1, -2
// unacceptable axis values: 0, 1, -3, -4 and anything else
// acceptable corrected axis values: 2, 3
// unacceptable corrected axis values: 0, 1 and anything else
// protect against "garbage" input arguments
bool axis_abort = (axis != 2 && axis != 3);
// must be at least one offset
// if only one offset, the same offset applies to all the dimensions
// including the chosen axis and trailing it
// if more than one offset, the number of offsets must match the number
// of dimensions consisting of the axis and all the dimensions trailing it
int num_offsets = p.offset_size();
// 1 + (3 - axis) = 4 - axis
// this is only valid for acceptable corrected axis values
// if !axis_abort then invariant that num_dims == 1 || num_dims == 2
int num_dims = 4 - axis;
bool offset_abort = (num_offsets != 0 && num_offsets != 1 && num_offsets != num_dims);
if (axis_abort)
{
std::cout << "Caffe Parser: Invalid axis in crop layer - only spatial cropping is supported" << std::endl;
return nullptr;
}
if (offset_abort)
{
std::cout << "Caffe Parser: Invalid number of offsets in crop layer" << std::endl;
return nullptr;
}
// get the offsets
// the offsets are zero by default (in case no offset is specified)
int offsetHeight = 0;
int offsetWidth = 0;
if (num_offsets != 0)
{
// offsetHeight will only be specified if the H channel is the chosen axis
// in this case, regardless of whether there are one or multiple offsets
// offsetHeight should always be the zero-indexed offset
offsetHeight = axis == 2 ? p.offset(0) : 0;
// offsetWidth should always be specified
// if there is only one offset, use the zero-indexed offset
// otherwise, use the one-indexed offset since the zero-indexed offet
// is for offsetHeight
offsetWidth = num_offsets == 1 ? p.offset(0) : p.offset(1);
}
// now compute the prePadding and postPadding required to perform the crop
// so that the first bottom is the same spatial size as the second bottom
// prePadding is the padding to the left/bottom (assuming origin is lower-left).
// postPadding is the padding to the right/top.
// - ( inputDims.h() - refDims.h() - offsetHeight ) = -inputDims.h() + refDims.h() + offsetHeight
// - ( inputDims.w() - refDims.w() - offsetWidth ) = -inputDims.w() + refDims.w() + offsetWidth
int prePadHeight = -offsetHeight;
int prePadWidth = -offsetWidth;
int postPadHeight = -inputDims.h() + refDims.h() + offsetHeight;
int postPadWidth = -inputDims.w() + refDims.w() + offsetWidth;
DimsHW prePadding = DimsHW{prePadHeight, prePadWidth};
DimsHW postPadding = DimsHW{postPadHeight, postPadWidth};
return network.addPadding(*tensors[msg.bottom(0)], prePadding, postPadding);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseDeconv.cpp
+74
View File
@@ -0,0 +1,74 @@
/*
* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseDeconvolution(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ConvolutionParameter& p = msg.convolution_param();
int nbOutputs = p.num_output();
int nbGroups = p.has_group() ? p.group() : 1;
int dilationH = p.dilation_size() > 0 ? p.dilation(0) : 1;
int dilationW = p.dilation_size() > 1 ? p.dilation(1) : p.dilation_size() > 0 ? p.dilation(0) : 1;
if (dilationH != 1 || dilationW != 1)
{
RETURN_AND_LOG_ERROR(nullptr, "Dilated deconvolution is not supported.");
}
int kernelW = p.has_kernel_w() ? p.kernel_w() : p.kernel_size(0);
int kernelH = p.has_kernel_h() ? p.kernel_h() : p.kernel_size_size() > 1 ? p.kernel_size(1) : p.kernel_size(0);
int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
float std_dev = 1.0F / sqrtf(kernelW * kernelH * sqrtf(C * nbOutputs));
Weights kernelWeights = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(kernelW * kernelH * C * nbOutputs / nbGroups, std::normal_distribution<float>(0.0F, std_dev));
Weights biasWeights = !p.has_bias_term() || p.bias_term() ? (weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kBIAS) : weightFactory.allocateWeights(nbOutputs)) : weightFactory.getNullWeights();
weightFactory.convert(kernelWeights);
weightFactory.convert(biasWeights);
auto layer = network.addDeconvolution(*tensors[msg.bottom(0)], nbOutputs, DimsHW{kernelH, kernelW}, kernelWeights, biasWeights);
if (layer)
{
int strideW = p.has_stride_w() ? p.stride_w() : p.stride_size() > 0 ? p.stride(0) : 1;
int strideH = p.has_stride_h() ? p.stride_h() : p.stride_size() > 1 ? p.stride(1) : p.stride_size() > 0 ? p.stride(0) : 1;
int padW = p.has_pad_w() ? p.pad_w() : p.pad_size() > 0 ? p.pad(0) : 0;
int padH = p.has_pad_h() ? p.pad_h() : p.pad_size() > 1 ? p.pad(1) : p.pad_size() > 0 ? p.pad(0) : 0;
layer->setStride(DimsHW{strideH, strideW});
layer->setPadding(DimsHW{padH, padW});
layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_DOWN);
layer->setNbGroups(nbGroups);
layer->setKernelWeights(kernelWeights);
if (!p.has_bias_term() || p.bias_term())
{
layer->setBiasWeights(biasWeights);
}
}
return layer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseELU.cpp
+40
View File
@@ -0,0 +1,40 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseELU(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /* weightFactory */, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ELUParameter& p = msg.elu_param();
float alpha = 1.f; // default parameter
if (p.has_alpha())
{
alpha = p.alpha();
}
auto newLayer = network.addActivation(*tensors[msg.bottom(0)], ActivationType::kELU);
newLayer->setAlpha(alpha);
return newLayer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseEltwise.cpp
+41
View File
@@ -0,0 +1,41 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseEltwise(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 2, 1))
{
return nullptr;
}
const trtcaffe::EltwiseParameter& p = msg.eltwise_param();
ElementWiseOperation op = ElementWiseOperation::kSUM;
switch (p.operation())
{
case trtcaffe::EltwiseParameter_EltwiseOp_SUM: op = ElementWiseOperation::kSUM; break;
case trtcaffe::EltwiseParameter_EltwiseOp_PROD: op = ElementWiseOperation::kPROD; break;
case trtcaffe::EltwiseParameter_EltwiseOp_MAX: op = ElementWiseOperation::kMAX; break;
}
return network.addElementWise(*tensors[msg.bottom(0)], *tensors[msg.bottom(1)], op);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseInnerProduct.cpp
+37
View File
@@ -0,0 +1,37 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseInnerProduct(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
const trtcaffe::InnerProductParameter& p = msg.inner_product_param();
int64_t nbInputs = parserutils::volume(parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()));
int64_t nbOutputs = p.num_output();
float std_dev = 1.0F / sqrtf(nbInputs * nbOutputs);
Weights kernelWeights = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(nbInputs * nbOutputs, std::normal_distribution<float>(0.0F, std_dev));
Weights biasWeights = !p.has_bias_term() || p.bias_term() ? (weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kBIAS) : weightFactory.allocateWeights(nbOutputs)) : weightFactory.getNullWeights();
weightFactory.convert(kernelWeights);
weightFactory.convert(biasWeights);
return network.addFullyConnected(*tensors[msg.bottom(0)], p.num_output(), kernelWeights, biasWeights);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseLRN.cpp
+37
View File
@@ -0,0 +1,37 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseLRN(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::LRNParameter& p = msg.lrn_param();
int localSize = p.has_local_size() ? p.local_size() : 5;
float alpha = p.has_alpha() ? p.alpha() : 1;
float beta = p.has_beta() ? p.beta() : 5;
float k = p.has_k() ? p.k() : 1;
return network.addLRN(*tensors[msg.bottom(0)], localSize, alpha, beta, k);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parsePermute.cpp
+80
View File
@@ -0,0 +1,80 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parsePermute(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::PermuteParameter& p = msg.permute_param();
Dims bottomDims = tensors[msg.bottom(0)]->getDimensions();
Dims topDims = tensors[msg.bottom(0)]->getDimensions();
int nbDims = bottomDims.nbDims;
std::vector<int> orders;
std::vector<bool> knownOrders(nbDims + 1, false);
bool orderAbort = (p.order(0) != 0); // First order must be 0 (batch dimension)
for (int i = 0; i < p.order_size(); i++)
{
int order = p.order(i);
orderAbort |= (order > nbDims) || (std::find(orders.begin(), orders.end(), order) != orders.end());
orders.push_back(order);
knownOrders[order] = true;
}
if (orderAbort)
{
std::cout << "Caffe Parser: Invalid permute param. TensorRT does not support permute in N (batch) dimension, and order index must be within the tensor dimensions. no duplicate order allowed." << std::endl;
return nullptr;
}
// Keep the rest of the order
for (int i = 0; i < nbDims; i++)
{
if (!knownOrders[i])
{
orders.push_back(i);
}
}
// Remove the first order (batch)
orders.erase(orders.begin());
for (int i = 0; i < nbDims; i++)
{
topDims.d[i] = bottomDims.d[orders[i] - 1];
}
assert(parserutils::volume(topDims) == parserutils::volume(bottomDims));
nvinfer1::Permutation permuteOrder;
for (int i = 0; i < nbDims; i++)
{
permuteOrder.order[i] = orders[i] - 1;
}
auto permute = network.addShuffle(*tensors[msg.bottom(0)]);
permute->setReshapeDimensions(topDims);
permute->setFirstTranspose(permuteOrder);
return permute;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parsePooling.cpp
+77
View File
@@ -0,0 +1,77 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parsePooling(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::PoolingParameter& p = msg.pooling_param();
if (p.pool() != trtcaffe::PoolingParameter::MAX && p.pool() != trtcaffe::PoolingParameter::AVE)
{
std::cout << "Caffe Parser: only AVE and MAX pool operations are supported" << std::endl;
return nullptr;
}
int kernelH, kernelW;
if (p.has_global_pooling() && p.global_pooling())
{
DimsCHW dims = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions());
kernelH = dims.h();
kernelW = dims.w();
}
else
{
// mandatory
kernelH = p.has_kernel_h() ? p.kernel_h() : p.kernel_size();
kernelW = p.has_kernel_w() ? p.kernel_w() : p.kernel_size();
}
PoolingType type = p.has_pool() && p.pool() == trtcaffe::PoolingParameter::AVE ? PoolingType::kAVERAGE : PoolingType::kMAX;
auto layer = network.addPooling(*tensors[msg.bottom(0)], type, DimsHW{kernelH, kernelW});
if (layer)
{
int stride = p.has_stride() ? p.stride() : 1;
layer->setStride(DimsHW{p.has_stride_h() ? int(p.stride_h()) : stride, p.has_stride_w() ? int(p.stride_w()) : stride});
int pad = p.has_pad() ? p.pad() : 0;
layer->setPadding(DimsHW{p.has_pad_h() ? int(p.pad_h()) : pad, p.has_pad_w() ? int(p.pad_w()) : pad});
layer->setName(msg.name().c_str());
layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_UP); // caffe pool use ceil mode by default
// FB pooling parameters
// Use floor((height + 2 * padding - kernel) / stride) + 1
// instead of ceil((height + 2 * padding - kernel) / stride) + 1
if (p.has_torch_pooling() ? p.torch_pooling() : false)
{
layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_DOWN); // facebook torch pool use floor mode
}
tensors[msg.top(0)] = layer->getOutput(0);
layer->setAverageCountExcludesPadding(false); // unlike other frameworks, caffe use inclusive counting for padded averaging
}
return layer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parsePower.cpp
+63
View File
@@ -0,0 +1,63 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parsePower(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::PowerParameter& p = msg.power_param();
float shift = p.has_shift() ? p.shift() : 0.0f;
float scale = p.has_scale() ? p.scale() : 1.0f;
float power = p.has_power() ? p.power() : 1.0f;
DataType dataType = weightFactory.getDataType();
assert(dataType == DataType::kFLOAT || dataType == DataType::kHALF);
Weights wShift, wScale, wPower;
if (dataType == DataType::kHALF)
{
auto* t = reinterpret_cast<float16*>(malloc(3 * sizeof(float16)));
t[0] = float16(shift), t[1] = float16(scale), t[2] = float16(power);
wShift = Weights{DataType::kHALF, &t[0], 1};
wScale = Weights{DataType::kHALF, &t[1], 1};
wPower = Weights{DataType::kHALF, &t[2], 1};
weightFactory.getTmpAllocs().push_back(t);
}
else
{
auto* t = reinterpret_cast<float*>(malloc(3 * sizeof(float)));
t[0] = shift, t[1] = scale, t[2] = power;
wShift = Weights{DataType::kFLOAT, &t[0], 1};
wScale = Weights{DataType::kFLOAT, &t[1], 1};
wPower = Weights{DataType::kFLOAT, &t[2], 1};
weightFactory.getTmpAllocs().push_back(t);
}
weightFactory.convert(wShift);
weightFactory.convert(wScale);
weightFactory.convert(wPower);
return network.addScale(*tensors[msg.bottom(0)], ScaleMode::kUNIFORM, wShift, wScale, wPower);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseReLU.cpp
+39
View File
@@ -0,0 +1,39 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseReLU(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ReLUParameter& p = msg.relu_param();
if (p.has_negative_slope() && p.negative_slope() != 0)
{
auto newLayer = network.addActivation(*tensors[msg.bottom(0)], ActivationType::kLEAKY_RELU);
newLayer->setAlpha(p.negative_slope());
return newLayer;
}
return network.addActivation(*tensors[msg.bottom(0)], ActivationType::kRELU);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseReduction.cpp
+143
View File
@@ -0,0 +1,143 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseReduction(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
// The first axis to reduce to a scalar -- may be negative to index from the
// end (e.g., -1 for the last axis).
// (Currently, only reduction along ALL "tail" axes is supported; reduction
// of axis M through N, where N < num_axes - 1, is unsupported.)
// Suppose we have an n-axis bottom Blob with shape:
// (d0, d1, d2, ..., d(m-1), dm, d(m+1), ..., d(n-1)).
// If axis == m, the output Blob will have shape
// (d0, d1, d2, ..., d(m-1)),
// and the ReductionOp operation is performed (d0 * d1 * d2 * ... * d(m-1))
// times, each including (dm * d(m+1) * ... * d(n-1)) individual data.
// If axis == 0 (the default), the output Blob always has the empty shape
// (count 1), performing reduction across the entire input --
// often useful for creating new loss functions.
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
// operation == 1 is SUM -> ReduceOperation::kSUM
const int SUM = 1;
// operation == 2 is ASUM -> UnaryOperation::kABS and ReduceOperation::kSUM
const int ASUM = 2;
// operation == 3 is SUMSQ -> ElementWiseOperation::kPROD and ReduceOperation::kSUM
const int SUMSQ = 3;
// operation == 4 is MEAN -> ReduceOperation::kAVG
const int MEAN = 4;
const trtcaffe::ReductionParameter& p = msg.reduction_param();
bool hasOperation = p.has_operation(); // optional parameter
bool hasAxis = p.has_axis(); // optional parameter
bool hasCoeff = p.has_coeff(); // optional parameter
int operation = hasOperation ? p.operation() : SUM; // default is SUM
int axis = hasAxis ? p.axis() : 0; // default is 0
axis = (axis < 0) ? 4 + axis : axis; // axis negative number correction
float coeff = hasCoeff ? p.coeff() : 1.0; // default is 1
// acceptable axis values: 1, 2, 3, -1, -2, -3
// unacceptable axis values: 0 and anything else
// acceptable corrected axis values: 1, 2, 3
// unacceptable corrected axis values: 0 and anything else
// protect against "garbage" input arguments
bool axisAbort = (axis != 1 && axis != 2 && axis != 3);
if (axisAbort)
{
std::cout << "Caffe Parser: Invalid axis in reduction layer - cannot reduce over batch size dimension and can only reduce NCHW input" << std::endl;
return nullptr;
}
ReduceOperation op = (operation == MEAN ? ReduceOperation::kAVG : ReduceOperation::kSUM);
// corrected axis values are 1, 2, 3
// only reduction along tail dimensions is supported
// 1 means 111 or 4 + 2 + 1 = 7
// 2 means 110 or 4 + 2 = 6
// 3 means 100 or 4
// Let's employ a bit shift trick instead
// 1000 = 8
// axis == 1: 1u << (axis - 1) is 1 and so 8 - 1 = 7 or 111
// axis == 2: 1u << (axis - 1) is 2 and so 8 - 2 = 6 or 110
// axis == 3: 1u << (axis - 1) is 4 and so 8 - 4 = 4 or 100
uint32_t reduceAxes = 8 - (1u << (axis - 1));
ITensor* input = tensors[msg.bottom(0)];
ILayer* returnVal = nullptr;
// need to add in layer before for ASUM and SUMSQ
if (operation == ASUM)
{
returnVal = network.addUnary(*input, UnaryOperation::kABS);
input = returnVal->getOutput(0);
std::string layerName = msg.name() + std::string("/reductionLayer/unaryLayer");
returnVal->setName(layerName.c_str());
}
else if (operation == SUMSQ)
{
returnVal = network.addElementWise(*input, *input, ElementWiseOperation::kPROD);
input = returnVal->getOutput(0);
std::string layerName = msg.name() + std::string("/reductionLayer/elementWiseLayer");
returnVal->setName(layerName.c_str());
}
// add in the actual reduce layer
#define GIE_3111 0
#if GIE_3111
returnVal = network.addReduce(*input, op, reduceAxes, false);
#else
returnVal = network.addReduce(*input, op, reduceAxes, true);
// output a warning
std::cout << "Warning: The Reduce layer does not discard reduced dimensions. The reduced dimensions are treated as dimensions of size one in the output of the Reduce layer." << std::endl;
#endif
input = returnVal->getOutput(0);
std::string reduceLayerName = msg.name() + std::string("/reductionLayer/reduceLayer");
returnVal->setName(reduceLayerName.c_str());
// need to add in layer after for coeff != 1.0
if (coeff != 1.0f)
{
auto* shiftArr = (float*) malloc(sizeof(float));
auto* scaleArr = (float*) malloc(sizeof(float));
auto* powerArr = (float*) malloc(sizeof(float));
weightFactory.getTmpAllocs().push_back(shiftArr);
weightFactory.getTmpAllocs().push_back(scaleArr);
weightFactory.getTmpAllocs().push_back(powerArr);
*shiftArr = 0.0f;
*scaleArr = coeff;
*powerArr = 1.0f;
Weights wShift, wScale, wPower;
wShift = Weights{DataType::kFLOAT, shiftArr, 1};
wScale = Weights{DataType::kFLOAT, scaleArr, 1};
wPower = Weights{DataType::kFLOAT, powerArr, 1};
returnVal = network.addScale(*input, ScaleMode::kUNIFORM, wShift, wScale, wPower);
std::string layerName = msg.name() + std::string("/reductionLayer/scaleLayer");
returnVal->setName(layerName.c_str());
}
return returnVal;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseReshape.cpp
+106
View File
@@ -0,0 +1,106 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseReshape(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ReshapeParameter& p = msg.reshape_param();
Dims bottomDims = tensors[msg.bottom(0)]->getDimensions();
int axis = p.has_axis() ? p.axis() : 0;
const ::trtcaffe::BlobShape& shape = p.shape();
// Check that N (batch dim) is 0. TensorRT does not support reshape in batch dimension
if ((axis == 0) && (shape.dim(0) != 0))
{
std::cout << "Caffe Parser: Invalid reshape param. TensorRT does not support reshape in N (batch) dimension" << std::endl;
return nullptr;
}
// Handle axis and dims parameters
int axStart = std::max(0, axis - 1);
int axEnd = p.has_num_axes() ? std::max(0, axis - 1 + p.num_axes()) : bottomDims.nbDims;
std::vector<int> reshapeDims;
reshapeDims.reserve(axStart);
for (int i = 0; i < axStart; i++)
{
reshapeDims.push_back(bottomDims.d[i]);
}
for (int i = 0; i < shape.dim_size(); i++)
{
// skip first 0 (batch)
if (axis == 0 && i == 0)
{
continue;
}
if (shape.dim(i) == 0)
{
// If there is no bottom dimension corresponding to the current axis, then the params are invalid
assert(static_cast<int>(reshapeDims.size()) <= bottomDims.nbDims);
reshapeDims.push_back(bottomDims.d[reshapeDims.size()]);
}
else
{
reshapeDims.push_back(shape.dim(i));
}
}
for (int i = axEnd; i < bottomDims.nbDims; i++)
{
reshapeDims.push_back(bottomDims.d[i]);
}
Dims topDims{};
topDims.nbDims = static_cast<int>(reshapeDims.size());
for (int i = 0; i < topDims.nbDims; i++)
{
topDims.d[i] = reshapeDims[i];
}
// Check there is at most one -1, and handle such case
int countMinusOne = 0;
for (int i = 0; i < topDims.nbDims; i++)
{
if (topDims.d[i] == -1)
{
countMinusOne += 1;
// Inferred dimension
int64_t newDim = parserutils::volume(bottomDims) / -parserutils::volume(topDims);
topDims.d[i] = newDim;
}
}
if (countMinusOne > 1)
{
std::cout << "Caffe Parser: Invalid reshape param. At most one axis can be inferred from other dimensions" << std::endl;
return nullptr;
}
auto layer = network.addShuffle(*tensors[msg.bottom(0)]);
layer->setReshapeDimensions(topDims);
return layer;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseScale.cpp
+41
View File
@@ -0,0 +1,41 @@
/*
* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseScale(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::ScaleParameter& p = msg.scale_param();
int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
Weights scale = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(C, std::uniform_real_distribution<float>(0.9F, 1.1F));
Weights shift = !p.has_bias_term() || p.bias_term() ? (weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kBIAS) : weightFactory.allocateWeights(C)) : weightFactory.getNullWeights();
Weights power = weightFactory.getNullWeights();
weightFactory.convert(shift);
weightFactory.convert(scale);
weightFactory.convert(power);
return network.addScale(*tensors[msg.bottom(0)], ScaleMode::kCHANNEL, shift, scale, power);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseSigmoid.cpp
+31
View File
@@ -0,0 +1,31 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseSigmoid(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
return network.addActivation(*tensors[msg.bottom(0)], ActivationType::kSIGMOID);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseSoftMax.cpp
+62
View File
@@ -0,0 +1,62 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseSoftMax(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
const trtcaffe::SoftmaxParameter& p = msg.softmax_param();
// Caffe supports negative axis, indexing from the last dimension
// However, there is a discrepancy in the internal tensor dimension in some cases.
// For example. InnerProduct produces flat 1D blob in Caffe, while TensorRT still
// produces CHW format. MNIST sample generates input to Softmax as,
// Caffe = n x 10
// TensorRT = n x 10 x 1 x 1
// To make sure we do not run into issues, negative axis won't be supported in TensorRT
int nbDims = tensors[msg.bottom(0)]->getDimensions().nbDims;
bool hasAxis = p.has_axis(); // optional parameter
int axis = hasAxis ? p.axis() : 1; // default is 1
bool axisAbort = (axis <= 0) || (axis > 3) || (axis > nbDims);
if (axisAbort)
{
std::cout << "Caffe Parser: Invalid axis in softmax layer - Cannot perform softmax along batch size dimension and expects NCHW input. Negative axis is not supported in TensorRT, please use positive axis indexing" << std::endl;
return nullptr;
}
auto softmax = network.addSoftMax(*tensors[msg.bottom(0)]);
// Do this so that setAxes is not used when the default axis is needed
// This is necessary to preserve correct roll-into-the-batch dimension behaviour for samples like FasterRCNN
// NCHW -> default axis when setAxes is not called will be 1 (the C dimension)
// NPCHW -> default axis when setAxes is not called will be 2 (the C dimension)
if (hasAxis)
{
uint32_t axes = 1u << (axis - 1);
softmax->setAxes(axes);
}
return softmax;
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/opParsers/parseTanH.cpp
+31
View File
@@ -0,0 +1,31 @@
/* Copyright (c) 2019, 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.
*/
#include "opParsers.h"
using namespace nvinfer1;
namespace nvcaffeparser1
{
ILayer* parseTanH(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
if (!checkBlobs(msg, 1, 1))
{
return nullptr;
}
return network.addActivation(*tensors[msg.bottom(0)], ActivationType::kTANH);
}
} //namespace nvcaffeparser1
parsers/caffe/caffeParser/readProto.h
+82
View File
@@ -0,0 +1,82 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_READ_PROTO_H
#define TRT_CAFFE_PARSER_READ_PROTO_H
#include <fstream>
#include "google/protobuf/io/coded_stream.h"
#include "google/protobuf/io/zero_copy_stream_impl.h"
#include "google/protobuf/text_format.h"
#include "caffeMacros.h"
#include "trtcaffe.pb.h"
namespace nvcaffeparser1
{
// There are some challenges associated with importing caffe models. One is that
// a .caffemodel file just consists of layers and doesn't have the specs for its
// input and output blobs.
//
// So we need to read the deploy file to get the input
bool readBinaryProto(trtcaffe::NetParameter* net, const char* file, size_t bufSize)
{
CHECK_NULL_RET_VAL(net, false)
CHECK_NULL_RET_VAL(file, false)
using namespace google::protobuf::io;
std::ifstream stream(file, std::ios::in | std::ios::binary);
if (!stream)
{
RETURN_AND_LOG_ERROR(false, "Could not open file " + std::string(file));
}
IstreamInputStream rawInput(&stream);
CodedInputStream codedInput(&rawInput);
codedInput.SetTotalBytesLimit(int(bufSize), -1);
bool ok = net->ParseFromCodedStream(&codedInput);
stream.close();
if (!ok)
{
RETURN_AND_LOG_ERROR(false, "Could not parse binary model file");
}
return ok;
}
bool readTextProto(trtcaffe::NetParameter* net, const char* file)
{
CHECK_NULL_RET_VAL(net, false)
CHECK_NULL_RET_VAL(file, false)
using namespace google::protobuf::io;
std::ifstream stream(file, std::ios::in);
if (!stream)
{
RETURN_AND_LOG_ERROR(false, "Could not open file " + std::string(file));
}
IstreamInputStream input(&stream);
bool ok = google::protobuf::TextFormat::Parse(&input, net);
stream.close();
return ok;
}
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_READ_PROTO_H
parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp
+411
View File
@@ -0,0 +1,411 @@
/*
* Copyright (c) 2019, 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.
*/
#include "caffeMacros.h"
#include "caffeWeightFactory.h"
#include "half.h"
using namespace nvinfer1;
using namespace nvcaffeparser1;
template <typename INPUT, typename OUTPUT>
void* convertInternal(void** ptr, int64_t count, bool* mOK)
{
assert(ptr != nullptr);
if (*ptr == nullptr)
{
return nullptr;
}
if (!count)
{
return nullptr;
}
auto* iPtr = static_cast<INPUT*>(*ptr);
auto* oPtr = static_cast<OUTPUT*>(malloc(count * sizeof(OUTPUT)));
for (int i = 0; i < count; ++i)
{
if (static_cast<OUTPUT>(iPtr[i]) > std::numeric_limits<OUTPUT>::max()
|| static_cast<OUTPUT>(iPtr[i]) < std::numeric_limits<OUTPUT>::lowest())
{
std::cout << "Error: Weight " << iPtr[i] << " is outside of [" << std::numeric_limits<OUTPUT>::max()
<< ", " << std::numeric_limits<OUTPUT>::lowest() << "]." << std::endl;
if (mOK)
{
(*mOK) = false;
}
break;
}
oPtr[i] = iPtr[i];
}
(*ptr) = oPtr;
return oPtr;
}
CaffeWeightFactory::CaffeWeightFactory(const trtcaffe::NetParameter& msg, DataType dataType, std::vector<void*>& tmpAllocs, bool isInitialized)
: mMsg(msg)
, mTmpAllocs(tmpAllocs)
, mDataType(dataType)
, mInitialized(isInitialized)
{
mRef = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
}
DataType CaffeWeightFactory::getDataType() const
{
return mDataType;
}
size_t CaffeWeightFactory::getDataTypeSize() const
{
switch (getDataType())
{
case DataType::kFLOAT:
case DataType::kINT32:
return 4;
case DataType::kHALF:
return 2;
case DataType::kINT8:
return 1;
}
return 0;
}
std::vector<void*>& CaffeWeightFactory::getTmpAllocs()
{
return mTmpAllocs;
}
int CaffeWeightFactory::getBlobsSize(const std::string& layerName)
{
for (int i = 0, n = mMsg.layer_size(); i < n; ++i)
{
if (mMsg.layer(i).name() == layerName)
{
return mMsg.layer(i).blobs_size();
}
}
return 0;
}
const trtcaffe::BlobProto* CaffeWeightFactory::getBlob(const std::string& layerName, int index)
{
if (mMsg.layer_size() > 0)
{
for (int i = 0, n = mMsg.layer_size(); i < n; i++)
{
if (mMsg.layer(i).name() == layerName && index < mMsg.layer(i).blobs_size())
{
return &mMsg.layer(i).blobs(index);
}
}
}
else
{
for (int i = 0, n = mMsg.layers_size(); i < n; i++)
{
if (mMsg.layers(i).name() == layerName && index < mMsg.layers(i).blobs_size())
{
return &mMsg.layers(i).blobs(index);
}
}
}
return nullptr;
}
std::vector<Weights> CaffeWeightFactory::getAllWeights(const std::string& layerName)
{
std::vector<Weights> v;
for (int i = 0;; i++)
{
auto b = getBlob(layerName, i);
if (b == nullptr)
{
break;
}
auto weights = getWeights(*b, layerName);
convert(weights, DataType::kFLOAT);
v.push_back(weights);
}
return v;
}
Weights CaffeWeightFactory::operator()(const std::string& layerName, WeightType weightType)
{
const trtcaffe::BlobProto* blobMsg = getBlob(layerName, int(weightType));
if (blobMsg == nullptr)
{
std::cout << "Weights for layer " << layerName << " doesn't exist" << std::endl;
RETURN_AND_LOG_ERROR(getNullWeights(), "ERROR: Attempting to access NULL weights");
assert(0);
}
return getWeights(*blobMsg, layerName);
}
void CaffeWeightFactory::convert(Weights& weights, DataType targetType)
{
void* tmpAlloc{nullptr};
if (weights.type == DataType::kFLOAT && targetType == DataType::kHALF)
{
tmpAlloc = convertInternal<float, float16>(const_cast<void**>(&weights.values), weights.count, &mOK);
weights.type = targetType;
}
if (weights.type == DataType::kHALF && targetType == DataType::kFLOAT)
{
tmpAlloc = convertInternal<float16, float>(const_cast<void**>(&weights.values), weights.count, &mOK);
weights.type = targetType;
}
if (tmpAlloc)
{
mTmpAllocs.push_back(tmpAlloc);
}
}
void CaffeWeightFactory::convert(Weights& weights)
{
convert(weights, getDataType());
}
bool CaffeWeightFactory::isOK()
{
return mOK;
}
bool CaffeWeightFactory::isInitialized()
{
return mInitialized;
}
Weights CaffeWeightFactory::getNullWeights()
{
return Weights{mDataType, nullptr, 0};
}
Weights CaffeWeightFactory::allocateWeights(int64_t elems, std::uniform_real_distribution<float> distribution)
{
void* data = malloc(elems * getDataTypeSize());
switch (getDataType())
{
case DataType::kFLOAT:
for (int64_t i = 0; i < elems; ++i)
{
((float*) data)[i] = distribution(generator);
}
break;
case DataType::kHALF:
for (int64_t i = 0; i < elems; ++i)
{
((float16*) data)[i] = (float16)(distribution(generator));
}
break;
default:
break;
}
mTmpAllocs.push_back(data);
return Weights{getDataType(), data, elems};
}
Weights CaffeWeightFactory::allocateWeights(int64_t elems, std::normal_distribution<float> distribution)
{
void* data = malloc(elems * getDataTypeSize());
switch (getDataType())
{
case DataType::kFLOAT:
for (int64_t i = 0; i < elems; ++i)
{
((float*) data)[i] = distribution(generator);
}
break;
case DataType::kHALF:
for (int64_t i = 0; i < elems; ++i)
{
((float16*) data)[i] = (float16)(distribution(generator));
}
break;
default:
break;
}
mTmpAllocs.push_back(data);
return Weights{getDataType(), data, elems};
}
trtcaffe::Type CaffeWeightFactory::getBlobProtoDataType(const trtcaffe::BlobProto& blobMsg)
{
if (blobMsg.has_raw_data())
{
assert(blobMsg.has_raw_data_type());
return blobMsg.raw_data_type();
}
if (blobMsg.double_data_size() > 0)
{
return trtcaffe::DOUBLE;
}
return trtcaffe::FLOAT;
}
size_t CaffeWeightFactory::sizeOfCaffeType(trtcaffe::Type type)
{
if (type == trtcaffe::FLOAT)
{
return sizeof(float);
}
if (type == trtcaffe::FLOAT16)
{
return sizeof(uint16_t);
}
return sizeof(double);
}
// The size returned here is the number of array entries, not bytes
std::pair<const void*, size_t> CaffeWeightFactory::getBlobProtoData(const trtcaffe::BlobProto& blobMsg,
trtcaffe::Type type, std::vector<void*>& tmpAllocs)
{
// NVCaffe new binary format. It may carry any type.
if (blobMsg.has_raw_data())
{
assert(blobMsg.has_raw_data_type());
if (blobMsg.raw_data_type() == type)
{
return std::make_pair(&blobMsg.raw_data().front(),
blobMsg.raw_data().size() / sizeOfCaffeType(type));
}
}
// Old BVLC format.
if (blobMsg.data_size() > 0 && type == trtcaffe::FLOAT)
{
return std::make_pair(&blobMsg.data().Get(0), blobMsg.data_size());
}
// Converting to the target type otherwise
const int count = blobMsg.has_raw_data() ? blobMsg.raw_data().size() / sizeOfCaffeType(blobMsg.raw_data_type()) : (blobMsg.data_size() > 0 ? blobMsg.data_size() : blobMsg.double_data_size());
if (count > 0)
{
void* new_memory = malloc(count * sizeOfCaffeType(type));
tmpAllocs.push_back(new_memory);
if (type == trtcaffe::FLOAT)
{
auto* dst = reinterpret_cast<float*>(new_memory);
if (blobMsg.has_raw_data())
{
if (blobMsg.raw_data_type() == trtcaffe::FLOAT16)
{
const auto* src = reinterpret_cast<const float16*>(&blobMsg.raw_data().front());
for (int i = 0; i < count; ++i)
{
dst[i] = float(src[i]);
}
}
else if (blobMsg.raw_data_type() == trtcaffe::DOUBLE)
{
const auto* src = reinterpret_cast<const double*>(&blobMsg.raw_data().front());
for (int i = 0; i < count; ++i)
{
dst[i] = float(src[i]);
}
}
}
else if (blobMsg.double_data_size() == count)
{
for (int i = 0; i < count; ++i)
{
dst[i] = float(blobMsg.double_data(i));
}
}
return std::make_pair(new_memory, count);
}
if (type == trtcaffe::FLOAT16)
{
auto* dst = reinterpret_cast<float16*>(new_memory);
if (blobMsg.has_raw_data())
{
if (blobMsg.raw_data_type() == trtcaffe::FLOAT)
{
const auto* src = reinterpret_cast<const float*>(&blobMsg.raw_data().front());
for (int i = 0; i < count; ++i)
{
dst[i] = float16(src[i]);
}
}
else if (blobMsg.raw_data_type() == trtcaffe::DOUBLE)
{
const auto* src = reinterpret_cast<const double*>(&blobMsg.raw_data().front());
for (int i = 0; i < count; ++i)
{
dst[i] = float16(float(src[i]));
}
}
}
else if (blobMsg.data_size() == count)
{
for (int i = 0; i < count; ++i)
{
dst[i] = float16(blobMsg.data(i));
}
}
else if (blobMsg.double_data_size() == count)
{
for (int i = 0; i < count; ++i)
{
dst[i] = float16(float(blobMsg.double_data(i)));
}
}
return std::make_pair(new_memory, count);
}
}
return std::make_pair(nullptr, 0UL);
}
template <typename T>
bool CaffeWeightFactory::checkForNans(const void* values, int count, const std::string& layerName)
{
const T* v = reinterpret_cast<const T*>(values);
for (int i = 0; i < count; i++)
{
if (std::isnan(float(v[i])))
{
std::cout << layerName << ": Nan detected in weights" << std::endl;
return false;
}
}
return true;
}
Weights CaffeWeightFactory::getWeights(const trtcaffe::BlobProto& blobMsg, const std::string& layerName)
{
// Always load weights into FLOAT format
const auto blobProtoData = getBlobProtoData(blobMsg, trtcaffe::FLOAT, mTmpAllocs);
if (blobProtoData.first == nullptr)
{
const int bits = mDataType == DataType::kFLOAT ? 32 : 16;
std::cout << layerName << ": ERROR - " << bits << "-bit weights not found for "
<< bits << "-bit model" << std::endl;
mOK = false;
return Weights{DataType::kFLOAT, nullptr, 0};
}
mOK &= checkForNans<float>(blobProtoData.first, int(blobProtoData.second), layerName);
return Weights{DataType::kFLOAT, blobProtoData.first, int(blobProtoData.second)};
}
parsers/caffe/caffeWeightFactory/caffeWeightFactory.h
+68
View File
@@ -0,0 +1,68 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_CAFFE_WEIGHT_FACTORY_H
#define TRT_CAFFE_PARSER_CAFFE_WEIGHT_FACTORY_H
#include <vector>
#include <string>
#include <random>
#include <memory>
#include "NvInfer.h"
#include "weightType.h"
#include "trtcaffe.pb.h"
namespace nvcaffeparser1
{
class CaffeWeightFactory
{
public:
CaffeWeightFactory(const trtcaffe::NetParameter& msg, nvinfer1::DataType dataType, std::vector<void*>& tmpAllocs, bool isInitialized);
nvinfer1::DataType getDataType() const;
size_t getDataTypeSize() const;
std::vector<void*>& getTmpAllocs();
int getBlobsSize(const std::string& layerName);
const trtcaffe::BlobProto* getBlob(const std::string& layerName, int index);
std::vector<nvinfer1::Weights> getAllWeights(const std::string& layerName);
virtual nvinfer1::Weights operator()(const std::string& layerName, WeightType weightType);
void convert(nvinfer1::Weights& weights, nvinfer1::DataType targetType);
void convert(nvinfer1::Weights& weights);
bool isOK();
bool isInitialized();
nvinfer1::Weights getNullWeights();
nvinfer1::Weights allocateWeights(int64_t elems, std::uniform_real_distribution<float> distribution = std::uniform_real_distribution<float>(-0.01f, 0.01F));
nvinfer1::Weights allocateWeights(int64_t elems, std::normal_distribution<float> distribution);
static trtcaffe::Type getBlobProtoDataType(const trtcaffe::BlobProto& blobMsg);
static size_t sizeOfCaffeType(trtcaffe::Type type);
// The size returned here is the number of array entries, not bytes
static std::pair<const void*, size_t> getBlobProtoData(const trtcaffe::BlobProto& blobMsg, trtcaffe::Type type, std::vector<void*>& tmpAllocs);
private:
template <typename T>
bool checkForNans(const void* values, int count, const std::string& layerName);
nvinfer1::Weights getWeights(const trtcaffe::BlobProto& blobMsg, const std::string& layerName);
const trtcaffe::NetParameter& mMsg;
std::unique_ptr<trtcaffe::NetParameter> mRef;
std::vector<void*>& mTmpAllocs;
nvinfer1::DataType mDataType;
// bool mQuantize;
bool mInitialized;
std::default_random_engine generator;
bool mOK{true};
};
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_CAFFE_WEIGHT_FACTORY_H
parsers/caffe/caffeWeightFactory/weightType.h
+42
View File
@@ -0,0 +1,42 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_CAFFE_PARSER_WEIGHT_TYPE_H
#define TRT_CAFFE_PARSER_WEIGHT_TYPE_H
namespace nvcaffeparser1
{
enum class WeightType
{
// types for convolution, deconv, fully connected
kGENERIC = 0, // typical weights for the layer: e.g. filter (for conv) or matrix weights (for innerproduct)
kBIAS = 1, // bias weights
// These enums are for BVLCCaffe, which are incompatible with nvCaffe enums below.
// See batch_norm_layer.cpp in BLVC source of Caffe
kMEAN = 0,
kVARIANCE = 1,
kMOVING_AVERAGE = 2,
// These enums are for nvCaffe, which are incompatible with BVLCCaffe enums above
// See batch_norm_layer.cpp in NVidia fork of Caffe
kNVMEAN = 0,
kNVVARIANCE = 1,
kNVSCALE = 3,
kNVBIAS = 4
};
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_WEIGHT_TYPE_H
parsers/caffe/proto/trtcaffe.proto
+1819
View File
File diff suppressed because it is too large Load Diff
parsers/common/ParserCommonSrcs.txt
+21
View File
@@ -0,0 +1,21 @@
#
# Copyright (c) 2019, 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.
#
set(PARSER_COMMON_SRCS
common/half.h
common/ieee_half.h
common/parserHelper.h
)
parsers/common/half.h
+42
View File
@@ -0,0 +1,42 @@
/*
* Copyright (c) 2019, 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.
*/
//
// Custom wrapper around external half-precision header
//
// Header has some "extra parentheses" warnings when different rounding modes are used.
#if defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wparentheses"
#endif
#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wmismatched-tags"
#endif
#include "ieee_half.h"
typedef half_float::half float16;
#if defined(__clang__)
#pragma clang diagnostic pop
#endif
#if defined(__GNUC__)
#pragma GCC diagnostic pop
#endif
parsers/common/ieee_half.h
+3077
View File
File diff suppressed because it is too large Load Diff
parsers/common/parserUtils.h
+143
View File
@@ -0,0 +1,143 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_PARSER_UTILS_H
#define TRT_PARSER_UTILS_H
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cstdio>
#include <iostream>
#include <memory>
#ifndef _MSC_VER
#include <unistd.h>
#else
#define NOMINMAX
#include <windows.h>
#endif
#include "NvInfer.h"
namespace parserutils
{
#define RETURN_AND_LOG_ERROR_IMPL(ret, message, parserName) \
do \
{ \
std::string errorMsg = parserName + std::string{message}; \
if (getLogger()) getLogger()->log(nvinfer1::ILogger::Severity::kERROR, errorMsg.c_str()); \
else std::cerr << "WARNING: Logger not found, logging to stderr.\n" << errorMsg << std::endl; \
return (ret); \
} while (0)
// Helper function to compute unpadded volume of a Dims (1 if 0 dimensional)
inline int64_t volume(const nvinfer1::Dims& d)
{
int64_t v = 1;
for (int64_t i = 0; i < d.nbDims; i++)
v *= d.d[i];
return v;
}
// Show some debugging output about how much memory is free
inline void printMem(const char* where)
{
#if !defined _MSC_VER && !defined __QNX__
const unsigned mb = 1024 * 1024;
auto pages = static_cast<uint64_t>(sysconf(_SC_PHYS_PAGES));
auto avPages = static_cast<uint64_t>(sysconf(_SC_AVPHYS_PAGES));
auto pageSize = static_cast<uint64_t>(sysconf(_SC_PAGE_SIZE));
std::cout << " (memory) " << where << " : Free(MB) = " << avPages * pageSize / mb << " total(MB)=" << pages * pageSize / mb << std::endl;
#elif !defined __QNX__
const unsigned mb = 1024 * 1024;
MEMORYSTATUSEX statex;
statex.dwLength = sizeof(statex);
GlobalMemoryStatusEx(&statex);
std::cout << " (memory) " << where << " : Free(MB) = " << statex.ullAvailPhys / mb << " total(MB)=" << statex.ullTotalPhys / mb << std::endl;
#endif
}
// Compute size of datatypes
inline unsigned int elementSize(nvinfer1::DataType t)
{
switch (t)
{
case nvinfer1::DataType::kINT32: return 4;
case nvinfer1::DataType::kFLOAT: return 4;
case nvinfer1::DataType::kHALF: return 2;
case nvinfer1::DataType::kINT8: return 1;
}
assert(0);
return 0;
}
inline std::ostream& operator<<(std::ostream& o, const nvinfer1::Dims& dims)
{
o << "[";
for (int i = 0; i < dims.nbDims; i++)
o << (i ? "," : "") << dims.d[i];
o << "]";
return o;
}
inline std::ostream& operator<<(std::ostream& o, nvinfer1::DataType dt)
{
switch (dt)
{
case nvinfer1::DataType::kINT32: o << "Int32"; break;
case nvinfer1::DataType::kFLOAT: o << "Float"; break;
case nvinfer1::DataType::kHALF: o << "Half"; break;
case nvinfer1::DataType::kINT8: o << "Int8"; break;
}
return o;
}
inline nvinfer1::DimsCHW getCHW(const nvinfer1::Dims& d)
{
assert(d.nbDims >= 3);
return nvinfer1::DimsCHW(d.d[d.nbDims - 3], d.d[d.nbDims - 2], d.d[d.nbDims - 1]);
}
inline nvinfer1::DimsCHW getCHWWithExpansion(const nvinfer1::Dims& d, int filler)
{
if (d.nbDims == 0)
return nvinfer1::DimsCHW(filler, filler, filler);
else if (d.nbDims == 1)
return nvinfer1::DimsCHW(filler, filler, d.d[0]);
else if (d.nbDims == 2)
return nvinfer1::DimsCHW(filler, d.d[0], d.d[1]);
else
return nvinfer1::DimsCHW(d.d[d.nbDims - 3], d.d[d.nbDims - 2], d.d[d.nbDims - 1]);
}
inline int combineIndexDimensions(int batchSize, const nvinfer1::Dims& d)
{
int x = batchSize;
for (int i = 0; i < d.nbDims - 3; i++)
x *= d.d[i];
return x;
}
template <typename A, typename B>
inline A divUp(A m, B n)
{
return (m + n - 1) / n;
}
} // namespace parserhelper
#endif // PARSER_HELPER_H
parsers/onnx
Submodule
+1
Submodule parsers/onnx added at 8b52755beb
plugin/CMakeLists.txt
+132
View File
@@ -0,0 +1,132 @@
#
# Copyright (c) 2019, 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.
#
add_custom_target(plugin)
set(TARGET_NAME nvinfer_plugin)
set(SHARED_TARGET ${TARGET_NAME})
set(STATIC_TARGET ${TARGET_NAME}_static)
set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR})
set(PLUGIN_EXPORT_MAP ${TARGET_DIR}/exports.map)
if(${CMAKE_BUILD_TYPE} MATCHES "Debug")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g")
endif()
set(PLUGIN_SRCS)
set(PLUGIN_SOURCES)
set(CUDA_SRCS)
set(COMMON_SRCS)
set(PLUGIN_LISTS
nmsPlugin
normalizePlugin
priorBoxPlugin
reorgPlugin
gridAnchorPlugin
regionPlugin
nvFasterRCNN
batchedNMSPlugin
flattenConcat
cropAndResizePlugin
proposalPlugin
)
include_directories(common common/kernels)
foreach(PLUGIN_ITER ${PLUGIN_LISTS})
include_directories(${PLUGIN_ITER})
add_subdirectory(${PLUGIN_ITER})
endforeach(PLUGIN_ITER)
# Add common
add_subdirectory(common)
list(APPEND PLUGIN_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/InferPlugin.cpp")
################################## SHARED LIBRARY #######################################
add_library(${SHARED_TARGET} SHARED
${PLUGIN_SOURCES}
)
target_include_directories(${SHARED_TARGET}
PUBLIC ${PROJECT_SOURCE_DIR}/include
PUBLIC ${CUB_ROOT_DIR}
PRIVATE ${PROJECT_SOURCE_DIR}/common
PUBLIC ${CUDA_INSTALL_DIR}/include
PRIVATE ${TARGET_DIR}
)
set_target_properties(${SHARED_TARGET} PROPERTIES
CXX_STANDARD "11"
CXX_STANDARD_REQUIRED "YES"
CXX_EXTENSIONS "NO"
ARCHIVE_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
LIBRARY_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
RUNTIME_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
)
set_target_properties(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-Wl,--exclude-libs,ALL -Wl,--version-script=${PLUGIN_EXPORT_MAP}")
set_target_properties(${SHARED_TARGET} PROPERTIES DEBUG_POSTFIX ${TRT_DEBUG_POSTFIX})
set_target_properties(${SHARED_TARGET} PROPERTIES VERSION ${TRT_VERSION} SOVERSION ${TRT_SOVERSION} )
set_property(TARGET ${SHARED_TARGET} PROPERTY CUDA_STANDARD 11)
################################## STATIC LIBRARY #######################################
add_library(${STATIC_TARGET} STATIC
${PLUGIN_SOURCES}
)
target_include_directories(${STATIC_TARGET}
PUBLIC ${PROJECT_SOURCE_DIR}/include
PUBLIC ${CUB_ROOT_DIR}
PRIVATE ${PROJECT_SOURCE_DIR}/common
PUBLIC ${CUDA_INSTALL_DIR}/include
PRIVATE ${TARGET_DIR}
)
set_target_properties(${STATIC_TARGET} PROPERTIES
CXX_STANDARD "11"
CXX_STANDARD_REQUIRED "YES"
CXX_EXTENSIONS "NO"
ARCHIVE_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
LIBRARY_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
RUNTIME_OUTPUT_DIRECTORY "${TRT_BIN_DIR}"
)
set_target_properties(${STATIC_TARGET} PROPERTIES LINK_FLAGS "-Wl,--exclude-libs,ALL")
set_target_properties(${STATIC_TARGET} PROPERTIES DEBUG_POSTFIX ${TRT_DEBUG_POSTFIX})
set_target_properties(${STATIC_TARGET} PROPERTIES VERSION ${TRT_VERSION} SOVERSION ${TRT_SOVERSION} )
set_property(TARGET ${STATIC_TARGET} PROPERTY CUDA_STANDARD 11)
#########################################################################################
add_dependencies(plugin ${SHARED_TARGET} ${STATIC_TARGET})
################################### INSTALLATION ########################################
install(TARGETS ${TARGET_NAME}
RUNTIME DESTINATION bin
LIBRARY DESTINATION lib
ARCHIVE DESTINATION lib
)
plugin/InferPlugin.cpp
+111
View File
@@ -0,0 +1,111 @@
/*
* Copyright (c) 2019, 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.
*/
#include "NvInfer.h"
#include "NvInferPlugin.h"
#include <algorithm>
#include <array>
#include <iostream>
#include <memory>
using namespace nvinfer1;
using namespace nvinfer1::plugin;
#include "batchedNMSPlugin/batchedNMSPlugin.h"
#include "cropAndResizePlugin/cropAndResizePlugin.h"
#include "flattenConcat/flattenConcat.h"
#include "gridAnchorPlugin/gridAnchorPlugin.h"
#include "nmsPlugin/nmsPlugin.h"
#include "normalizePlugin/normalizePlugin.h"
#include "nvFasterRCNN/nvFasterRCNNPlugin.h"
#include "priorBoxPlugin/priorBoxPlugin.h"
#include "proposalPlugin/proposalPlugin.h"
#include "regionPlugin/regionPlugin.h"
#include "reorgPlugin/reorgPlugin.h"
using nvinfer1::plugin::RPROIParams;
namespace nvinfer1
{
namespace internal
{
extern ILogger* gLogger;
}
namespace plugin
{
// Instances of this class are statically constructed in initializePlugin.
// This ensures that each plugin is only registered a single time, as further calls to
// initializePlugin will be no-ops.
template <typename CreatorType>
class InitializePlugin
{
public:
InitializePlugin(void* logger, const char* libNamespace)
: mCreator{new CreatorType{}}
{
mCreator->setPluginNamespace(libNamespace);
bool status = getPluginRegistry()->registerCreator(*mCreator, libNamespace);
if (logger)
{
nvinfer1::internal::gLogger = static_cast<nvinfer1::ILogger*>(logger);
if (!status)
{
std::string errorMsg{"Could not register plugin creator: " + std::string(mCreator->getPluginName())
+ " in namespace: " + std::string{mCreator->getPluginNamespace()}};
nvinfer1::internal::gLogger->log(ILogger::Severity::kERROR, errorMsg.c_str());
}
else
{
std::string verboseMsg{
"Plugin Creator registration succeeded - " + std::string{mCreator->getPluginName()}};
nvinfer1::internal::gLogger->log(ILogger::Severity::kVERBOSE, verboseMsg.c_str());
}
}
}
InitializePlugin(const InitializePlugin&) = delete;
InitializePlugin(InitializePlugin&&) = delete;
private:
std::unique_ptr<CreatorType> mCreator;
};
template <typename CreatorType>
void initializePlugin(void* logger, const char* libNamespace)
{
static InitializePlugin<CreatorType> plugin{logger, libNamespace};
}
} // namespace plugin
} // namespace nvinfer1
extern "C" {
bool initLibNvInferPlugins(void* logger, const char* libNamespace)
{
initializePlugin<nvinfer1::plugin::GridAnchorPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::NMSPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::ReorgPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::RegionPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::PriorBoxPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::NormalizePluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::RPROIPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::BatchedNMSPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::FlattenConcatPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::CropAndResizePluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::ProposalPluginCreator>(logger, libNamespace);
return true;
}
} // extern "C"
plugin/batchedNMSPlugin/CMakeLists.txt
+20
View File
@@ -0,0 +1,20 @@
#
# Copyright (c) 2019, 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.
#
file(GLOB SRCS *.cpp *.cu)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
plugin/batchedNMSPlugin/README.md
+113
View File
@@ -0,0 +1,113 @@
# batchedNMSPlugin
**Table Of Contents**
- [Description](#description)
* [Structure](#structure)
- [Parameters](#parameters)
- [Algorithms](#algorithms)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
The `batchedNMSPlugin` implements a non-maximum suppression (NMS) step over boxes for object detection networks.
Non-maximum suppression is typically the universal step in object detection inference. This plugin is used after you’ve processed the bounding box prediction and object classification to get the final bounding boxes for objects.
With this plugin, you can incorporate the non-maximum suppression step during TensorRT inference. During inference, the neural network generates a fixed number of bounding boxes with box coordinates, identified class and confidence levels. Not all bounding boxes, but the most representative ones, have to be drawn on the original image.
Non-maximum suppression is the way to eliminate the boxes which have low confidence or do not have object in and keep the most representative ones. For example, the objects within an image might be covered by many boxes with different levels of confidence. The goal of the non-maximum suppression step is to find the most confident box for the object and remove all the less confident ones.
This plugin accelerates this non maximum suppression step during TensorRT inference on GPU.
### Structure
The `batchedNMSPlugin` takes two inputs, boxes input and scores input.
**Boxes input**
The boxes input are of shape `[batch_size, number_boxes, number_classes, number_box_parameters]`. The box location usually consists of four parameters such as `[x_min, y_min, x_max, y_max]`. For example, if your model outputs `8732` bounding boxes given one image, there are `100` candidate classes, the shape of boxes input will be `[8732, 100, 4]`.
**Scores input**
The scores input are of shape `[batch_size, number_boxes, number_classes]`. Each box has an array of probability for each candidate class.
The boxes input and scores input generates the following four outputs:
- `num_detections`
The `num_detections` input are of shape `[batch_size, 1]`. The last dimension of size 1 is an INT32 scalar indicating the number of valid detections per batch item. It can be less than `keepTopK`. Only the top `num_detections[i]` entries in `nmsed_boxes[i]`, `nmsed_scores[i]` and `nmsed_classes[i]` are valid.
- `nmsed_boxes`
A `[batch_size, keepTopK, 4]` float32 tensor containing the coordinates of non-max suppressed boxes.
- `nmsed_scores`
A `[batch_size, keepTopK]` float32 tensor containing the scores for the boxes.
- `nmsed_classes`
A `[batch_size, keepTopK]` float32 tensor containing the classes for the boxes.
## Parameters
The `batchedNMSPlugin` has plugin creator class `BatchedNMSPluginCreator` and plugin class `BatchedNMSPlugin`.
The `batchedNMSPlugin` is created using `BatchedNMSPluginCreator` with `NMSParameters` typed parameters. The `NMSParameters` data structure is listed as follows and is defined in the [NvInferPlugin.h header file](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/_nv_infer_plugin_8h_source.html).
| Type | Parameter | Description
|----------|--------------------------|--------------------------------------------------------
|`bool` |`shareLocation` |If set to `true`, the boxes input are shared across all classes. If set to `false`, the boxes input should account for per-class box data.
|`int` |`backgroundLabelId` |The label ID for the background class. If there is no background class, set it to `-1`.
|`int` |`numClasses` |The number of classes in the network.
|`int` |`topK` |The number of bounding boxes to be fed into the NMS step.
|`int` |`keepTopK` |The number of total bounding boxes to be kept per-image after the NMS step. Should be less than or equal to the `topK` value.
|`float` |`scoreThreshold` |The scalar threshold for score (low scoring boxes are removed).
|`float` |`iouThreshold` |The 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, meaning they are not in the range `[0,1]`. Defaults to `true`.
## Algorithms
The NMS algorithm used in this particular plugin first sorts the bounding boxes indices by the score for each class, then sorts the bounding boxes by the updated scores, and finally collects the desired number of bounding boxes with the highest scores.
It is mainly accelerated using the `nmsInference` kernel defined in the `batchedNMSInference.cu` file.
Specifically, the NMS algorithm:
- Sorts the bounding box indices by the score for each class. Before sorting, the bounding boxes with a score less than `scoreThreshold` are discarded by setting their indices to `-1` and their scores to `0`. This is using the `sortScoresPerClass` kernel defined in the `sortScoresPerClass.cu` file.
- Finds the most confident box for the object and removes all the less confident ones using the iterative non-maximum suppression step step for each class. Starting from the bounding box with the highest score in each class, the bounding boxes that has overlap higher than `iouThreshold` is suppressed by setting their indices to `-1` and their scores to `0`. Then all the less confident bounding boxes were suppressed for each class. This is using the `allClassNMS` kernel defined in the `allClassNMS.cu` file.
- Sorts the bounding boxes per image using the updated scores. At this time, all the classes were mixed before sort. Discarded and suppressed bounding boxes will go to the end of the sorted array since their score is `0`. This is using the `sortScoresPerImage` kernel defined in the `sortScoresPerImage.cu` file.
- Collects the desired number, `keepTopK`, of bounding box indices with the highest scores from the top of the sorted array, their bounding box coordinates, and their object classification information. This is using the `gatherNMSOutputs` kernel defined in the `gatherNMSOutputs.cu` file.
## Additional resources
The following resources provide a deeper understanding of the `batchedNMSPlugin` plugin:
**Networks**
- [SSD: Single Shot MultiBox Detector](https://arxiv.org/abs/1512.02325)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://arxiv.org/abs/1506.01497)
- [Mask R-CNN](https://arxiv.org/abs/1703.06870)
**Documentation**
- [NMSParameter detailed descriptions](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/structnvinfer1_1_1plugin_1_1_n_m_s_parameters.html)
- [NMS algorithm](https://www.coursera.org/lecture/convolutional-neural-networks/non-max-suppression-dvrjH)
## License
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
documentation.
## Changelog
May 2019
This is the first release of this `README.md` file.
## Known issues
There are no known issues in this plugin.
plugin/batchedNMSPlugin/batchedNMSInference.cpp
+126
View File
@@ -0,0 +1,126 @@
/*
* Copyright (c) 2019, 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.
*/
#include "bboxUtils.h"
#include "cuda_runtime_api.h"
#include "gatherNMSOutputs.h"
#include "kernel.h"
#include "nmsUtils.h"
pluginStatus_t nmsInference(cudaStream_t stream, const int N, const int perBatchBoxesSize, const int perBatchScoresSize,
const bool shareLocation, const int backgroundLabelId, const int numPredsPerClass, const int numClasses,
const int topK, const int keepTopK, const float scoreThreshold, const float iouThreshold, const DataType DT_BBOX,
const void* locData, const DataType DT_SCORE, const void* confData, void* keepCount, void* nmsedBoxes,
void* nmsedScores, void* nmsedClasses, void* workspace, bool isNormalized, bool confSigmoid, bool clipBoxes)
{
// locCount = batch_size * number_boxes_per_sample * 4
const int locCount = N * perBatchBoxesSize;
/*
* shareLocation
* Bounding box are shared among all classes, i.e., a bounding box could be classified as any candidate class.
* Otherwise
* Bounding box are designed for specific classes, i.e., a bounding box could be classified as one certain class or
* not (binary classification).
*/
const int numLocClasses = shareLocation ? 1 : numClasses;
size_t bboxDataSize = detectionForwardBBoxDataSize(N, perBatchBoxesSize, DataType::kFLOAT);
void* bboxDataRaw = workspace;
cudaMemcpyAsync(bboxDataRaw, locData, bboxDataSize, cudaMemcpyDeviceToDevice, stream);
pluginStatus_t status;
/*
* bboxDataRaw format:
* [batch size, numPriors (per sample), numLocClasses, 4]
*/
// float for now
void* bboxData;
size_t bboxPermuteSize = detectionForwardBBoxPermuteSize(shareLocation, N, perBatchBoxesSize, DataType::kFLOAT);
void* bboxPermute = nextWorkspacePtr((int8_t*) bboxDataRaw, bboxDataSize);
/*
* After permutation, bboxData format:
* [batch_size, numLocClasses, numPriors (per sample) (numPredsPerClass), 4]
* This is equivalent to swapping axis
*/
if (!shareLocation)
{
status = permuteData(
stream, locCount, numLocClasses, numPredsPerClass, 4, DataType::kFLOAT, false, bboxDataRaw, bboxPermute);
ASSERT_FAILURE(status == STATUS_SUCCESS);
bboxData = bboxPermute;
}
/*
* If shareLocation, numLocClasses = 1
* No need to permute data on linear memory
*/
else
{
bboxData = bboxDataRaw;
}
/*
* Conf data format
* [batch size, numPriors * param.numClasses, 1, 1]
*/
const int numScores = N * perBatchScoresSize;
size_t totalScoresSize = detectionForwardPreNMSSize(N, perBatchScoresSize);
void* scores = nextWorkspacePtr((int8_t*) bboxPermute, bboxPermuteSize);
// need a conf_scores
/*
* After permutation, bboxData format:
* [batch_size, numClasses, numPredsPerClass, 1]
*/
status = permuteData(
stream, numScores, numClasses, numPredsPerClass, 1, DataType::kFLOAT, confSigmoid, confData, scores);
ASSERT_FAILURE(status == STATUS_SUCCESS);
size_t indicesSize = detectionForwardPreNMSSize(N, perBatchScoresSize);
void* indices = nextWorkspacePtr((int8_t*) scores, totalScoresSize);
size_t postNMSScoresSize = detectionForwardPostNMSSize(N, numClasses, topK);
size_t postNMSIndicesSize = detectionForwardPostNMSSize(N, numClasses, topK);
void* postNMSScores = nextWorkspacePtr((int8_t*) indices, indicesSize);
void* postNMSIndices = nextWorkspacePtr((int8_t*) postNMSScores, postNMSScoresSize);
void* sortingWorkspace = nextWorkspacePtr((int8_t*) postNMSIndices, postNMSIndicesSize);
// Sort the scores so that the following NMS could be applied.
status = sortScoresPerClass(stream, N, numClasses, numPredsPerClass, backgroundLabelId, scoreThreshold,
DataType::kFLOAT, scores, indices, sortingWorkspace);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// This is set to true as the input bounding boxes are of the format [ymin,
// xmin, ymax, xmax]. The default implementation assumes [xmin, ymin, xmax, ymax]
bool flipXY = true;
// NMS
status = allClassNMS(stream, N, numClasses, numPredsPerClass, topK, iouThreshold, shareLocation, isNormalized,
DataType::kFLOAT, DataType::kFLOAT, bboxData, scores, indices, postNMSScores, postNMSIndices, flipXY);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// Sort the bounding boxes after NMS using scores
status = sortScoresPerImage(stream, N, numClasses * topK, DataType::kFLOAT, postNMSScores, postNMSIndices, scores,
indices, sortingWorkspace);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// Gather data from the sorted bounding boxes after NMS
status = gatherNMSOutputs(stream, shareLocation, N, numPredsPerClass, numClasses, topK, keepTopK, DataType::kFLOAT,
DataType::kFLOAT, indices, scores, bboxData, keepCount, nmsedBoxes, nmsedScores, nmsedClasses, clipBoxes);
ASSERT_FAILURE(status == STATUS_SUCCESS);
return STATUS_SUCCESS;
}
plugin/batchedNMSPlugin/batchedNMSInference.h
+28
View File
@@ -0,0 +1,28 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_BATCHED_NMS_INFERENCE_H
#define TRT_BATCHED_NMS_INFERENCE_H
#include "plugin.h"
using namespace nvinfer1;
using namespace nvinfer1::plugin;
pluginStatus_t nmsInference(cudaStream_t stream, int N, int boxesSize, int scoresSize, bool shareLocation,
int backgroundLabelId, int numPredsPerClass, int numClasses, int topK, int keepTopK, float scoreThreshold,
float iouThreshold, DataType DT_BBOX, const void* locData, DataType DT_SCORE, const void* confData, void* keepCount,
void* nmsedBoxes, void* nmsedScores, void* nmsedClasses, void* workspace, bool isNormalized = true,
bool confSigmoid = false, bool clipBoxes = true);
#endif
plugin/batchedNMSPlugin/batchedNMSPlugin.cpp
+320
View File
@@ -0,0 +1,320 @@
/*
* Copyright (c) 2019, 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.
*/
#include "batchedNMSPlugin.h"
#include <algorithm>
#include <cstring>
#include <iostream>
#include <sstream>
#include <vector>
using namespace nvinfer1;
using nvinfer1::plugin::BatchedNMSPlugin;
using nvinfer1::plugin::BatchedNMSPluginCreator;
using nvinfer1::plugin::NMSParameters;
namespace
{
const char* NMS_PLUGIN_VERSION{"1"};
const char* NMS_PLUGIN_NAME{"BatchedNMS_TRT"};
} // namespace
PluginFieldCollection BatchedNMSPluginCreator::mFC{};
std::vector<PluginField> BatchedNMSPluginCreator::mPluginAttributes;
BatchedNMSPlugin::BatchedNMSPlugin(NMSParameters params)
: param(params)
{
}
BatchedNMSPlugin::BatchedNMSPlugin(const void* data, size_t length)
{
const char *d = reinterpret_cast<const char*>(data), *a = d;
param = read<NMSParameters>(d);
boxesSize = read<int>(d);
scoresSize = read<int>(d);
numPriors = read<int>(d);
mClipBoxes = read<bool>(d);
ASSERT(d == a + length);
}
int BatchedNMSPlugin::getNbOutputs() const
{
return 4;
}
int BatchedNMSPlugin::initialize()
{
return STATUS_SUCCESS;
}
void BatchedNMSPlugin::terminate() {}
Dims BatchedNMSPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
{
ASSERT(nbInputDims == 2);
ASSERT(index >= 0 && index < this->getNbOutputs());
ASSERT(inputs[0].nbDims == 3);
ASSERT(inputs[1].nbDims == 2);
// boxesSize: number of box coordinates for one sample
boxesSize = inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2];
// scoresSize: number of scores for one sample
scoresSize = inputs[1].d[0] * inputs[1].d[1];
// num_detections
if (index == 0)
{
Dims dim0;
dim0.nbDims = 0;
return dim0;
}
// nmsed_boxes
if (index == 1)
{
return DimsHW(param.keepTopK, 4);
}
// nmsed_scores or nmsed_classes
Dims dim1;
dim1.nbDims = 1;
dim1.d[0] = param.keepTopK;
return dim1;
}
size_t BatchedNMSPlugin::getWorkspaceSize(int maxBatchSize) const
{
return detectionInferenceWorkspaceSize(param.shareLocation, maxBatchSize, boxesSize, scoresSize, param.numClasses,
numPriors, param.topK, DataType::kFLOAT, DataType::kFLOAT);
}
int BatchedNMSPlugin::enqueue(
int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)
{
const void* const locData = inputs[0];
const void* const confData = inputs[1];
void* keepCount = outputs[0];
void* nmsedBoxes = outputs[1];
void* nmsedScores = outputs[2];
void* nmsedClasses = outputs[3];
pluginStatus_t status = nmsInference(stream, batchSize, boxesSize, scoresSize, param.shareLocation,
param.backgroundLabelId, numPriors, param.numClasses, param.topK, param.keepTopK, param.scoreThreshold,
param.iouThreshold, DataType::kFLOAT, locData, DataType::kFLOAT, confData, keepCount, nmsedBoxes, nmsedScores,
nmsedClasses, workspace, param.isNormalized, false, mClipBoxes);
ASSERT(status == STATUS_SUCCESS);
return 0;
}
size_t BatchedNMSPlugin::getSerializationSize() const
{
// NMSParameters, boxesSize,scoresSize,numPriors
return sizeof(NMSParameters) + sizeof(int) * 3 + sizeof(bool);
}
void BatchedNMSPlugin::serialize(void* buffer) const
{
char *d = reinterpret_cast<char*>(buffer), *a = d;
write(d, param);
write(d, boxesSize);
write(d, scoresSize);
write(d, numPriors);
write(d, mClipBoxes);
ASSERT(d == a + getSerializationSize());
}
void BatchedNMSPlugin::configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, nvinfer1::PluginFormat format, int maxBatchSize)
{
ASSERT(nbInputs == 2);
ASSERT(nbOutputs == 4);
ASSERT(inputDims[0].nbDims == 3);
ASSERT(inputDims[1].nbDims == 2);
ASSERT(std::none_of(inputIsBroadcast, inputIsBroadcast + nbInputs, [](bool b) { return b; }));
ASSERT(std::none_of(outputIsBroadcast, outputIsBroadcast + nbInputs, [](bool b) { return b; }));
boxesSize = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2];
scoresSize = inputDims[1].d[0] * inputDims[1].d[1];
// num_boxes
numPriors = inputDims[0].d[0];
const int numLocClasses = param.shareLocation ? 1 : param.numClasses;
// Third dimension of boxes must be either 1 or num_classes
ASSERT(inputDims[0].d[1] == numLocClasses);
ASSERT(inputDims[0].d[2] == 4);
}
bool BatchedNMSPlugin::supportsFormat(DataType type, PluginFormat format) const
{
return ((type == DataType::kFLOAT || type == DataType::kINT32) && format == PluginFormat::kNCHW);
}
const char* BatchedNMSPlugin::getPluginType() const
{
return NMS_PLUGIN_NAME;
}
const char* BatchedNMSPlugin::getPluginVersion() const
{
return NMS_PLUGIN_VERSION;
}
void BatchedNMSPlugin::destroy()
{
delete this;
}
IPluginV2Ext* BatchedNMSPlugin::clone() const
{
auto* plugin = new BatchedNMSPlugin(param);
plugin->boxesSize = boxesSize;
plugin->scoresSize = scoresSize;
plugin->numPriors = numPriors;
plugin->setPluginNamespace(mNamespace.c_str());
plugin->setClipParam(mClipBoxes);
return plugin;
}
void BatchedNMSPlugin::setPluginNamespace(const char* pluginNamespace)
{
mPluginNamespace = pluginNamespace;
}
const char* BatchedNMSPlugin::getPluginNamespace() const
{
return mPluginNamespace;
}
nvinfer1::DataType BatchedNMSPlugin::getOutputDataType(
int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
{
if (index == 0)
{
return nvinfer1::DataType::kINT32;
}
return inputTypes[0];
}
void BatchedNMSPlugin::setClipParam(bool clip)
{
mClipBoxes = clip;
}
bool BatchedNMSPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
{
return false;
}
bool BatchedNMSPlugin::canBroadcastInputAcrossBatch(int inputIndex) const
{
return false;
}
BatchedNMSPluginCreator::BatchedNMSPluginCreator()
: params{}
{
mPluginAttributes.emplace_back(PluginField("shareLocation", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("backgroundLabelId", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("numClasses", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("topK", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("keepTopK", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("scoreThreshold", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("iouThreshold", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("isNormalized", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("clipBoxes", nullptr, PluginFieldType::kINT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
const char* BatchedNMSPluginCreator::getPluginName() const
{
return NMS_PLUGIN_NAME;
}
const char* BatchedNMSPluginCreator::getPluginVersion() const
{
return NMS_PLUGIN_VERSION;
}
const PluginFieldCollection* BatchedNMSPluginCreator::getFieldNames()
{
return &mFC;
}
IPluginV2Ext* BatchedNMSPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
const PluginField* fields = fc->fields;
mClipBoxes = true;
for (int i = 0; i < fc->nbFields; ++i)
{
const char* attrName = fields[i].name;
if (!strcmp(attrName, "shareLocation"))
{
params.shareLocation = *(static_cast<const bool*>(fields[i].data));
}
else if (!strcmp(attrName, "backgroundLabelId"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
params.backgroundLabelId = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "numClasses"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
params.numClasses = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "topK"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
params.topK = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "keepTopK"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
params.keepTopK = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "scoreThreshold"))
{
ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
params.scoreThreshold = *(static_cast<const float*>(fields[i].data));
}
else if (!strcmp(attrName, "iouThreshold"))
{
ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
params.iouThreshold = *(static_cast<const float*>(fields[i].data));
}
else if (!strcmp(attrName, "isNormalized"))
{
params.isNormalized = *(static_cast<const bool*>(fields[i].data));
}
else if (!strcmp(attrName, "clipBoxes"))
{
mClipBoxes = *(static_cast<const bool*>(fields[i].data));
}
}
BatchedNMSPlugin* plugin = new BatchedNMSPlugin(params);
plugin->setClipParam(mClipBoxes);
plugin->setPluginNamespace(mNamespace.c_str());
return plugin;
}
IPluginV2Ext* BatchedNMSPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
{
// This object will be deleted when the network is destroyed, which will
// call NMS::destroy()
BatchedNMSPlugin* plugin = new BatchedNMSPlugin(serialData, serialLength);
plugin->setPluginNamespace(mNamespace.c_str());
return plugin;
}
plugin/batchedNMSPlugin/batchedNMSPlugin.h
+118
View File
@@ -0,0 +1,118 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_BATCHED_NMS_PLUGIN_H
#define TRT_BATCHED_NMS_PLUGIN_H
#include "batchedNMSPlugin/batchedNMSInference.h"
#include "batchedNMSPlugin/gatherNMSOutputs.h"
#include "kernel.h"
#include "nmsUtils.h"
#include "plugin.h"
#include <string>
#include <vector>
using namespace nvinfer1::plugin;
namespace nvinfer1
{
namespace plugin
{
class BatchedNMSPlugin : public IPluginV2Ext
{
public:
BatchedNMSPlugin(NMSParameters param);
BatchedNMSPlugin(const void* data, size_t length);
~BatchedNMSPlugin() override = default;
int getNbOutputs() const override;
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
int initialize() override;
void terminate() override;
size_t getWorkspaceSize(int maxBatchSize) const override;
int enqueue(
int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override;
size_t getSerializationSize() const override;
void serialize(void* buffer) const override;
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) override;
bool supportsFormat(DataType type, PluginFormat format) const override;
const char* getPluginType() const override;
const char* getPluginVersion() const override;
void destroy() override;
IPluginV2Ext* clone() const override;
nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputType, int nbInputs) const override;
void setPluginNamespace(const char* libNamespace) override;
const char* getPluginNamespace() const override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
bool canBroadcastInputAcrossBatch(int inputIndex) const override;
void setClipParam(bool clip);
private:
NMSParameters param;
int boxesSize, scoresSize, numPriors;
std::string mNamespace;
bool mClipBoxes;
const char* mPluginNamespace;
};
class BatchedNMSPluginCreator : public BaseCreator
{
public:
BatchedNMSPluginCreator();
~BatchedNMSPluginCreator() override = default;
const char* getPluginName() const override;
const char* getPluginVersion() const override;
const PluginFieldCollection* getFieldNames() override;
IPluginV2Ext* createPlugin(const char* name, const PluginFieldCollection* fc) override;
IPluginV2Ext* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
private:
static PluginFieldCollection mFC;
NMSParameters params;
static std::vector<PluginField> mPluginAttributes;
bool mClipBoxes;
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_BATCHED_NMS_PLUGIN_H
plugin/batchedNMSPlugin/gatherNMSOutputs.cu
+219
View File
@@ -0,0 +1,219 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "plugin.h"
#include "gatherNMSOutputs.h"
#include <vector>
template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void gatherNMSOutputs_kernel(
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const int* indices,
const T_SCORE* scores,
const T_BBOX* bboxData,
int* numDetections,
T_BBOX* nmsedBoxes,
T_BBOX* nmsedScores,
T_BBOX* nmsedClasses,
bool clipBoxes
)
{
if (keepTopK > topK)
return;
for (int i = blockIdx.x * nthds_per_cta + threadIdx.x;
i < numImages * keepTopK;
i += gridDim.x * nthds_per_cta)
{
const int imgId = i / keepTopK;
const int detId = i % keepTopK;
const int offset = imgId * numClasses * topK;
const int index = indices[offset + detId];
const T_SCORE score = scores[offset + detId];
if (index == -1)
{
nmsedClasses[i] = -1;
nmsedScores[i] = 0;
nmsedBoxes[i * 4] = 0;
nmsedBoxes[i * 4 + 1] = 0;
nmsedBoxes[i * 4 + 2] = 0;
nmsedBoxes[i * 4 + 3] = 0;
}
else
{
const int bboxOffset = imgId * (shareLocation ? numPredsPerClass : (numClasses * numPredsPerClass));
const int bboxId = ((shareLocation ? (index % numPredsPerClass)
: index % (numClasses * numPredsPerClass)) + bboxOffset) * 4;
nmsedClasses[i] = (index % (numClasses * numPredsPerClass)) / numPredsPerClass; // label
nmsedScores[i] = score; // confidence score
// clipped bbox xmin
nmsedBoxes[i * 4] = clipBoxes ? max(min(bboxData[bboxId],
T_BBOX(1.)), T_BBOX(0.)) : bboxData[bboxId];
// clipped bbox ymin
nmsedBoxes[i * 4 + 1] = clipBoxes ? max(min(bboxData[bboxId + 1],
T_BBOX(1.)), T_BBOX(0.)) : bboxData[bboxId + 1];
// clipped bbox xmax
nmsedBoxes[i * 4 + 2] = clipBoxes ? max(min(bboxData[bboxId + 2],
T_BBOX(1.)), T_BBOX(0.)) : bboxData[bboxId + 2];
// clipped bbox ymax
nmsedBoxes[i * 4 + 3] = clipBoxes ? max(min(bboxData[bboxId + 3],
T_BBOX(1.)), T_BBOX(0.)) : bboxData[bboxId + 3];
atomicAdd(&numDetections[i / keepTopK], 1);
}
}
}
template <typename T_BBOX, typename T_SCORE>
pluginStatus_t gatherNMSOutputs_gpu(
cudaStream_t stream,
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const void* indices,
const void* scores,
const void* bboxData,
void* numDetections,
void* nmsedBoxes,
void* nmsedScores,
void* nmsedClasses,
bool clipBoxes
)
{
cudaMemsetAsync(numDetections, 0, numImages * sizeof(int), stream);
const int BS = 32;
const int GS = 32;
gatherNMSOutputs_kernel<T_BBOX, T_SCORE, BS><<<GS, BS, 0, stream>>>(shareLocation, numImages, numPredsPerClass,
numClasses, topK, keepTopK,
(int*) indices, (T_SCORE*) scores, (T_BBOX*) bboxData,
(int*) numDetections,
(T_BBOX*) nmsedBoxes,
(T_BBOX*) nmsedScores,
(T_BBOX*) nmsedClasses,
clipBoxes
);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// gatherNMSOutputs LAUNCH CONFIG {{{
typedef pluginStatus_t (*nmsOutFunc)(cudaStream_t,
const bool,
const int,
const int,
const int,
const int,
const int,
const void*,
const void*,
const void*,
void*,
void*,
void*,
void*,
bool);
struct nmsOutLaunchConfig
{
DataType t_bbox;
DataType t_score;
nmsOutFunc function;
nmsOutLaunchConfig(DataType t_bbox, DataType t_score)
: t_bbox(t_bbox)
, t_score(t_score)
{
}
nmsOutLaunchConfig(DataType t_bbox, DataType t_score, nmsOutFunc function)
: t_bbox(t_bbox)
, t_score(t_score)
, function(function)
{
}
bool operator==(const nmsOutLaunchConfig& other)
{
return t_bbox == other.t_bbox && t_score == other.t_score;
}
};
using nvinfer1::DataType;
static std::vector<nmsOutLaunchConfig> nmsOutFuncVec;
bool nmsOutputInit()
{
nmsOutFuncVec.push_back(nmsOutLaunchConfig(DataType::kFLOAT, DataType::kFLOAT,
gatherNMSOutputs_gpu<float, float>));
return true;
}
static bool initialized = nmsOutputInit();
//}}}
pluginStatus_t gatherNMSOutputs(
cudaStream_t stream,
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const DataType DT_BBOX,
const DataType DT_SCORE,
const void* indices,
const void* scores,
const void* bboxData,
void* numDetections,
void* nmsedBoxes,
void* nmsedScores,
void* nmsedClasses,
bool clipBoxes
)
{
nmsOutLaunchConfig lc = nmsOutLaunchConfig(DT_BBOX, DT_SCORE);
for (unsigned i = 0; i < nmsOutFuncVec.size(); ++i)
{
if (lc == nmsOutFuncVec[i])
{
DEBUG_PRINTF("gatherNMSOutputs kernel %d\n", i);
return nmsOutFuncVec[i].function(stream,
shareLocation,
numImages,
numPredsPerClass,
numClasses,
topK,
keepTopK,
indices,
scores,
bboxData,
numDetections,
nmsedBoxes,
nmsedScores,
nmsedClasses,
clipBoxes
);
}
}
return STATUS_BAD_PARAM;
}
plugin/batchedNMSPlugin/gatherNMSOutputs.h
+27
View File
@@ -0,0 +1,27 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_BATCHED_NMS_HELPER_H
#define TRT_BATCHED_NMS_HELPER_H
#include "plugin.h"
using namespace nvinfer1;
using namespace nvinfer1::plugin;
pluginStatus_t gatherNMSOutputs(cudaStream_t stream, bool shareLocation, int numImages, int numPredsPerClass,
int numClasses, int topK, int keepTopK, DataType DT_BBOX, DataType DT_SCORE, const void* indices,
const void* scores, const void* bboxData, void* keepCount, void* nmsedBoxes, void* nmsedScores, void* nmsedClasses,
bool clipBoxes = true);
#endif
plugin/common/CMakeLists.txt
+22
View File
@@ -0,0 +1,22 @@
#
# Copyright (c) 2019, 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.
#
file(GLOB SRCS *.cpp *.cu)
add_subdirectory(kernels)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
plugin/common/bboxUtils.h
+76
View File
@@ -0,0 +1,76 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_BBOX_UTILS_H
#define TRT_BBOX_UTILS_H
#include "plugin.h"
using namespace nvinfer1;
using namespace nvinfer1::plugin;
template <typename T>
struct Bbox
{
T xmin, ymin, xmax, ymax;
Bbox(T xmin, T ymin, T xmax, T ymax)
: xmin(xmin)
, ymin(ymin)
, xmax(xmax)
, ymax(ymax)
{
}
Bbox() = default;
};
template <typename T>
struct BboxInfo
{
T conf_score;
int label;
int bbox_idx;
bool kept;
BboxInfo(T conf_score, int label, int bbox_idx, bool kept)
: conf_score(conf_score)
, label(label)
, bbox_idx(bbox_idx)
, kept(kept)
{
}
BboxInfo() = default;
};
template <typename TFloat>
bool operator<(const Bbox<TFloat>& lhs, const Bbox<TFloat>& rhs)
{
return lhs.x1 < rhs.x1;
}
template <typename TFloat>
bool operator==(const Bbox<TFloat>& lhs, const Bbox<TFloat>& rhs)
{
return lhs.x1 == rhs.x1 && lhs.y1 == rhs.y1 && lhs.x2 == rhs.x2 && lhs.y2 == rhs.y2;
}
// }}}
int8_t* alignPtr(int8_t* ptr, uintptr_t to);
int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize);
size_t dataTypeSize(DataType dtype);
void setUniformOffsets(cudaStream_t stream, int num_segments, int offset, int* d_offsets);
#endif
plugin/common/cub_helper.h
+27
View File
@@ -0,0 +1,27 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
template <typename KeyT, typename ValueT>
size_t cubSortPairsWorkspaceSize(int num_items, int num_segments)
{
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending((void*) NULL, temp_storage_bytes, (const KeyT*) NULL,
(KeyT*) NULL, (const ValueT*) NULL, (ValueT*) NULL,
num_items, // # items
num_segments, // # segments
(const int*) NULL, (const int*) NULL);
return temp_storage_bytes;
}
plugin/common/kernels/CMakeLists.txt
+18
View File
@@ -0,0 +1,18 @@
#
# Copyright (c) 2019, 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.
#
file(GLOB SRCS *.cpp *.cu)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
plugin/common/kernels/allClassNMS.cu
+367
View File
@@ -0,0 +1,367 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "bboxUtils.h"
#include <vector>
template <typename T_BBOX>
__device__ T_BBOX bboxSize(
const Bbox<T_BBOX>& bbox,
const bool normalized)
{
if (bbox.xmax < bbox.xmin || bbox.ymax < bbox.ymin)
{
// If bbox is invalid (e.g. xmax < xmin or ymax < ymin), return 0.
return 0;
}
else
{
T_BBOX width = bbox.xmax - bbox.xmin;
T_BBOX height = bbox.ymax - bbox.ymin;
if (normalized)
{
return width * height;
}
else
{
// If bbox is not within range [0, 1].
return (width + 1) * (height + 1);
}
}
}
template <typename T_BBOX>
__device__ void intersectBbox(
const Bbox<T_BBOX>& bbox1,
const Bbox<T_BBOX>& bbox2,
Bbox<T_BBOX>* intersect_bbox)
{
if (bbox2.xmin > bbox1.xmax || bbox2.xmax < bbox1.xmin || bbox2.ymin > bbox1.ymax || bbox2.ymax < bbox1.ymin)
{
// Return [0, 0, 0, 0] if there is no intersection.
intersect_bbox->xmin = T_BBOX(0);
intersect_bbox->ymin = T_BBOX(0);
intersect_bbox->xmax = T_BBOX(0);
intersect_bbox->ymax = T_BBOX(0);
}
else
{
intersect_bbox->xmin = max(bbox1.xmin, bbox2.xmin);
intersect_bbox->ymin = max(bbox1.ymin, bbox2.ymin);
intersect_bbox->xmax = min(bbox1.xmax, bbox2.xmax);
intersect_bbox->ymax = min(bbox1.ymax, bbox2.ymax);
}
}
template <typename T_BBOX>
__device__ float jaccardOverlap(
const Bbox<T_BBOX>& bbox1,
const Bbox<T_BBOX>& bbox2,
const bool normalized)
{
Bbox<T_BBOX> intersect_bbox;
intersectBbox(bbox1, bbox2, &intersect_bbox);
float intersect_width, intersect_height;
if (normalized)
{
intersect_width = intersect_bbox.xmax - intersect_bbox.xmin;
intersect_height = intersect_bbox.ymax - intersect_bbox.ymin;
}
else
{
intersect_width = intersect_bbox.xmax - intersect_bbox.xmin + 1;
intersect_height = intersect_bbox.ymax - intersect_bbox.ymin + 1;
}
if (intersect_width > 0 && intersect_height > 0)
{
float intersect_size = intersect_width * intersect_height;
float bbox1_size = bboxSize(bbox1, normalized);
float bbox2_size = bboxSize(bbox2, normalized);
return intersect_size / (bbox1_size + bbox2_size - intersect_size);
}
else
{
return 0.;
}
}
template <typename T_BBOX>
__device__ void emptyBboxInfo(
BboxInfo<T_BBOX>* bbox_info)
{
bbox_info->conf_score = T_BBOX(0);
bbox_info->label = -2; // -1 is used for all labels when shared_location is ture
bbox_info->bbox_idx = -1;
bbox_info->kept = false;
}
/********** new NMS for only score and index array **********/
template <typename T_SCORE, typename T_BBOX, int TSIZE>
__global__ void allClassNMS_kernel(
const int num,
const int num_classes,
const int num_preds_per_class,
const int top_k,
const float nms_threshold,
const bool share_location,
const bool isNormalized,
T_BBOX* bbox_data, // bbox_data should be float to preserve location information
T_SCORE* beforeNMS_scores,
int* beforeNMS_index_array,
T_SCORE* afterNMS_scores,
int* afterNMS_index_array,
bool flipXY = false)
{
//__shared__ bool kept_bboxinfo_flag[CAFFE_CUDA_NUM_THREADS * TSIZE];
extern __shared__ bool kept_bboxinfo_flag[];
for (int i = 0; i < num; i++)
{
const int offset = i * num_classes * num_preds_per_class + blockIdx.x * num_preds_per_class;
const int max_idx = offset + top_k; // put top_k bboxes into NMS calculation
const int bbox_idx_offset = share_location ? (i * num_preds_per_class) : (i * num_classes * num_preds_per_class);
// local thread data
int loc_bboxIndex[TSIZE];
Bbox<T_BBOX> loc_bbox[TSIZE];
// initialize Bbox, Bboxinfo, kept_bboxinfo_flag
#pragma unroll
for (int t = 0; t < TSIZE; t++)
{
const int cur_idx = threadIdx.x + blockDim.x * t;
const int item_idx = offset + cur_idx;
if (item_idx < max_idx)
{
loc_bboxIndex[t] = beforeNMS_index_array[item_idx];
if (loc_bboxIndex[t] != -1)
{
const int bbox_data_idx = share_location ? (loc_bboxIndex[t] % num_preds_per_class + bbox_idx_offset) : loc_bboxIndex[t];
loc_bbox[t].xmin = flipXY ? bbox_data[bbox_data_idx * 4 + 1]
: bbox_data[bbox_data_idx * 4 + 0];
loc_bbox[t].ymin = flipXY ? bbox_data[bbox_data_idx * 4 + 0]
: bbox_data[bbox_data_idx * 4 + 1];
loc_bbox[t].xmax = flipXY ? bbox_data[bbox_data_idx * 4 + 3]
: bbox_data[bbox_data_idx * 4 + 2];
loc_bbox[t].ymax = flipXY ? bbox_data[bbox_data_idx * 4 + 2]
: bbox_data[bbox_data_idx * 4 + 3];
kept_bboxinfo_flag[cur_idx] = true;
}
else
{
kept_bboxinfo_flag[cur_idx] = false;
}
}
else
{
kept_bboxinfo_flag[cur_idx] = false;
}
}
// filter out overlapped boxes with lower scores
int ref_item_idx = offset;
int ref_bbox_idx = share_location ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset) : beforeNMS_index_array[ref_item_idx];
while ((ref_bbox_idx != -1) && ref_item_idx < max_idx)
{
Bbox<T_BBOX> ref_bbox;
ref_bbox.xmin = flipXY ? bbox_data[ref_bbox_idx * 4 + 1] : bbox_data[ref_bbox_idx * 4 + 0];
ref_bbox.ymin = flipXY ? bbox_data[ref_bbox_idx * 4 + 0] : bbox_data[ref_bbox_idx * 4 + 1];
ref_bbox.xmax = flipXY ? bbox_data[ref_bbox_idx * 4 + 3] : bbox_data[ref_bbox_idx * 4 + 2];
ref_bbox.ymax = flipXY ? bbox_data[ref_bbox_idx * 4 + 2] : bbox_data[ref_bbox_idx * 4 + 3];
for (int t = 0; t < TSIZE; t++)
{
const int cur_idx = threadIdx.x + blockDim.x * t;
const int item_idx = offset + cur_idx;
if ((kept_bboxinfo_flag[cur_idx]) && (item_idx > ref_item_idx))
{
// TODO: may need to add bool normalized as argument, HERE true means normalized
if (jaccardOverlap(ref_bbox, loc_bbox[t], isNormalized) > nms_threshold)
{
kept_bboxinfo_flag[cur_idx] = false;
}
}
}
__syncthreads();
do
{
ref_item_idx++;
} while (ref_item_idx < max_idx && !kept_bboxinfo_flag[ref_item_idx - offset]);
ref_bbox_idx = share_location ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset) : beforeNMS_index_array[ref_item_idx];
}
// store data
for (int t = 0; t < TSIZE; t++)
{
const int cur_idx = threadIdx.x + blockDim.x * t;
const int read_item_idx = offset + cur_idx;
const int write_item_idx = (i * num_classes * top_k + blockIdx.x * top_k) + cur_idx;
/*
* If not not keeping the bbox
* Set the score to 0
* Set the bounding box index to -1
*/
if (read_item_idx < max_idx)
{
afterNMS_scores[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? beforeNMS_scores[read_item_idx] : 0.0f;
afterNMS_index_array[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? loc_bboxIndex[t] : -1;
}
}
}
}
template <typename T_SCORE, typename T_BBOX>
pluginStatus_t allClassNMS_gpu(
cudaStream_t stream,
const int num,
const int num_classes,
const int num_preds_per_class,
const int top_k,
const float nms_threshold,
const bool share_location,
const bool isNormalized,
void* bbox_data,
void* beforeNMS_scores,
void* beforeNMS_index_array,
void* afterNMS_scores,
void* afterNMS_index_array,
bool flipXY = false)
{
#define P(tsize) allClassNMS_kernel<T_SCORE, T_BBOX, (tsize)>
void (*kernel[8])(const int, const int, const int, const int, const float,
const bool, const bool, float*, T_SCORE*, int*, T_SCORE*,
int*, bool)
= {
P(1), P(2), P(3), P(4), P(5), P(6), P(7), P(8),
};
const int BS = 512;
const int GS = num_classes;
const int t_size = (top_k + BS - 1) / BS;
kernel[t_size - 1]<<<GS, BS, BS * t_size * sizeof(bool), stream>>>(num, num_classes, num_preds_per_class,
top_k, nms_threshold, share_location, isNormalized,
(T_BBOX*) bbox_data,
(T_SCORE*) beforeNMS_scores,
(int*) beforeNMS_index_array,
(T_SCORE*) afterNMS_scores,
(int*) afterNMS_index_array,
flipXY);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// allClassNMS LAUNCH CONFIG
typedef pluginStatus_t (*nmsFunc)(cudaStream_t,
const int,
const int,
const int,
const int,
const float,
const bool,
const bool,
void*,
void*,
void*,
void*,
void*,
bool);
struct nmsLaunchConfigSSD
{
DataType t_score;
DataType t_bbox;
nmsFunc function;
nmsLaunchConfigSSD(DataType t_score, DataType t_bbox)
: t_score(t_score)
, t_bbox(t_bbox)
{
}
nmsLaunchConfigSSD(DataType t_score, DataType t_bbox, nmsFunc function)
: t_score(t_score)
, t_bbox(t_bbox)
, function(function)
{
}
bool operator==(const nmsLaunchConfigSSD& other)
{
return t_score == other.t_score && t_bbox == other.t_bbox;
}
};
static std::vector<nmsLaunchConfigSSD> nmsFuncVec;
bool nmsInit()
{
nmsFuncVec.push_back(nmsLaunchConfigSSD(DataType::kFLOAT, DataType::kFLOAT,
allClassNMS_gpu<float, float>));
return true;
}
static bool initialized = nmsInit();
pluginStatus_t allClassNMS(cudaStream_t stream,
const int num,
const int num_classes,
const int num_preds_per_class,
const int top_k,
const float nms_threshold,
const bool share_location,
const bool isNormalized,
const DataType DT_SCORE,
const DataType DT_BBOX,
void* bbox_data,
void* beforeNMS_scores,
void* beforeNMS_index_array,
void* afterNMS_scores,
void* afterNMS_index_array,
bool flipXY)
{
nmsLaunchConfigSSD lc = nmsLaunchConfigSSD(DT_SCORE, DT_BBOX, allClassNMS_gpu<float, float>);
for (unsigned i = 0; i < nmsFuncVec.size(); ++i)
{
if (lc == nmsFuncVec[i])
{
DEBUG_PRINTF("all class nms kernel %d\n", i);
return nmsFuncVec[i].function(stream,
num,
num_classes,
num_preds_per_class,
top_k,
nms_threshold,
share_location,
isNormalized,
bbox_data,
beforeNMS_scores,
beforeNMS_index_array,
afterNMS_scores,
afterNMS_index_array,
flipXY);
}
}
return STATUS_BAD_PARAM;
}
plugin/common/kernels/bboxDeltas2Proposals.cu
Executable
+292
View File
@@ -0,0 +1,292 @@
/*
* Copyright (c) 2019, 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.
*/
#include <algorithm>
#include <math.h>
#include <stdio.h>
#include <vector>
#include "kernel.h"
using std::max;
using std::min;
// BBD2P KERNEL
template <typename T_DELTAS,
DLayout_t L_DELTAS,
typename TV_PROPOSALS,
DLayout_t L_PROPOSALS,
typename T_FGSCORES,
DLayout_t L_FGSCORES>
__global__ void bboxDeltas2Proposals_kernel(
int N,
int A,
int H,
int W,
const float* __restrict__ anchors,
const float* __restrict__ imInfo,
int featureStride,
float minSize,
const T_DELTAS* __restrict__ deltas,
TV_PROPOSALS* __restrict__ proposals,
T_FGSCORES* __restrict__ scores)
{
int tid = threadIdx.x + blockIdx.x * blockDim.x;
if (tid < N * A * H * W)
{ // TODO this can be a loop.
// Find out the index of bounding box for the output
int cnt = tid;
// width index
int w = cnt % W;
cnt = cnt / W;
// height index
int h = cnt % H;
cnt = cnt / H;
// anchor box index
int a = cnt % A;
cnt = cnt / A;
// batch index
int n = cnt;
int hw = h * W + w;
// Get the height and width of the original input image
float imHeight = imInfo[3 * n];
float imWidth = imInfo[3 * n + 1];
// Point to the right anchor box
float4 anchor = ((float4*) anchors)[a];
// Get anchor box coordinates
float a_ctr_x = anchor.x;
float a_ctr_y = anchor.y;
float a_w = anchor.z;
float a_h = anchor.w;
// NCHW format
// Find out the starting position of the bounding box in the input (predicted bounding box offsets)
int id = ((tid - hw) * 4) + hw;
T_DELTAS dx;
T_DELTAS dy;
T_DELTAS dw;
T_DELTAS dh;
if (L_DELTAS == NCHW)
{
// The offsets between adjacent coordinates on linear memory is H * W
dx = deltas[id];
dy = deltas[id + 1 * H * W];
dw = deltas[id + 2 * H * W];
dh = deltas[id + 3 * H * W];
}
// NC4HW format
else if (L_DELTAS == NC4HW)
{
dx = deltas[tid * 4 + 0];
dy = deltas[tid * 4 + 1];
dw = deltas[tid * 4 + 2];
dh = deltas[tid * 4 + 3];
}
/*
* Calculate the coordinates of decoded bounding box on the original input image scale
* Only works if param.minBoxSize == param.featureStride
*/
float ctr_x = a_ctr_x + w * featureStride;
float ctr_y = a_ctr_y + h * featureStride;
// float ctr_x = (w + 0.5) * featureStride;
// float ctr_y = (h + 0.5) * featureStride;
/*
* Decode the predicted bounding box
* The decoded bounding boxes has coordinates of [x_topleft, y_topleft, x_bottomright, y_bottomright]
* The units are in pixels
*/
ctr_x = ctr_x + dx * a_w;
ctr_y = ctr_y + dy * a_h;
float b_w = __expf(dw) * a_w;
float b_h = __expf(dh) * a_h;
float bx = ctr_x - (b_w / 2);
float by = ctr_y - (b_h / 2);
float bz = ctr_x + (b_w / 2);
float bw = ctr_y + (b_h / 2);
TV_PROPOSALS bbox;
// Make sure that the decoded bouding box go outside of the original input image
bbox.x = fminf(fmaxf(bx, 0.0f), imWidth - 1.0f);
bbox.y = fminf(fmaxf(by, 0.0f), imHeight - 1.0f);
bbox.z = fminf(fmaxf(bz, 0.0f), imWidth - 1.0f);
bbox.w = fminf(fmaxf(bw, 0.0f), imHeight - 1.0f);
// Put the decoded bounding box information to the outputs
if (L_PROPOSALS == NC4HW)
{
proposals[tid] = bbox;
}
int ininf = 0xff800000;
float ninf = *(float*) &ininf;
// minBoxSize at the original input image scale
float scaledMinSize = minSize * imInfo[3 * n + 2];
// Set the objectness score to -inf if the predicted bounding box has edgth length less than the minimal box size expected.
if (bbox.z - bbox.x + 1 < scaledMinSize || bbox.w - bbox.y + 1 < scaledMinSize)
{
if (L_FGSCORES == NCHW)
scores[tid] = ninf;
}
}
}
// BBD2P KERNEL LAUNCHER
template <typename T_DELTAS,
DLayout_t L_DELTAS,
typename TV_PROPOSALS,
DLayout_t L_PROPOSALS,
typename T_FGSCORES,
DLayout_t L_FGSCORES>
pluginStatus_t bboxDeltas2Proposals_gpu(cudaStream_t stream,
int N,
int A,
int H,
int W,
const float* imInfo,
int featureStride,
float minBoxSize,
const float* anchors,
const void* deltas,
void* propos,
void* scores)
{
const int BS = 32;
const int GS = ((N * A * H * W) + BS - 1) / BS;
bboxDeltas2Proposals_kernel<T_DELTAS, L_DELTAS, TV_PROPOSALS, L_PROPOSALS, T_FGSCORES, L_FGSCORES><<<GS, BS, 0, stream>>>(N, A, H, W,
anchors,
imInfo,
featureStride,
minBoxSize,
(T_DELTAS*) deltas,
(TV_PROPOSALS*) propos,
(T_FGSCORES*) scores);
DEBUG_PRINTF("&&&& [bboxD2P] POST LAUNCH\n");
DEBUG_PRINTF("&&&& [bboxD2P] PROPOS %u\n", hash(propos, N * A * H * W * 4 * sizeof(float)));
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// BBD2P LAUNCH CONFIG {{{
typedef pluginStatus_t (*bd2pFun)(cudaStream_t,
int,
int,
int,
int,
const float*,
int,
float,
const float*,
const void*,
void*,
void*);
struct bd2pLaunchConfig
{
DataType t_deltas;
DLayout_t l_deltas;
DataType t_proposals;
DLayout_t l_proposals;
DataType t_scores;
DLayout_t l_scores;
bd2pFun function;
bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals, DataType t_scores, DLayout_t l_scores)
: t_deltas(t_deltas)
, l_deltas(l_deltas)
, t_proposals(t_proposals)
, l_proposals(l_proposals)
, t_scores(t_scores)
, l_scores(l_scores)
{
}
bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals, DataType t_scores, DLayout_t l_scores, bd2pFun function)
: t_deltas(t_deltas)
, l_deltas(l_deltas)
, t_proposals(t_proposals)
, l_proposals(l_proposals)
, t_scores(t_scores)
, l_scores(l_scores)
, function(function)
{
}
bool operator==(const bd2pLaunchConfig& other)
{
return t_deltas == other.t_deltas && l_deltas == other.l_deltas && t_proposals == other.t_proposals && l_proposals == other.l_proposals && t_scores == other.t_scores && l_scores == other.l_scores;
}
};
static std::vector<bd2pLaunchConfig> bd2pFunVec;
#define FLOAT32 nvinfer1::DataType::kFLOAT
bool init()
{
bd2pFunVec.push_back(bd2pLaunchConfig(FLOAT32, NC4HW,
FLOAT32, NC4HW,
FLOAT32, NCHW,
bboxDeltas2Proposals_gpu<float, NC4HW, float4, NC4HW, float, NCHW>));
bd2pFunVec.push_back(bd2pLaunchConfig(FLOAT32, NCHW,
FLOAT32, NC4HW,
FLOAT32, NCHW,
bboxDeltas2Proposals_gpu<float, NCHW, float4, NC4HW, float, NCHW>));
return true;
}
static bool initialized = init();
// BBD2P
pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
const int N,
const int A,
const int H,
const int W,
const int featureStride,
const float minBoxSize,
const float* imInfo,
const float* anchors,
const DataType t_deltas,
const DLayout_t l_deltas,
const void* deltas,
const DataType t_proposals,
const DLayout_t l_proposals,
void* proposals,
const DataType t_scores,
const DLayout_t l_scores,
void* scores)
{
bd2pLaunchConfig lc = bd2pLaunchConfig(t_deltas, l_deltas, t_proposals, l_proposals, t_scores, l_scores);
for (unsigned i = 0; i < bd2pFunVec.size(); i++)
{
if (lc == bd2pFunVec[i])
{
DEBUG_PRINTF("BBD2P kernel %d\n", i);
return bd2pFunVec[i].function(stream,
N, A, H, W,
imInfo,
featureStride,
minBoxSize,
anchors,
deltas,
proposals,
scores);
}
}
return STATUS_BAD_PARAM;
}
plugin/common/kernels/common.cu
Executable
+182
View File
@@ -0,0 +1,182 @@
/*
* Copyright (c) 2019, 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.
*/
#include "cublas_v2.h"
#include <cub/cub.cuh>
#include <stdint.h>
#include "kernel.h"
#include "bboxUtils.h"
#define CUDA_MEM_ALIGN 256
// HASH
unsigned int hash(const void* array_, size_t size)
{
// Apply hashing only when debugging RPN codes.
if (DEBUG_ENABLE)
{
const char* array_const;
char* array;
cudaMallocHost((void**) &array, size);
cudaMemcpy(array, array_, size, cudaMemcpyDeviceToHost);
array_const = array;
unsigned int hash = 45599;
for (size_t i = 0; i < size; i++)
{
unsigned int value = array_const[i];
hash = hash * 1487 + value;
hash = hash * 317;
hash = hash % 105359;
}
return hash;
}
else
{
return 0;
}
}
// ALIGNPTR
int8_t* alignPtr(int8_t* ptr, uintptr_t to)
{
uintptr_t addr = (uintptr_t) ptr;
if (addr % to)
{
addr += to - addr % to;
}
return (int8_t*) addr;
}
// NEXTWORKSPACEPTR
int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize)
{
uintptr_t addr = (uintptr_t) ptr;
addr += previousWorkspaceSize;
return alignPtr((int8_t*) addr, CUDA_MEM_ALIGN);
}
// CALCULATE TOTAL WORKSPACE SIZE
size_t calculateTotalWorkspaceSize(size_t* workspaces, int count)
{
size_t total = 0;
for (int i = 0; i < count; i++)
{
total += workspaces[i];
if (workspaces[i] % CUDA_MEM_ALIGN)
{
total += CUDA_MEM_ALIGN - (workspaces[i] % CUDA_MEM_ALIGN);
}
}
return total;
}
using nvinfer1::DataType;
// DATA TYPE SIZE
size_t dataTypeSize(const DataType dtype)
{
switch (dtype)
{
case DataType::kINT8: return sizeof(char);
case DataType::kHALF: return sizeof(short);
case DataType::kFLOAT: return sizeof(float);
default: return 0;
}
}
// CUB
/*
size_t cubSortFloatIntPairsWorkspaceSize(int num_items, int num_segments)
{
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending(
(int *)NULL, temp_storage_bytes,
(const float *)NULL, (float *)NULL,
(const int *)NULL, (int *)NULL,
num_items, // # items
num_segments, // # segments
(const int *)NULL, (const int *)NULL);
return temp_storage_bytes;
}
size_t cubSortFloatBboxInfoPairsWorkspaceSize(int num_items, int num_segments)
{
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending(
(int *)NULL, temp_storage_bytes,
(const float *)NULL, (float *)NULL,
(const BboxInfo<float> *)NULL, (BboxInfo<float> *)NULL,
num_items, // # items
num_segments, // # segments
(const int *)NULL, (const int *)NULL);
return temp_storage_bytes;
}
*/
template <unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void setUniformOffsets_kernel(
const int num_segments,
const int offset,
int* d_offsets)
{
const int idx = blockIdx.x * nthds_per_cta + threadIdx.x;
if (idx <= num_segments)
d_offsets[idx] = idx * offset;
}
void setUniformOffsets(
cudaStream_t stream,
const int num_segments,
const int offset,
int* d_offsets)
{
const int BS = 32;
const int GS = (num_segments + 1 + BS - 1) / BS;
setUniformOffsets_kernel<BS><<<GS, BS, 0, stream>>>(num_segments, offset, d_offsets);
}
const char* cublasGetErrorString(cublasStatus_t error)
{
switch (error)
{
case CUBLAS_STATUS_SUCCESS:
return "CUBLAS_STATUS_SUCCESS";
case CUBLAS_STATUS_NOT_INITIALIZED:
return "CUBLAS_STATUS_NOT_INITIALIZED";
case CUBLAS_STATUS_ALLOC_FAILED:
return "CUBLAS_STATUS_ALLOC_FAILED";
case CUBLAS_STATUS_INVALID_VALUE:
return "CUBLAS_STATUS_INVALID_VALUE";
case CUBLAS_STATUS_ARCH_MISMATCH:
return "CUBLAS_STATUS_ARCH_MISMATCH";
case CUBLAS_STATUS_MAPPING_ERROR:
return "CUBLAS_STATUS_MAPPING_ERROR";
case CUBLAS_STATUS_EXECUTION_FAILED:
return "CUBLAS_STATUS_EXECUTION_FAILED";
case CUBLAS_STATUS_INTERNAL_ERROR:
return "CUBLAS_STATUS_INTERNAL_ERROR";
#if CUDA_VERSION >= 6000
case CUBLAS_STATUS_NOT_SUPPORTED:
return "CUBLAS_STATUS_NOT_SUPPORTED";
#endif
#if CUDA_VERSION >= 6050
case CUBLAS_STATUS_LICENSE_ERROR:
return "CUBLAS_STATUS_LICENSE_ERROR";
#endif
}
return "Unknown cublas status";
}
plugin/common/kernels/cropAndResizeKernel.cu
Executable
+140
View File
@@ -0,0 +1,140 @@
/*
* Copyright (c) 2019, 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.
*/
#include <cuda.h>
#include <cuda_runtime.h>
template <typename T>
__global__ void cropAndResizeKernel(const int nthreads, const T* image_ptr, const float* boxes_ptr,
int num_boxes, int batch, int image_height, int image_width,
int crop_height, int crop_width, int depth,
float extrapolation_value, float* crops_ptr)
{
for (int out_idx = threadIdx.x + blockIdx.x * blockDim.x ; out_idx < nthreads;
out_idx += blockDim.x * gridDim.x)
{
int idx = out_idx;
const int x = idx % crop_width;
idx /= crop_width;
const int y = idx % crop_height;
idx /= crop_height;
const int d = idx % depth;
const int b = idx / depth;
const float y1 = boxes_ptr[b * 4];
const float x1 = boxes_ptr[b * 4 + 1];
const float y2 = boxes_ptr[b * 4 + 2];
const float x2 = boxes_ptr[b * 4 + 3];
//each image has num_boxes of boxes, so we simply divide to get the box index.
const int b_in = b / num_boxes;
if (b_in < 0 || b_in >= batch)
{
continue;
}
const float height_scale =
(crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1)
: 0;
const float width_scale =
(crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) : 0;
const float in_y = (crop_height > 1)
? y1 * (image_height - 1) + y * height_scale
: 0.5 * (y1 + y2) * (image_height - 1);
if (in_y < 0 || in_y > image_height - 1)
{
crops_ptr[out_idx] = extrapolation_value;
continue;
}
const float in_x = (crop_width > 1)
? x1 * (image_width - 1) + x * width_scale
: 0.5 * (x1 + x2) * (image_width - 1);
if (in_x < 0 || in_x > image_width - 1)
{
crops_ptr[out_idx] = extrapolation_value;
continue;
}
const int top_y_index = floorf(in_y);
const int bottom_y_index = ceilf(in_y);
const float y_lerp = in_y - top_y_index;
const int left_x_index = floorf(in_x);
const int right_x_index = ceilf(in_x);
const float x_lerp = in_x - left_x_index;
const float top_left(static_cast<float>(
image_ptr[((b_in * depth + d) * image_height +
top_y_index) *
image_width +
left_x_index]));
const float top_right(static_cast<float>(
image_ptr[((b_in * depth + d) * image_height +
top_y_index) *
image_width +
right_x_index]));
const float bottom_left(static_cast<float>(
image_ptr[((b_in * depth + d) * image_height +
bottom_y_index) *
image_width +
left_x_index]));
const float bottom_right(static_cast<float>(
image_ptr[((b_in * depth + d) * image_height +
bottom_y_index) *
image_width +
right_x_index]));
const float top = top_left + (top_right - top_left) * x_lerp;
const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp;
crops_ptr[out_idx] = top + (bottom - top) * y_lerp;
}
}
int cropAndResizeInference(
cudaStream_t stream,
int n,
const void* image,
const void* rois,
int batch_size,
int input_height,
int input_width,
int num_boxes,
int crop_height,
int crop_width,
int depth,
void* output)
{
int output_volume = batch_size * num_boxes * crop_height * crop_width * depth;
int block_size = 1024;
int grid_size = (output_volume + block_size - 1 ) / block_size;
cropAndResizeKernel<float> <<< grid_size, block_size, 0, stream>>>(output_volume,
static_cast<const float*>(image),
static_cast<const float*>(rois),
num_boxes,
batch_size,
input_height,
input_width,
crop_height,
crop_width,
depth,
0.0f,
static_cast<float*>(output));
return 0;
}
plugin/common/kernels/decodeBBoxes.cu
+328
View File
@@ -0,0 +1,328 @@
/*
* Copyright (c) 2019, 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.
*/
#include <vector>
#include "kernel.h"
template <typename T_BBOX, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void decodeBBoxes_kernel(
const int nthreads,
const CodeTypeSSD code_type,
const bool variance_encoded_in_target,
const int num_priors,
const bool share_location,
const int num_loc_classes,
const int background_label_id,
const bool clip_bbox,
const T_BBOX* loc_data,
const T_BBOX* prior_data,
T_BBOX* bbox_data)
{
for (int index = blockIdx.x * nthds_per_cta + threadIdx.x;
index < nthreads;
index += nthds_per_cta * gridDim.x)
{
// Bounding box coordinate index {0, 1, 2, 3}
const int i = index % 4;
// Bounding box class index
const int c = (index / 4) % num_loc_classes;
// Prior box id corresponding to the bounding box
const int d = (index / 4 / num_loc_classes) % num_priors;
// If bounding box was not shared among all the classes and the bounding box is corresponding to the background class
if (!share_location && c == background_label_id)
{
// Ignore background class if not share_location.
return;
}
// Index to the right anchor box corresponding to the current bounding box
const int pi = d * 4;
// Index to the right variances corresponding to the current bounding box
const int vi = pi + num_priors * 4;
// Encoding method: CodeTypeSSD::CORNER
//if (code_type == PriorBoxParameter_CodeType_CORNER){
if (code_type == CodeTypeSSD::CORNER)
{
// Do not want to use variances to adjust the bounding box decoding
if (variance_encoded_in_target)
{
// variance is encoded in target, we simply need to add the offset
// predictions.
// prior_data[pi + i]: prior box coordinates corresponding to the current bounding box coordinate
bbox_data[index] = prior_data[pi + i] + loc_data[index];
}
else
{
// variance is encoded in bbox, we need to scale the offset accordingly.
// prior_data[vi + i]: variance corresponding to the current bounding box coordinate
bbox_data[index] = prior_data[pi + i] + loc_data[index] * prior_data[vi + i];
}
//} else if (code_type == PriorBoxParameter_CodeType_CENTER_SIZE) {
}
// Encoding method: CodeTypeSSD::CENTER_SIZE
else if (code_type == CodeTypeSSD::CENTER_SIZE)
{
// Get prior box coordinates
const T_BBOX p_xmin = prior_data[pi];
const T_BBOX p_ymin = prior_data[pi + 1];
const T_BBOX p_xmax = prior_data[pi + 2];
const T_BBOX p_ymax = prior_data[pi + 3];
// Calculate prior box center, height, and width
const T_BBOX prior_width = p_xmax - p_xmin;
const T_BBOX prior_height = p_ymax - p_ymin;
const T_BBOX prior_center_x = (p_xmin + p_xmax) / 2.;
const T_BBOX prior_center_y = (p_ymin + p_ymax) / 2.;
// Get the current bounding box coordinates
const T_BBOX xmin = loc_data[index - i];
const T_BBOX ymin = loc_data[index - i + 1];
const T_BBOX xmax = loc_data[index - i + 2];
const T_BBOX ymax = loc_data[index - i + 3];
// Declare decoded bounding box coordinates
T_BBOX decode_bbox_center_x, decode_bbox_center_y;
T_BBOX decode_bbox_width, decode_bbox_height;
// Do not want to use variances to adjust the bounding box decoding
if (variance_encoded_in_target)
{
// variance is encoded in target, we simply need to retore the offset
// predictions.
decode_bbox_center_x = xmin * prior_width + prior_center_x;
decode_bbox_center_y = ymin * prior_height + prior_center_y;
decode_bbox_width = exp(xmax) * prior_width;
decode_bbox_height = exp(ymax) * prior_height;
}
else
{
// variance is encoded in bbox, we need to scale the offset accordingly.
decode_bbox_center_x = prior_data[vi] * xmin * prior_width + prior_center_x;
decode_bbox_center_y = prior_data[vi + 1] * ymin * prior_height + prior_center_y;
decode_bbox_width = exp(prior_data[vi + 2] * xmax) * prior_width;
decode_bbox_height = exp(prior_data[vi + 3] * ymax) * prior_height;
}
// Use [x_topleft, y_topleft, x_bottomright, y_bottomright] as coordinates for final decoded bounding box output
switch (i)
{
case 0:
bbox_data[index] = decode_bbox_center_x - decode_bbox_width / 2.;
break;
case 1:
bbox_data[index] = decode_bbox_center_y - decode_bbox_height / 2.;
break;
case 2:
bbox_data[index] = decode_bbox_center_x + decode_bbox_width / 2.;
break;
case 3:
bbox_data[index] = decode_bbox_center_y + decode_bbox_height / 2.;
break;
}
//} else if (code_type == PriorBoxParameter_CodeType_CORNER_SIZE) {
}
// Encoding method: CodeTypeSSD::CORNER_SIZE
else if (code_type == CodeTypeSSD::CORNER_SIZE)
{
// Get prior box coordinates
const T_BBOX p_xmin = prior_data[pi];
const T_BBOX p_ymin = prior_data[pi + 1];
const T_BBOX p_xmax = prior_data[pi + 2];
const T_BBOX p_ymax = prior_data[pi + 3];
// Get prior box width and height
const T_BBOX prior_width = p_xmax - p_xmin;
const T_BBOX prior_height = p_ymax - p_ymin;
T_BBOX p_size;
if (i == 0 || i == 2)
{
p_size = prior_width;
}
else
{
p_size = prior_height;
}
// Do not want to use variances to adjust the bounding box decoding
if (variance_encoded_in_target)
{
// variance is encoded in target, we simply need to add the offset
// predictions.
bbox_data[index] = prior_data[pi + i] + loc_data[index] * p_size;
}
else
{
// variance is encoded in bbox, we need to scale the offset accordingly.
bbox_data[index] = prior_data[pi + i] + loc_data[index] * prior_data[vi + i] * p_size;
}
}
// Exactly the same to CodeTypeSSD::CENTER_SIZE with using variance to adjust the bounding box decoding
else if (code_type == CodeTypeSSD::TF_CENTER)
{
const T_BBOX pXmin = prior_data[pi];
const T_BBOX pYmin = prior_data[pi + 1];
const T_BBOX pXmax = prior_data[pi + 2];
const T_BBOX pYmax = prior_data[pi + 3];
const T_BBOX priorWidth = pXmax - pXmin;
const T_BBOX priorHeight = pYmax - pYmin;
const T_BBOX priorCenterX = (pXmin + pXmax) / 2.;
const T_BBOX priorCenterY = (pYmin + pYmax) / 2.;
const T_BBOX ymin = loc_data[index - i];
const T_BBOX xmin = loc_data[index - i + 1];
const T_BBOX ymax = loc_data[index - i + 2];
const T_BBOX xmax = loc_data[index - i + 3];
T_BBOX bboxCenterX, bboxCenterY;
T_BBOX bboxWidth, bboxHeight;
bboxCenterX = prior_data[vi] * xmin * priorWidth + priorCenterX;
bboxCenterY = prior_data[vi + 1] * ymin * priorHeight + priorCenterY;
bboxWidth = exp(prior_data[vi + 2] * xmax) * priorWidth;
bboxHeight = exp(prior_data[vi + 3] * ymax) * priorHeight;
switch (i)
{
case 0:
bbox_data[index] = bboxCenterX - bboxWidth / 2.;
break;
case 1:
bbox_data[index] = bboxCenterY - bboxHeight / 2.;
break;
case 2:
bbox_data[index] = bboxCenterX + bboxWidth / 2.;
break;
case 3:
bbox_data[index] = bboxCenterY + bboxHeight / 2.;
break;
}
}
else
{
// Unknown code type.
assert("Unknown Box decode code type");
}
// Clip bounding box or not
if (clip_bbox)
{
bbox_data[index] = max(min(bbox_data[index], T_BBOX(1.)), T_BBOX(0.));
}
}
}
template <typename T_BBOX>
pluginStatus_t decodeBBoxes_gpu(
cudaStream_t stream,
const int nthreads,
const CodeTypeSSD code_type,
const bool variance_encoded_in_target,
const int num_priors,
const bool share_location,
const int num_loc_classes,
const int background_label_id,
const bool clip_bbox,
const void* loc_data,
const void* prior_data,
void* bbox_data)
{
const int BS = 512;
const int GS = (nthreads + BS - 1) / BS;
decodeBBoxes_kernel<T_BBOX, BS><<<GS, BS, 0, stream>>>(nthreads, code_type, variance_encoded_in_target,
num_priors, share_location, num_loc_classes,
background_label_id, clip_bbox,
(const T_BBOX*) loc_data, (const T_BBOX*) prior_data,
(T_BBOX*) bbox_data);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// decodeBBoxes LAUNCH CONFIG
typedef pluginStatus_t (*dbbFunc)(cudaStream_t,
const int,
const CodeTypeSSD,
const bool,
const int,
const bool,
const int,
const int,
const bool,
const void*,
const void*,
void*);
struct dbbLaunchConfig
{
DataType t_bbox;
dbbFunc function;
dbbLaunchConfig(DataType t_bbox)
: t_bbox(t_bbox)
{
}
dbbLaunchConfig(DataType t_bbox, dbbFunc function)
: t_bbox(t_bbox)
, function(function)
{
}
bool operator==(const dbbLaunchConfig& other)
{
return t_bbox == other.t_bbox;
}
};
static std::vector<dbbLaunchConfig> dbbFuncVec;
bool decodeBBoxesInit()
{
dbbFuncVec.push_back(dbbLaunchConfig(DataType::kFLOAT,
decodeBBoxes_gpu<float>));
return true;
}
static bool initialized = decodeBBoxesInit();
pluginStatus_t decodeBBoxes(
cudaStream_t stream,
const int nthreads,
const CodeTypeSSD code_type,
const bool variance_encoded_in_target,
const int num_priors,
const bool share_location,
const int num_loc_classes,
const int background_label_id,
const bool clip_bbox,
const DataType DT_BBOX,
const void* loc_data,
const void* prior_data,
void* bbox_data)
{
dbbLaunchConfig lc = dbbLaunchConfig(DT_BBOX);
for (unsigned i = 0; i < dbbFuncVec.size(); ++i)
{
if (lc == dbbFuncVec[i])
{
DEBUG_PRINTF("decodeBBox kernel %d\n", i);
return dbbFuncVec[i].function(stream,
nthreads,
code_type,
variance_encoded_in_target,
num_priors,
share_location,
num_loc_classes,
background_label_id,
clip_bbox,
loc_data,
prior_data,
bbox_data);
}
}
return STATUS_BAD_PARAM;
}
plugin/common/kernels/detectionForward.cu
+208
View File
@@ -0,0 +1,208 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "bboxUtils.h"
pluginStatus_t detectionInference(
cudaStream_t stream,
const int N,
const int C1,
const int C2,
const bool shareLocation,
const bool varianceEncodedInTarget,
const int backgroundLabelId,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const float confidenceThreshold,
const float nmsThreshold,
const CodeTypeSSD codeType,
const DataType DT_BBOX,
const void* locData,
const void* priorData,
const DataType DT_SCORE,
const void* confData,
void* keepCount,
void* topDetections,
void* workspace,
bool isNormalized,
bool confSigmoid)
{
// Batch size * number bbox per sample * 4 = total number of bounding boxes * 4
const int locCount = N * C1;
// Do not clip the bounding box that goes outside the image
const bool clipBBox = false;
/*
* shareLocation
* Bounding box are shared among all classes, i.e., a bounding box could be classified as any candidate class.
* Otherwise
* Bounding box are designed for specific classes, i.e., a bounding box could be classified as one certain class or not (binary classification).
*/
const int numLocClasses = shareLocation ? 1 : numClasses;
size_t bboxDataSize = detectionForwardBBoxDataSize(N, C1, DataType::kFLOAT);
void* bboxDataRaw = workspace;
pluginStatus_t status = decodeBBoxes(stream,
locCount,
codeType,
varianceEncodedInTarget,
numPredsPerClass,
shareLocation,
numLocClasses,
backgroundLabelId,
clipBBox,
DataType::kFLOAT,
locData,
priorData,
bboxDataRaw);
ASSERT_FAILURE(status == STATUS_SUCCESS);
/*
* bboxDataRaw format:
* [batch size, numPriors (per sample), numLocClasses, 4]
*/
// float for now
void* bboxData;
size_t bboxPermuteSize = detectionForwardBBoxPermuteSize(shareLocation, N, C1, DataType::kFLOAT);
void* bboxPermute = nextWorkspacePtr((int8_t*) bboxDataRaw, bboxDataSize);
/*
* After permutation, bboxData format:
* [batch_size, numLocClasses, numPriors (per sample) (numPredsPerClass), 4]
* This is equivalent to swapping axis
*/
if (!shareLocation)
{
status = permuteData(stream,
locCount,
numLocClasses,
numPredsPerClass,
4,
DataType::kFLOAT,
false,
bboxDataRaw,
bboxPermute);
ASSERT_FAILURE(status == STATUS_SUCCESS);
bboxData = bboxPermute;
}
/*
* If shareLocation, numLocClasses = 1
* No need to permute data on linear memory
*/
else
{
bboxData = bboxDataRaw;
}
/*
* Conf data format
* [batch size, numPriors * param.numClasses, 1, 1]
*/
const int numScores = N * C2;
size_t scoresSize = detectionForwardPreNMSSize(N, C2);
void* scores = nextWorkspacePtr((int8_t*) bboxPermute, bboxPermuteSize);
// need a conf_scores
/*
* After permutation, bboxData format:
* [batch_size, numClasses, numPredsPerClass, 1]
*/
status = permuteData(stream,
numScores,
numClasses,
numPredsPerClass,
1,
DataType::kFLOAT,
confSigmoid,
confData,
scores);
ASSERT_FAILURE(status == STATUS_SUCCESS);
size_t indicesSize = detectionForwardPreNMSSize(N, C2);
void* indices = nextWorkspacePtr((int8_t*) scores, scoresSize);
size_t postNMSScoresSize = detectionForwardPostNMSSize(N, numClasses, topK);
size_t postNMSIndicesSize = detectionForwardPostNMSSize(N, numClasses, topK);
void* postNMSScores = nextWorkspacePtr((int8_t*) indices, indicesSize);
void* postNMSIndices = nextWorkspacePtr((int8_t*) postNMSScores, postNMSScoresSize);
//size_t sortingWorkspaceSize = sortScoresPerClassWorkspaceSize(N, numClasses, numPredsPerClass, FLOAT32);
void* sortingWorkspace = nextWorkspacePtr((int8_t*) postNMSIndices, postNMSIndicesSize);
// Sort the scores so that the following NMS could be applied.
status = sortScoresPerClass(stream,
N,
numClasses,
numPredsPerClass,
backgroundLabelId,
confidenceThreshold,
DataType::kFLOAT,
scores,
indices,
sortingWorkspace);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// NMS
status = allClassNMS(stream,
N,
numClasses,
numPredsPerClass,
topK,
nmsThreshold,
shareLocation,
isNormalized,
DataType::kFLOAT,
DataType::kFLOAT,
bboxData,
scores,
indices,
postNMSScores,
postNMSIndices,
false);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// Sort the bounding boxes after NMS using scores
status = sortScoresPerImage(stream,
N,
numClasses * topK,
DataType::kFLOAT,
postNMSScores,
postNMSIndices,
scores,
indices,
sortingWorkspace);
ASSERT_FAILURE(status == STATUS_SUCCESS);
// Gather data from the sorted bounding boxes after NMS
status = gatherTopDetections(stream,
shareLocation,
N,
numPredsPerClass,
numClasses,
topK,
keepTopK,
DataType::kFLOAT,
DataType::kFLOAT,
indices,
scores,
bboxData,
keepCount,
topDetections);
ASSERT_FAILURE(status == STATUS_SUCCESS);
return STATUS_SUCCESS;
}
plugin/common/kernels/extractFgScores.cu
Executable
+81
View File
@@ -0,0 +1,81 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
template <typename T>
pluginStatus_t extractFgScores_gpu(cudaStream_t stream,
int N,
int A,
int H,
int W,
const void* scores,
void* fgScores)
{
// Copy all the objectness scores for one batch
size_t size = A * H * W * sizeof(T);
for (int n = 0; n < N; n++)
{
// Find out the starting pointer of the objectness scores in the input
size_t offset_ld = (n * 2 + 1) * A * H * W;
// Find out the starting pointer of the objectness scores in the output
size_t offset_st = n * A * H * W;
CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size, cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
}
return STATUS_SUCCESS;
}
template <typename T>
pluginStatus_t extractFgScores_cpu(int N,
int A,
int H,
int W,
const void* scores,
void* fgScores)
{
size_t size = A * H * W * sizeof(T);
for (int n = 0; n < N; n++)
{
size_t offset_ld = (n * 2 + 1) * A * H * W;
size_t offset_st = n * A * H * W;
memcpy(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size);
}
return STATUS_SUCCESS;
}
pluginStatus_t extractFgScores(cudaStream_t stream,
const int N,
const int A,
const int H,
const int W,
const DataType t_scores,
const DLayout_t l_scores,
const void* scores,
const DataType t_fgScores,
const DLayout_t l_fgScores,
void* fgScores)
{
if (l_fgScores != NCHW || l_scores != NCHW)
return STATUS_BAD_PARAM;
if (t_fgScores != DataType::kFLOAT)
return STATUS_BAD_PARAM;
if (t_scores != DataType::kFLOAT)
return STATUS_BAD_PARAM;
return extractFgScores_gpu<float>(stream, N, A, H, W, scores, fgScores);
}
plugin/common/kernels/gatherTopDetections.cu
+200
View File
@@ -0,0 +1,200 @@
/*
* Copyright (c) 2019, 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.
*/
#include <vector>
#include "plugin.h"
#include "kernel.h"
template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void gatherTopDetections_kernel(
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const int* indices,
const T_SCORE* scores,
const T_BBOX* bboxData,
int* keepCount,
T_BBOX* topDetections)
{
if (keepTopK > topK)
return;
for (int i = blockIdx.x * nthds_per_cta + threadIdx.x;
i < numImages * keepTopK;
i += gridDim.x * nthds_per_cta)
{
const int imgId = i / keepTopK;
const int detId = i % keepTopK;
const int offset = imgId * numClasses * topK;
const int index = indices[offset + detId];
const T_SCORE score = scores[offset + detId];
/*
* It is also likely that there is "bad bounding boxes" in the keepTopK bounding boxes.
* We set the bounding boxes parameters as the parameters shown below.
* These data will only show up at the end of keepTopK bounding boxes since the bounding boxes were sorted previously.
* It is also not going to affect the count of valid bounding boxes (keepCount).
* These data will probably never be used (because we have keepCount).
*/
if (index == -1)
{
topDetections[i * 7] = imgId; // image id
topDetections[i * 7 + 1] = -1; // label
topDetections[i * 7 + 2] = 0; // confidence score
// score==0 will not pass the VisualizeBBox check
topDetections[i * 7 + 3] = 0; // bbox xmin
topDetections[i * 7 + 4] = 0; // bbox ymin
topDetections[i * 7 + 5] = 0; // bbox xmax
topDetections[i * 7 + 6] = 0; // bbox ymax
}
else
{
const int bboxOffset = imgId * (shareLocation ? numPredsPerClass : (numClasses * numPredsPerClass));
const int bboxId = ((shareLocation ? (index % numPredsPerClass)
: index % (numClasses * numPredsPerClass)) + bboxOffset) * 4;
topDetections[i * 7] = imgId; // image id
topDetections[i * 7 + 1] = (index % (numClasses * numPredsPerClass)) / numPredsPerClass; // label
topDetections[i * 7 + 2] = score; // confidence score
// clipped bbox xmin
topDetections[i * 7 + 3] = max(min(bboxData[bboxId], T_BBOX(1.)), T_BBOX(0.));
// clipped bbox ymin
topDetections[i * 7 + 4] = max(min(bboxData[bboxId + 1], T_BBOX(1.)), T_BBOX(0.));
// clipped bbox xmax
topDetections[i * 7 + 5] = max(min(bboxData[bboxId + 2], T_BBOX(1.)), T_BBOX(0.));
// clipped bbox ymax
topDetections[i * 7 + 6] = max(min(bboxData[bboxId + 3], T_BBOX(1.)), T_BBOX(0.));
// Atomic add to increase the count of valid keepTopK bounding boxes
// Without having to do manual sync.
atomicAdd(&keepCount[i / keepTopK], 1);
}
}
}
template <typename T_BBOX, typename T_SCORE>
pluginStatus_t gatherTopDetections_gpu(
cudaStream_t stream,
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const void* indices,
const void* scores,
const void* bboxData,
void* keepCount,
void* topDetections)
{
cudaMemsetAsync(keepCount, 0, numImages * sizeof(int), stream);
const int BS = 32;
const int GS = 32;
gatherTopDetections_kernel<T_BBOX, T_SCORE, BS><<<GS, BS, 0, stream>>>(shareLocation, numImages, numPredsPerClass,
numClasses, topK, keepTopK,
(int*) indices, (T_SCORE*) scores, (T_BBOX*) bboxData,
(int*) keepCount, (T_BBOX*) topDetections);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// gatherTopDetections LAUNCH CONFIG
typedef pluginStatus_t (*gtdFunc)(cudaStream_t,
const bool,
const int,
const int,
const int,
const int,
const int,
const void*,
const void*,
const void*,
void*,
void*);
struct gtdLaunchConfig
{
DataType t_bbox;
DataType t_score;
gtdFunc function;
gtdLaunchConfig(DataType t_bbox, DataType t_score)
: t_bbox(t_bbox)
, t_score(t_score)
{
}
gtdLaunchConfig(DataType t_bbox, DataType t_score, gtdFunc function)
: t_bbox(t_bbox)
, t_score(t_score)
, function(function)
{
}
bool operator==(const gtdLaunchConfig& other)
{
return t_bbox == other.t_bbox && t_score == other.t_score;
}
};
using nvinfer1::DataType;
static std::vector<gtdLaunchConfig> gtdFuncVec;
bool gtdInit()
{
gtdFuncVec.push_back(gtdLaunchConfig(DataType::kFLOAT, DataType::kFLOAT,
gatherTopDetections_gpu<float, float>));
return true;
}
static bool initialized = gtdInit();
pluginStatus_t gatherTopDetections(
cudaStream_t stream,
const bool shareLocation,
const int numImages,
const int numPredsPerClass,
const int numClasses,
const int topK,
const int keepTopK,
const DataType DT_BBOX,
const DataType DT_SCORE,
const void* indices,
const void* scores,
const void* bboxData,
void* keepCount,
void* topDetections)
{
gtdLaunchConfig lc = gtdLaunchConfig(DT_BBOX, DT_SCORE);
for (unsigned i = 0; i < gtdFuncVec.size(); ++i)
{
if (lc == gtdFuncVec[i])
{
DEBUG_PRINTF("gatherTopDetections kernel %d\n", i);
return gtdFuncVec[i].function(stream,
shareLocation,
numImages,
numPredsPerClass,
numClasses,
topK,
keepTopK,
indices,
scores,
bboxData,
keepCount,
topDetections);
}
}
return STATUS_BAD_PARAM;
}
plugin/common/kernels/generateAnchors.cu
Executable
+100
View File
@@ -0,0 +1,100 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include <cstdio>
pluginStatus_t generateAnchors_cpu(int numRatios,
float* ratios,
int numScales,
float* scales,
int baseSize,
float* anchors)
{
#ifdef DEBUG
DEBUG_PRINTF("Generating Anchors with:\n");
DEBUG_PRINTF("Scales:");
for (int s = 0; s < numScales; ++s)
{
DEBUG_PRINTF("%f\t", scales[s]);
}
DEBUG_PRINTF("\n");
DEBUG_PRINTF("Ratios:");
for (int r = 0; r < numRatios; ++r)
{
DEBUG_PRINTF("%f\t", ratios[r]);
}
DEBUG_PRINTF("\n");
#endif
if ((numScales <= 0) || (numRatios <= 0) || (baseSize <= 0))
{
return STATUS_BAD_PARAM;
}
// Generate parameters for numRatios * numScales general anchor boxes
for (int r = 0; r < numRatios; ++r)
{
for (int s = 0; s < numScales; ++s)
{
int id = r * numScales + s;
float scale = scales[s];
float ratio = ratios[r];
float bs = baseSize;
float ws = round(sqrt((float) (bs * bs) / ratio));
float hs = round(ws * ratio);
// Width: bs / sqrt(ratio) * scale
// Height: bs * sqrt(ratio) * scale
ws *= scale;
hs *= scale;
// x_anchor_ctr
/*
* This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
* Because the center of anchor box in the original input raw image scale will not be dependent on this.
*/
anchors[id * 4] = (bs - 1) / 2;
// y_anchor_ctr
/*
* This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
* Because the center of anchor box in the original input raw image scale will not be dependent on this.
*/
anchors[id * 4 + 1] = (bs - 1) / 2;
// w_anchor
anchors[id * 4 + 2] = ws;
// h_anchor
anchors[id * 4 + 3] = hs;
}
}
return STATUS_SUCCESS;
}
pluginStatus_t generateAnchors(cudaStream_t stream,
int numRatios,
float* ratios,
int numScales,
float* scales,
int baseSize,
float* anchors)
{
// Each anchor box has 4 parameters
int ac = numRatios * numScales * 4;
float* anchors_cpu;
cudaMallocHost((void**) &anchors_cpu, sizeof(float) * ac);
pluginStatus_t status = generateAnchors_cpu(numRatios, ratios, numScales, scales, baseSize, anchors_cpu);
cudaMemcpyAsync(anchors, anchors_cpu, sizeof(float) * ac, cudaMemcpyHostToDevice, stream);
cudaFreeHost(anchors_cpu);
return status;
}
plugin/common/kernels/gridAnchorLayer.cu
+116
View File
@@ -0,0 +1,116 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "reducedMath.h"
#include <iostream>
using nvinfer1::rt::reduced_divisor;
template <unsigned nthdsPerCTA>
__launch_bounds__(nthdsPerCTA)
__global__ void gridAnchorKernel(
const GridAnchorParameters param,
const int numAspectRatios,
reduced_divisor divObj,
const float* widths,
const float* heights,
float* outputData
)
{
// output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
const int dim = param.H * param.W * numAspectRatios;
/*
* Parameters used to calculate the bounding box coordinates back to input image scale
* Normally we calculate the anchorStride = image_input_size (in pixel) / feature_map_size
* Here we do not use image_input_size for the moment
* Instead we use 1.0
* The coordinates calculated are scaled by the input image size.
* Most of the coordinates will be in a range of [0, 1], except for the bounding box coordinates going outside of the image
* Every coordinate will go back to the pixel coordinates in the input image if being multiplied by image_input_size
* Here we implicitly assumes the image input and feature map are square
*/
float anchorStride = (1.0 / param.H);
float anchorOffset = 0.5 * anchorStride;
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= dim)
return;
int arId, currIndex;
divObj.divmod(tid, currIndex, arId);
const int w = currIndex % param.W;
const int h = currIndex / param.W;
// Center coordinates
float yC = h * anchorStride + anchorOffset;
float xC = w * anchorStride + anchorOffset;
// x_min, y_min
float xMin = xC - 0.5 * widths[arId];
float yMin = yC - 0.5 * heights[arId];
// x_max, y_max
float xMax = xC + 0.5 * widths[arId];
float yMax = yC + 0.5 * heights[arId];
outputData[tid * 4] = xMin;
outputData[tid * 4 + 1] = yMin;
outputData[tid * 4 + 2] = xMax;
outputData[tid * 4 + 3] = yMax;
// Remember to move the output cursor
float* output = outputData + dim * 4;
// Simply copying the variance
output[tid * 4] = param.variance[0];
output[tid * 4 + 1] = param.variance[1];
output[tid * 4 + 2] = param.variance[2];
output[tid * 4 + 3] = param.variance[3];
}
pluginStatus_t anchorGridInference(
cudaStream_t stream,
const GridAnchorParameters param,
const int numAspectRatios,
const void* widths,
const void* heights,
void* outputData
)
{
const int dim = param.H * param.W * numAspectRatios;
reduced_divisor divObj(numAspectRatios);
if (dim > 5120)
{
const int BS = 128;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);
}
else
{
const int BS = 32;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);
}
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
plugin/common/kernels/kernel.cpp
+32
View File
@@ -0,0 +1,32 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "plugin.h"
size_t detectionInferenceWorkspaceSize(bool shareLocation, int N, int C1, int C2, int numClasses, int numPredsPerClass,
int topK, DataType DT_BBOX, DataType DT_SCORE)
{
size_t wss[7];
wss[0] = detectionForwardBBoxDataSize(N, C1, DT_BBOX);
wss[1] = detectionForwardBBoxPermuteSize(shareLocation, N, C1, DT_BBOX);
wss[2] = detectionForwardPreNMSSize(N, C2);
wss[3] = detectionForwardPreNMSSize(N, C2);
wss[4] = detectionForwardPostNMSSize(N, numClasses, topK);
wss[5] = detectionForwardPostNMSSize(N, numClasses, topK);
wss[6] = std::max(sortScoresPerClassWorkspaceSize(N, numClasses, numPredsPerClass, DT_SCORE),
sortScoresPerImageWorkspaceSize(N, numClasses * topK, DT_SCORE));
return calculateTotalWorkspaceSize(wss, 7);
}
plugin/common/kernels/kernel.h
+213
View File
@@ -0,0 +1,213 @@
/*
* Copyright (c) 2019, 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.
*/
#ifndef TRT_KERNEL_H
#define TRT_KERNEL_H
#include "cublas_v2.h"
#include "plugin.h"
#include <cassert>
#include <cstdio>
using namespace nvinfer1;
using namespace nvinfer1::plugin;
#define DEBUG_ENABLE 0
typedef enum
{
NCHW = 0,
NC4HW = 1
} DLayout_t;
pluginStatus_t allClassNMS(cudaStream_t stream, int num, int num_classes, int num_preds_per_class, int top_k,
float nms_threshold, bool share_location, bool isNormalized, DataType DT_SCORE, DataType DT_BBOX, void* bbox_data,
void* beforeNMS_scores, void* beforeNMS_index_array, void* afterNMS_scores, void* afterNMS_index_array,
bool flipXY = false);
pluginStatus_t detectionInference(cudaStream_t stream, int N, int C1, int C2, bool shareLocation,
bool varianceEncodedInTarget, int backgroundLabelId, int numPredsPerClass, int numClasses, int topK, int keepTopK,
float confidenceThreshold, float nmsThreshold, CodeTypeSSD codeType, DataType DT_BBOX, const void* locData,
const void* priorData, DataType DT_SCORE, const void* confData, void* keepCount, void* topDetections,
void* workspace, bool isNormalized = true, bool confSigmoid = false
);
pluginStatus_t gatherTopDetections(cudaStream_t stream, bool shareLocation, int numImages, int numPredsPerClass,
int numClasses, int topK, int keepTopK, DataType DT_BBOX, DataType DT_SCORE, const void* indices,
const void* scores, const void* bboxData, void* keepCount, void* topDetections);
size_t detectionForwardBBoxDataSize(int N, int C1, DataType DT_BBOX);
size_t detectionForwardBBoxPermuteSize(bool shareLocation, int N, int C1, DataType DT_BBOX);
size_t sortScoresPerClassWorkspaceSize(int num, int num_classes, int num_preds_per_class, DataType DT_CONF);
size_t sortScoresPerImageWorkspaceSize(int num_images, int num_items_per_image, DataType DT_SCORE);
pluginStatus_t sortScoresPerImage(cudaStream_t stream, int num_images, int num_items_per_image, DataType DT_SCORE,
void* unsorted_scores, void* unsorted_bbox_indices, void* sorted_scores, void* sorted_bbox_indices,
void* workspace);
pluginStatus_t sortScoresPerClass(cudaStream_t stream, int num, int num_classes, int num_preds_per_class,
int background_label_id, float confidence_threshold, DataType DT_SCORE, void* conf_scores_gpu,
void* index_array_gpu, void* workspace);
size_t calculateTotalWorkspaceSize(size_t* workspaces, int count);
const char* cublasGetErrorString(cublasStatus_t error);
pluginStatus_t permuteData(cudaStream_t stream, int nthreads, int num_classes, int num_data, int num_dim,
DataType DT_DATA, bool confSigmoid, const void* data, void* new_data);
size_t detectionForwardPreNMSSize(int N, int C2);
size_t detectionForwardPostNMSSize(int N, int numClasses, int topK);
pluginStatus_t decodeBBoxes(cudaStream_t stream, int nthreads, CodeTypeSSD code_type, bool variance_encoded_in_target,
int num_priors, bool share_location, int num_loc_classes, int background_label_id, bool clip_bbox, DataType DT_BBOX,
const void* loc_data, const void* prior_data, void* bbox_data);
size_t normalizePluginWorkspaceSize(bool acrossSpatial, int C, int H, int W);
pluginStatus_t normalizeInference(cudaStream_t stream, cublasHandle_t handle, bool acrossSpatial, bool channelShared,
int N, int C, int H, int W, float eps, const void* scale, const void* inputData, void* outputData, void* workspace);
pluginStatus_t priorBoxInference(cudaStream_t stream, PriorBoxParameters param, int H, int W, int numPriors,
int numAspectRatios, const void* minSize, const void* maxSize, const void* aspectRatios, void* outputData);
pluginStatus_t reorgInference(
cudaStream_t stream, int batch, int C, int H, int W, int stride, const void* input, void* output);
pluginStatus_t anchorGridInference(cudaStream_t stream, GridAnchorParameters param, int numAspectRatios,
const void* aspectRatios, const void* scales, void* outputData);
pluginStatus_t regionInference(cudaStream_t stream, int batch, int C, int H, int W, int num, int coords, int classes,
bool hasSoftmaxTree, const nvinfer1::plugin::softmaxTree* smTree, const void* input, void* output);
// GENERATE ANCHORS
// For now it takes host pointers - ratios and scales but
// in GPU MODE anchors should be device pointer
pluginStatus_t generateAnchors(cudaStream_t stream,
int numRatios, // number of ratios
float* ratios, // ratio array
int numScales, // number of scales
float* scales, // scale array
int baseSize, // size of the base anchor (baseSize x baseSize)
float* anchors); // output anchors (numRatios x numScales)
// BBD2P
pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
int N, // batch size
int A, // number of anchors
int H, // last feature map H
int W, // last feature map W
int featureStride, // feature stride
float minBoxSize, // minimum allowed box size before scaling
const float* imInfo, // image info (nrows, ncols, image scale)
const float* anchors, // input anchors
DataType tDeltas, // type of input deltas
DLayout_t lDeltas, // layout of input deltas
const void* deltas, // input deltas
DataType tProposals, // type of output proposals
DLayout_t lProposals, // layout of output proposals
void* proposals, // output proposals
DataType tScores, // type of output scores
DLayout_t lScores, // layout of output scores
void* scores); // output scores (the score associated with too small box will be set to -inf)
// NMS
pluginStatus_t nms(cudaStream_t stream,
int N, // batch size
int R, // number of ROIs (region of interest) per image
int preNmsTop, // number of proposals before applying NMS
int nmsMaxOut, // number of remaining proposals after applying NMS
float iouThreshold, // IoU threshold
DataType tFgScores, // type of foreground scores
DLayout_t lFgScores, // layout of foreground scores
void* fgScores, // foreground scores
DataType tProposals, // type of proposals
DLayout_t lProposals, // layout of proposals
const void* proposals, // proposals
void* workspace, // workspace
DataType tRois, // type of ROIs
void* rois); // ROIs
// WORKSPACE SIZES
size_t proposalsForwardNMSWorkspaceSize(int N, int A, int H, int W, int nmsMaxOut);
size_t proposalsForwardBboxWorkspaceSize(int N, int A, int H, int W);
size_t proposalForwardFgScoresWorkspaceSize(int N, int A, int H, int W);
size_t proposalsInferenceWorkspaceSize(int N, int A, int H, int W, int nmsMaxOut);
size_t RPROIInferenceFusedWorkspaceSize(int N, int A, int H, int W, int nmsMaxOut);
// PROPOSALS INFERENCE
pluginStatus_t proposalsInference(cudaStream_t stream, int N, int A, int H, int W, int featureStride, int preNmsTop,
int nmsMaxOut, float iouThreshold, float minBoxSize, const float* imInfo, const float* anchors, DataType tScores,
DLayout_t lScores, const void* scores, DataType tDeltas, DLayout_t lDeltas, const void* deltas, void* workspace,
DataType tRois, void* rois);
// EXTRACT FG SCORES
pluginStatus_t extractFgScores(cudaStream_t stream, int N, int A, int H, int W, DataType tScores, DLayout_t lScores,
const void* scores, DataType tFgScores, DLayout_t lFgScores, void* fgScores);
// ROI INFERENCE
pluginStatus_t roiInference(cudaStream_t stream,
const int R, // TOTAL number of rois -> ~nmsMaxOut * N
const int N, // Batch size
const int C, // Channels
const int H, // Input feature map H
const int W, // Input feature map W
const int poolingH, // Output feature map H
const int poolingW, // Output feature map W
const float spatialScale, const DataType tRois, const void* rois, const DataType tFeatureMap,
const DLayout_t lFeatureMap, const void* featureMap, const DataType tTop, const DLayout_t lTop, void* top);
// ROI FORWARD
pluginStatus_t roiForward(cudaStream_t stream,
int R, // TOTAL number of rois -> ~nmsMaxOut * N
int N, // Batch size
int C, // Channels
int H, // Input feature map H
int W, // Input feature map W
int poolingH, // Output feature map H
int poolingW, // Output feature map W
float spatialScale, DataType tRois, const void* rois, DataType tFeatureMap, DLayout_t lFeatureMap,
const void* featureMap, DataType tTop, DLayout_t lTop, void* top, int* maxIds);
// RP ROI Fused INFERENCE
pluginStatus_t RPROIInferenceFused(cudaStream_t stream, int N, int A, int C, int H, int W, int poolingH, int poolingW,
int featureStride, int preNmsTop, int nmsMaxOut, float iouThreshold, float minBoxSize, float spatialScale,
const float* imInfo, const float* anchors, DataType tScores, DLayout_t lScores, const void* scores,
DataType tDeltas, DLayout_t lDeltas, const void* deltas, DataType tFeatureMap, DLayout_t lFeatureMap,
const void* featureMap, void* workspace, DataType tRois, void* rois, DataType tTop, DLayout_t lTop, void* top);
// GENERATE ANCHORS CPU
pluginStatus_t generateAnchors_cpu(
int numRatios, float* ratios, int numScales, float* scales, int baseSize, float* anchors);
int cropAndResizeInference(cudaStream_t stream, int n, const void* image, const void* rois, int batch_size,
int input_height, int input_width, int num_boxes, int crop_height, int crop_width, int depth, void* output);
int proposalInference_gpu(cudaStream_t stream, const void* rpn_prob, const void* rpn_regr, int batch_size,
int input_height, int input_width, int rpn_height, int rpn_width, int MAX_BOX_NUM, int RPN_PRE_NMS_TOP_N,
float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float rpn_std_scaling,
int rpn_stride, float bbox_min_size, float nms_iou_threshold, void* workspace, void* output);
size_t _get_workspace_size(int N, int anc_size_num, int anc_ratio_num, int H, int W, int nmsMaxOut);
#endif
plugin/common/kernels/nmsLayer.cu
Executable
+445
View File
@@ -0,0 +1,445 @@
/*
* Copyright (c) 2019, 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.
*/
#include <algorithm>
#include "cuda_runtime_api.h"
#include <cub/cub.cuh>
#include <functional>
#include <stdint.h>
#include <stdio.h>
#include <vector>
#include "kernel.h"
#include "bboxUtils.h"
// CUB's bug workaround:
// To work properly for large batch size CUB segmented sort needs ridiculous
// workspace alignment.
const uintptr_t ALIGNMENT = 1 << 20;
// IOU
template <typename TFloat>
__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
{
TFloat left = max(a.xmin, b.xmin), right = min(a.xmax, b.xmax);
TFloat top = max(a.ymin, b.ymin), bottom = min(a.ymax, b.ymax);
TFloat width = max((TFloat)(right - left + (TFloat) 1.0), (TFloat) 0.0);
TFloat height = max((TFloat)(bottom - top + (TFloat) 1.0), (TFloat) 0.0);
TFloat interS = width * height;
TFloat Sa = (a.xmax - a.xmin + (TFloat) 1) * (a.ymax - a.ymin + (TFloat) 1);
TFloat Sb = (b.xmax - b.xmin + (TFloat) 1) * (b.ymax - b.ymin + (TFloat) 1);
return (float) interS / (float) (Sa + Sb - interS);
}
// NMS KERNEL FOR SMALL BATCH SIZE
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
__global__ __launch_bounds__(DIM) void nmsKernel1(const int propSize,
Bbox<T_PROPOSALS> const* __restrict__ preNmsProposals,
T_ROIS* __restrict__ afterNmsProposals,
const int preNmsTopN,
const float nmsThres,
const int afterNmsTopN)
{
__shared__ bool kept_boxes[TSIZE * DIM];
int kept = 0;
int batch_offset = blockIdx.x * propSize;
int max_box_idx = batch_offset + preNmsTopN;
int batch_offset_out = blockIdx.x * afterNmsTopN;
int flag_idx[TSIZE];
int boxes_idx[TSIZE];
Bbox<T_PROPOSALS> cur_boxes[TSIZE];
// initialize kept_boxes
#pragma unroll
for (int i = 0; i < TSIZE; i++)
{
boxes_idx[i] = threadIdx.x + batch_offset + DIM * i;
flag_idx[i] = threadIdx.x + DIM * i;
if (boxes_idx[i] < max_box_idx)
{
cur_boxes[i] = preNmsProposals[boxes_idx[i]];
kept_boxes[flag_idx[i]] = true;
}
else
{
kept_boxes[flag_idx[i]] = false;
boxes_idx[i] = -1.0f;
flag_idx[i] = -1.0f;
}
}
int ref_box_idx = 0 + batch_offset;
// remove the overlapped boxes
while ((kept < afterNmsTopN) && (ref_box_idx < max_box_idx))
{
Bbox<T_PROPOSALS> ref_box;
ref_box = preNmsProposals[ref_box_idx];
#pragma unroll
for (int i = 0; i < TSIZE; i++)
{
if (boxes_idx[i] > ref_box_idx)
{
if (IoU(ref_box, cur_boxes[i]) > nmsThres)
{
kept_boxes[flag_idx[i]] = false;
}
}
else if (boxes_idx[i] == ref_box_idx)
{
afterNmsProposals[(batch_offset_out + kept) * 4 + 0] = ref_box.xmin;
afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.ymin;
afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.xmax;
afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.ymax;
}
}
__syncthreads();
do
{
ref_box_idx++;
} while (!kept_boxes[ref_box_idx - batch_offset] && ref_box_idx < max_box_idx);
kept++;
}
}
// NMS KERNEL FOR LARGE BATCH SIZE
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
__global__ __launch_bounds__(DIM) void nmsKernel2(const int propSize,
Bbox<T_PROPOSALS> const* __restrict__ proposals,
T_ROIS* __restrict__ filtered,
const int preNmsTopN,
const float nmsThres,
const int afterNmsTopN)
{
Bbox<T_PROPOSALS> const* cProposals = proposals + blockIdx.x * propSize;
Bbox<T_PROPOSALS> t[TSIZE];
uint64_t del = 0;
for (int i = 0; i < TSIZE; i++)
{
if (i < TSIZE - 1 || i * DIM + threadIdx.x < preNmsTopN)
{
t[i] = cProposals[i * DIM + threadIdx.x];
}
}
__shared__ Bbox<T_PROPOSALS> last;
__shared__ bool kept;
__shared__ int foundBatch;
if (threadIdx.x == 0)
foundBatch = 0;
for (int i = 0; i < TSIZE; i++)
{
for (int j = 0; j < DIM; j++)
{
int offset = i * DIM;
int index = offset + j;
if (index >= preNmsTopN)
break;
__syncthreads();
if (threadIdx.x == j)
{
kept = 0 == (del & ((uint64_t) 1 << i));
last = t[i];
if (kept)
{
int cnt = blockIdx.x * afterNmsTopN + foundBatch;
filtered[cnt * 4 + 0] = t[i].xmin;
filtered[cnt * 4 + 1] = t[i].ymin;
filtered[cnt * 4 + 2] = t[i].xmax;
filtered[cnt * 4 + 3] = t[i].ymax;
foundBatch++;
}
}
__syncthreads();
if (foundBatch == afterNmsTopN)
{
return;
}
if (kept)
{
Bbox<T_PROPOSALS> test = last;
for (int k = 0; k < TSIZE; k++)
{
if (index < k * DIM + threadIdx.x
&& IoU<T_PROPOSALS>(test, t[k]) > nmsThres)
{
del |= (uint64_t) 1 << k;
}
}
}
}
}
}
// NMS LAUNCH
template <typename T_PROPOSALS, DLayout_t L_PROPOSALS, typename T_ROIS>
pluginStatus_t nmsLaunch(cudaStream_t stream,
const int batch,
const int propSize,
void* proposals,
void* filtered,
const int preNmsTopN,
const float nmsThres,
const int afterNmsTopN)
{
const int blockSize = 1024;
#define P1(tsize) nmsKernel1<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
#define P2(tsize) nmsKernel2<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
void (*kernel[64])(int, Bbox<T_PROPOSALS> const*, T_ROIS*, int, float, int) = {
P1(1), P1(2), P1(3), P1(4), P1(5), P1(6), P1(7), P1(8), P1(9), P1(10), P1(11), P1(12), P2(13), P2(14), P2(15), P2(16),
P2(17), P2(18), P2(19), P2(20), P2(21), P2(22), P2(23), P2(24), P2(25), P2(26), P2(27), P2(28), P2(29), P2(30), P2(31), P2(32),
P2(33), P2(34), P2(35), P2(36), P2(37), P2(38), P2(39), P2(40), P2(41), P2(42), P2(43), P2(44), P2(45), P2(46), P2(47), P2(48),
P2(49), P2(50), P2(51), P2(52), P2(53), P2(54), P2(55), P2(56), P2(57), P2(58), P2(59), P2(60), P2(61), P2(62), P2(63), P2(64)};
ASSERT_PARAM(preNmsTopN < 64 * blockSize);
CSC(cudaMemsetAsync(filtered, 0, batch * afterNmsTopN * 4 * sizeof(T_ROIS), stream), STATUS_FAILURE);
kernel[(preNmsTopN + blockSize - 1) / blockSize - 1]<<<batch, blockSize, 0, stream>>>(propSize,
(Bbox<T_PROPOSALS>*) proposals,
(T_ROIS*) filtered,
preNmsTopN,
nmsThres,
afterNmsTopN);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// SET OFFSET
// Works for up to 2Gi elements (cub's limitation)!
__global__ void setOffset(int stride, int size, int* output)
{
// One block, because batch size shouldn't be too large.
for (int i = threadIdx.x; i < size; i += blockDim.x)
{
output[i] = i * stride;
}
}
// NMS GPU
template <typename T_SCORES, typename T_ROIS>
pluginStatus_t nmsGpu(cudaStream_t stream,
const int N,
const int R,
const int preNmsTop,
const int nmsMaxOut,
const float iouThreshold,
//const float minBoxSize,
//const float * imInfo,
void* fgScores,
const void* proposals,
void* workspace,
void* rois)
{
int8_t* vworkspace = alignPtr((int8_t*) workspace, ALIGNMENT);
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
pluginStatus_t error;
DEBUG_PRINTF("&&&& [NMS] DISCARD\n");
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
// Generate offsets
int* offsets = (int*) vworkspace;
setOffset<<<1, 1024, 0, stream>>>(R, N + 1, offsets);
CSC(cudaGetLastError(), STATUS_FAILURE);
vworkspace = vworkspace + N + 1;
vworkspace = alignPtr(vworkspace, ALIGNMENT);
// Sort (batched)
std::size_t tempStorageBytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending(
NULL, tempStorageBytes,
(T_SCORES*) fgScores, (T_SCORES*) fgScores,
(Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposals,
N * R, N,
offsets, offsets + 1, 0, 8 * sizeof(T_SCORES), stream);
CSC(cudaGetLastError(), STATUS_FAILURE);
T_SCORES* scoresOut = (T_SCORES*) vworkspace;
vworkspace = (int8_t*) (scoresOut + N * R);
vworkspace = alignPtr(vworkspace, ALIGNMENT);
Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
vworkspace = (int8_t*) (proposalsOut + N * R);
vworkspace = alignPtr(vworkspace, ALIGNMENT);
cub::DeviceSegmentedRadixSort::SortPairsDescending(
vworkspace, tempStorageBytes,
(T_SCORES*) fgScores, (T_SCORES*) scoresOut,
(Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposalsOut,
N * R, N,
offsets, offsets + 1,
0, 8 * sizeof(T_SCORES), stream);
CSC(cudaGetLastError(), STATUS_FAILURE);
DEBUG_PRINTF("&&&& [NMS] POST CUB\n");
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposalsOut, N * R * 4 * sizeof(float)));
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(scoresOut, N * R * sizeof(float)));
error = nmsLaunch<T_ROIS, NC4HW, T_ROIS>(stream,
N,
R,
proposalsOut,
rois,
preNmsTop,
iouThreshold,
nmsMaxOut);
DEBUG_PRINTF("&&&& [NMS] POST LAUNCH\n");
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(rois, N * nmsMaxOut * 4 * sizeof(float)));
if (error != STATUS_SUCCESS)
{
return error;
}
return STATUS_SUCCESS;
}
// NMS LAUNCH CONFIG
typedef pluginStatus_t (*nmsFun)(cudaStream_t,
const int, // N
const int, // R
const int, // preNmsTop
const int, // nmsMaxOut
const float, // iouThreshold
//const float, // minBoxSize
//const float *, // imInfo
void*, // fgScores
const void*, // proposals,
void*, // workspace,
void*); // rois
struct nmsLaunchConfig
{
DataType t_fgScores;
DLayout_t l_fgScores;
DataType t_proposals;
DLayout_t l_proposals;
DataType t_rois;
nmsFun function;
nmsLaunchConfig(DataType t_fgScores,
DLayout_t l_fgScores,
DataType t_proposals,
DLayout_t l_proposals,
DataType t_rois,
nmsFun function)
: t_fgScores(t_fgScores)
, l_fgScores(l_fgScores)
, t_proposals(t_proposals)
, l_proposals(l_proposals)
, t_rois(t_rois)
, function(function)
{
}
nmsLaunchConfig(DataType t_fgScores,
DLayout_t l_fgScores,
DataType t_proposals,
DLayout_t l_proposals,
DataType t_rois)
: t_fgScores(t_fgScores)
, l_fgScores(l_fgScores)
, t_proposals(t_proposals)
, l_proposals(l_proposals)
, t_rois(t_rois)
{
}
bool operator==(const nmsLaunchConfig& other)
{
return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores) && (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals) && (t_rois == other.t_rois);
}
};
static std::vector<nmsLaunchConfig> nmsLCVec;
#define FLOAT32 nvinfer1::DataType::kFLOAT
bool initNmsLC()
{
nmsLCVec.reserve(1);
nmsLCVec.push_back(nmsLaunchConfig(FLOAT32, NCHW,
FLOAT32, NC4HW,
FLOAT32,
nmsGpu<float, float>));
return true;
}
static bool initializedNmsLC = initNmsLC();
// NMS
pluginStatus_t nms(cudaStream_t stream,
const int N,
const int R,
const int preNmsTop,
const int nmsMaxOut,
const float iouThreshold,
const DataType t_fgScores,
const DLayout_t l_fgScores,
void* fgScores,
const DataType t_proposals,
const DLayout_t l_proposals,
const void* proposals,
void* workspace,
const DataType t_rois,
void* rois)
{
if (!initializedNmsLC)
return STATUS_NOT_INITIALIZED;
nmsLaunchConfig lc(t_fgScores, l_fgScores, t_proposals, l_proposals, t_rois);
for (unsigned i = 0; i < nmsLCVec.size(); i++)
{
if (nmsLCVec[i] == lc)
{
DEBUG_PRINTF("NMS KERNEL %d\n", i);
return nmsLCVec[i].function(stream,
N, R,
preNmsTop,
nmsMaxOut,
iouThreshold,
fgScores,
proposals,
workspace,
rois);
}
}
return STATUS_BAD_PARAM;
}
plugin/common/kernels/normalizeLayer.cu
+219
View File
@@ -0,0 +1,219 @@
/*
* Copyright (c) 2019, 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.
*/
#include "kernel.h"
#include "bboxUtils.h"
#define CUBLAS_CHECK(condition) \
do \
{ \
cublasStatus_t status = condition; \
if (status != CUBLAS_STATUS_SUCCESS) \
{ \
printf("%s %d CUBLAS FAIL %s\n", __FILE__, __LINE__, cublasGetErrorString(status)); \
} \
} while (0)
size_t normalizePluginWorkspaceSize(bool acrossSpatial, int C, int H, int W)
{
if (acrossSpatial)
return sizeof(float) * C * H * W;
else
return (size_t) 0;
}
template <unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void normalizeNotAcrossSpatialKernel(
const bool channelShared,
const int N,
const int C,
const int H,
const int W,
const float eps,
const float* scale,
float* inputData,
float* outputData)
{
const int dim = C * H * W;
const int spatialDim = H * W;
const int tile = 32;
const int numTile = (spatialDim + tile - 1) / tile;
for (int n = blockIdx.x; n < N * numTile; n += gridDim.x)
{
float* input = inputData + (n / numTile) * dim;
float* output = outputData + (n / numTile) * dim;
__shared__ float sum[tile];
float localsum = 0.0F;
for (int i = threadIdx.x; i < tile; i += nthds_per_cta)
{
sum[i] = 0.0F;
}
__syncthreads();
for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
{
int row = i / tile;
int col = (n % numTile) * tile + i % tile;
float data = 0.0F;
if (col < spatialDim)
data = input[row * spatialDim + col];
localsum += data * data;
}
atomicAdd(&sum[threadIdx.x & 31], localsum);
__syncthreads();
for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
{
int row = i / tile;
int col = (n % numTile) * tile + i % tile;
if (col < spatialDim)
{
int offset = row * spatialDim + col;
output[offset] = input[offset] / sqrt(sum[threadIdx.x & 31] + eps);
}
}
if (channelShared)
{
for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
{
int row = i / tile;
int col = (n % numTile) * tile + i % tile;
if (col < spatialDim)
output[row * spatialDim + col] *= scale[0];
}
}
else
{
for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
{
int row = i / tile;
int col = (n % numTile) * tile + i % tile;
if (col < spatialDim)
output[row * spatialDim + col] *= scale[row];
}
}
}
}
pluginStatus_t normalizeNotAcrossSpatialGpu(
cudaStream_t stream,
const bool channelShared,
const int N,
const int C,
const int H,
const int W,
const float eps,
const void* scale,
const void* inputData,
void* outputData)
{
const int BS = 128;
const int GS = 256;
// assumes warp size == 32
ASSERT(BS % 32 == 0);
normalizeNotAcrossSpatialKernel<BS><<<GS, BS, 0, stream>>>(channelShared, N, C, H, W, eps,
(const float*) scale,
(float*) inputData,
(float*) outputData);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
__global__ void squareKernel(
const int n,
const float* x,
float* y)
{
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
i < n; i += gridDim.x * blockDim.x)
{
y[i] = x[i] * x[i];
}
}
__global__ void scalChannelKernel(
const int n,
const int spatialDim,
const float* inputData,
const float* scale,
float* outputData)
{
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
i < n; i += gridDim.x * blockDim.x)
{
// scale factors are indepedent across different channels
// scale[i / spatialDim]: find the right scale factor for specific channels
outputData[i] = inputData[i] / scale[i / spatialDim];
}
}
pluginStatus_t normalizeInference(
cudaStream_t stream,
cublasHandle_t handle,
const bool acrossSpatial,
const bool channelShared,
const int N,
const int C,
const int H,
const int W,
const float eps,
const void* scale,
const void* inputData,
void* outputData,
void* workspace)
{
const int dim = C * H * W;
// Normalization is conducted for each sample from the batch indepdently
if (acrossSpatial)
{
float* input = (float*) const_cast<void*>(inputData);
float* output = (float*) outputData;
float* buffer = (float*) workspace;
for (int n = 0; n < N; ++n)
{
// Take the square of each element in the input
squareKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, input, buffer);
float normsqr = 0.0F;
// Sum up all the squared elements
CUBLAS_CHECK(cublasSasum(handle, dim, buffer, 1, &normsqr));
// Make a copy of the input to the output
CUBLAS_CHECK(cublasScopy(handle, dim, input, 1, output, 1));
// Calculate the inverse of the square root of the sum
// Use eps to prevent being divided by zero
normsqr = 1 / sqrt(normsqr + eps);
// Scale all the outputs by normsqr
CUBLAS_CHECK(cublasSscal(handle, dim, &normsqr, output, 1));
// If channel shared is true, scale all the outputs
if (channelShared)
{
CUBLAS_CHECK(cublasSscal(handle, dim, (float*) scale, output, 1));
}
// Use different scale factors for different channels
else
{
// scale the output according to channels
scalChannelKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, H * W, output, (float*) scale, output);
}
// Move cursors
input += dim;
output += dim;
}
return STATUS_SUCCESS;
}
// Normalization ignoring the batch
else
{
return normalizeNotAcrossSpatialGpu(stream, channelShared, N, C, H, W, eps, scale, inputData, outputData);
}
}
plugin/common/kernels/permuteData.cu
+138
View File
@@ -0,0 +1,138 @@
/*
* Copyright (c) 2019, 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.
*/
#include <vector>
#include "kernel.h"
template <typename Dtype, unsigned nthds_per_cta>
__launch_bounds__(nthds_per_cta)
__global__ void permuteData_kernel(
const int nthreads,
const int num_classes,
const int num_data,
const int num_dim,
bool confSigmoid,
const Dtype* data,
Dtype* new_data)
{
// data format: [batch_size, num_data, num_classes, num_dim]
for (int index = blockIdx.x * nthds_per_cta + threadIdx.x;
index < nthreads;
index += nthds_per_cta * gridDim.x)
{
const int i = index % num_dim;
const int c = (index / num_dim) % num_classes;
const int d = (index / num_dim / num_classes) % num_data;
const int n = index / num_dim / num_classes / num_data;
const int new_index = ((n * num_classes + c) * num_data + d) * num_dim + i;
float result = data[index];
if (confSigmoid)
result = exp(result) / (1 + exp(result));
new_data[new_index] = result;
}
// new data format: [batch_size, num_classes, num_data, num_dim]
}
template <typename Dtype>
pluginStatus_t permuteData_gpu(
cudaStream_t stream,
const int nthreads,
const int num_classes,
const int num_data,
const int num_dim,
bool confSigmoid,
const void* data,
void* new_data)
{
const int BS = 512;
const int GS = (nthreads + BS - 1) / BS;
permuteData_kernel<Dtype, BS><<<GS, BS, 0, stream>>>(nthreads, num_classes, num_data, num_dim, confSigmoid,
(const Dtype*) data, (Dtype*) new_data);
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}
// permuteData LAUNCH CONFIG
typedef pluginStatus_t (*pdFunc)(cudaStream_t,
const int,
const int,
const int,
const int,
bool,
const void*,
void*);
struct pdLaunchConfig
{
DataType t_data;
pdFunc function;
pdLaunchConfig(DataType t_data)
: t_data(t_data)
{
}
pdLaunchConfig(DataType t_data, pdFunc function)
: t_data(t_data)
, function(function)
{
}
bool operator==(const pdLaunchConfig& other)
{
return t_data == other.t_data;
}
};
static std::vector<pdLaunchConfig> pdFuncVec;
bool permuteDataInit()
{
pdFuncVec.push_back(pdLaunchConfig(DataType::kFLOAT,
permuteData_gpu<float>));
return true;
}
static bool initialized = permuteDataInit();
pluginStatus_t permuteData(cudaStream_t stream,
const int nthreads,
const int num_classes,
const int num_data,
const int num_dim,
const DataType DT_DATA,
bool confSigmoid,
const void* data,
void* new_data)
{
pdLaunchConfig lc = pdLaunchConfig(DT_DATA);
for (unsigned i = 0; i < pdFuncVec.size(); ++i)
{
if (lc == pdFuncVec[i])
{
DEBUG_PRINTF("permuteData kernel %d\n", i);
return pdFuncVec[i].function(stream,
nthreads,
num_classes,
num_data,
num_dim,
confSigmoid,
data,
new_data);
}
}
return STATUS_BAD_PARAM;
}

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