haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

merge into: haotian:master
Branches Tags
haotian:master
haotian:yolov8
...
pull from: haotian:yolov8
Branches Tags
haotian:yolov8
haotian:master

No commits in common. "master" and "yolov8" have entirely different histories.
master ... yolov8

623 changed files with 119527 additions and 1233 deletions
Show Stats Expand all files Collapse all files

82
.clang-format Normal file
View File

@ -0,0 +1,82 @@
# Google C/C++ Code Style settings (with 4-space)
# Refered to https://github.com/kehanXue/google-style-clang-format/blob/master/.clang-format
Language: Cpp
BasedOnStyle: Google
AccessModifierOffset: -1
AlignAfterOpenBracket: Align
AlignConsecutiveAssignments: None
AlignOperands: Align
AllowAllArgumentsOnNextLine: true
AllowAllConstructorInitializersOnNextLine: true
AllowAllParametersOfDeclarationOnNextLine: false
AllowShortBlocksOnASingleLine: Empty
AllowShortCaseLabelsOnASingleLine: false
AllowShortFunctionsOnASingleLine: Inline
AllowShortIfStatementsOnASingleLine: Never # To avoid conflict, set this "Never" and each "if statement" should include brace when coding
AllowShortLambdasOnASingleLine: Inline
AllowShortLoopsOnASingleLine: false
AlwaysBreakAfterReturnType: None
AlwaysBreakTemplateDeclarations: Yes
BinPackArguments: true
BreakBeforeBraces: Custom
BraceWrapping:
AfterCaseLabel: false
AfterClass: false
AfterStruct: false
AfterControlStatement: Never
AfterEnum: false
AfterFunction: false
AfterNamespace: false
AfterUnion: false
AfterExternBlock: false
BeforeCatch: false
BeforeElse: false
BeforeLambdaBody: false
IndentBraces: false
SplitEmptyFunction: false
SplitEmptyRecord: false
SplitEmptyNamespace: false
BreakBeforeBinaryOperators: None
BreakBeforeTernaryOperators: true
BreakConstructorInitializers: BeforeColon
BreakInheritanceList: BeforeColon
ColumnLimit: 120
CompactNamespaces: false
ContinuationIndentWidth: 8
Cpp11BracedListStyle: true
DerivePointerAlignment: false # Make sure the * or & align on the left
EmptyLineBeforeAccessModifier: LogicalBlock
FixNamespaceComments: true
IncludeBlocks: Preserve
IndentCaseLabels: true
IndentPPDirectives: None
IndentWidth: 4
KeepEmptyLinesAtTheStartOfBlocks: true
MaxEmptyLinesToKeep: 1
NamespaceIndentation: None
ObjCSpaceAfterProperty: false
ObjCSpaceBeforeProtocolList: true
PointerAlignment: Left
ReflowComments: false
# SeparateDefinitionBlocks: Always # Only support since clang-format 14
SpaceAfterCStyleCast: false
SpaceAfterLogicalNot: false
SpaceAfterTemplateKeyword: true
SpaceBeforeAssignmentOperators: true
SpaceBeforeCpp11BracedList: false
SpaceBeforeCtorInitializerColon: true
SpaceBeforeInheritanceColon: true
SpaceBeforeParens: ControlStatements
SpaceBeforeRangeBasedForLoopColon: true
SpaceBeforeSquareBrackets: false
SpaceInEmptyParentheses: false
SpacesBeforeTrailingComments: 2
SpacesInAngles: false
SpacesInCStyleCastParentheses: false
SpacesInContainerLiterals: false
SpacesInParentheses: false
SpacesInSquareBrackets: false
Standard: c++11
TabWidth: 8
UseTab: Never

245
.cmake-format.yaml Normal file
View File

@ -0,0 +1,245 @@
_help_parse: Options affecting listfile parsing
parse:
_help_additional_commands:
- Specify structure for custom cmake functions
additional_commands:
foo:
flags:
- BAR
- BAZ
kwargs:
HEADERS: '*'
SOURCES: '*'
DEPENDS: '*'
_help_override_spec:
- Override configurations per-command where available
override_spec: {}
_help_vartags:
- Specify variable tags.
vartags: []
_help_proptags:
- Specify property tags.
proptags: []
_help_format: Options affecting formatting.
format:
_help_disable:
- Disable formatting entirely, making cmake-format a no-op
disable: false
_help_line_width:
- How wide to allow formatted cmake files
line_width: 80
_help_tab_size:
- How many spaces to tab for indent
tab_size: 2
_help_use_tabchars:
- If true, lines are indented using tab characters (utf-8
- 0x09) instead of <tab_size> space characters (utf-8 0x20).
- In cases where the layout would require a fractional tab
- character, the behavior of the fractional indentation is
- governed by <fractional_tab_policy>
use_tabchars: false
_help_fractional_tab_policy:
- If <use_tabchars> is True, then the value of this variable
- indicates how fractional indentions are handled during
- whitespace replacement. If set to 'use-space', fractional
- indentation is left as spaces (utf-8 0x20). If set to
- '`round-up` fractional indentation is replaced with a single'
- tab character (utf-8 0x09) effectively shifting the column
- to the next tabstop
fractional_tab_policy: use-space
_help_max_subgroups_hwrap:
- If an argument group contains more than this many sub-groups
- (parg or kwarg groups) then force it to a vertical layout.
max_subgroups_hwrap: 2
_help_max_pargs_hwrap:
- If a positional argument group contains more than this many
- arguments, then force it to a vertical layout.
max_pargs_hwrap: 6
_help_max_rows_cmdline:
- If a cmdline positional group consumes more than this many
- lines without nesting, then invalidate the layout (and nest)
max_rows_cmdline: 2
_help_separate_ctrl_name_with_space:
- If true, separate flow control names from their parentheses
- with a space
separate_ctrl_name_with_space: false
_help_separate_fn_name_with_space:
- If true, separate function names from parentheses with a
- space
separate_fn_name_with_space: false
_help_dangle_parens:
- If a statement is wrapped to more than one line, than dangle
- the closing parenthesis on its own line.
dangle_parens: false
_help_dangle_align:
- If the trailing parenthesis must be 'dangled' on its on
- 'line, then align it to this reference: `prefix`: the start'
- 'of the statement, `prefix-indent`: the start of the'
- 'statement, plus one indentation level, `child`: align to'
- the column of the arguments
dangle_align: prefix
_help_min_prefix_chars:
- If the statement spelling length (including space and
- parenthesis) is smaller than this amount, then force reject
- nested layouts.
min_prefix_chars: 4
_help_max_prefix_chars:
- If the statement spelling length (including space and
- parenthesis) is larger than the tab width by more than this
- amount, then force reject un-nested layouts.
max_prefix_chars: 10
_help_max_lines_hwrap:
- If a candidate layout is wrapped horizontally but it exceeds
- this many lines, then reject the layout.
max_lines_hwrap: 2
_help_line_ending:
- What style line endings to use in the output.
line_ending: unix
_help_command_case:
- Format command names consistently as 'lower' or 'upper' case
command_case: canonical
_help_keyword_case:
- Format keywords consistently as 'lower' or 'upper' case
keyword_case: unchanged
_help_always_wrap:
- A list of command names which should always be wrapped
always_wrap: []
_help_enable_sort:
- If true, the argument lists which are known to be sortable
- will be sorted lexicographicall
enable_sort: true
_help_autosort:
- If true, the parsers may infer whether or not an argument
- list is sortable (without annotation).
autosort: false
_help_require_valid_layout:
- By default, if cmake-format cannot successfully fit
- everything into the desired linewidth it will apply the
- last, most agressive attempt that it made. If this flag is
- True, however, cmake-format will print error, exit with non-
- zero status code, and write-out nothing
require_valid_layout: false
_help_layout_passes:
- A dictionary mapping layout nodes to a list of wrap
- decisions. See the documentation for more information.
layout_passes: {}
_help_markup: Options affecting comment reflow and formatting.
markup:
_help_bullet_char:
- What character to use for bulleted lists
bullet_char: '*'
_help_enum_char:
- What character to use as punctuation after numerals in an
- enumerated list
enum_char: .
_help_first_comment_is_literal:
- If comment markup is enabled, don't reflow the first comment
- block in each listfile. Use this to preserve formatting of
- your copyright/license statements.
first_comment_is_literal: false
_help_literal_comment_pattern:
- If comment markup is enabled, don't reflow any comment block
- which matches this (regex) pattern. Default is `None`
- (disabled).
literal_comment_pattern: null
_help_fence_pattern:
- Regular expression to match preformat fences in comments
- default= ``r'^\s*([`~]{3}[`~]*)(.*)$'``
fence_pattern: ^\s*([`~]{3}[`~]*)(.*)$
_help_ruler_pattern:
- Regular expression to match rulers in comments default=
- '``r''^\s*[^\w\s]{3}.*[^\w\s]{3}$''``'
ruler_pattern: ^\s*[^\w\s]{3}.*[^\w\s]{3}$
_help_explicit_trailing_pattern:
- If a comment line matches starts with this pattern then it
- is explicitly a trailing comment for the preceeding
- argument. Default is '#<'
explicit_trailing_pattern: '#<'
_help_hashruler_min_length:
- If a comment line starts with at least this many consecutive
- hash characters, then don't lstrip() them off. This allows
- for lazy hash rulers where the first hash char is not
- separated by space
hashruler_min_length: 10
_help_canonicalize_hashrulers:
- If true, then insert a space between the first hash char and
- remaining hash chars in a hash ruler, and normalize its
- length to fill the column
canonicalize_hashrulers: true
_help_enable_markup:
- enable comment markup parsing and reflow
enable_markup: true
_help_lint: Options affecting the linter
lint:
_help_disabled_codes:
- a list of lint codes to disable
disabled_codes: []
_help_function_pattern:
- regular expression pattern describing valid function names
function_pattern: '[0-9a-z_]+'
_help_macro_pattern:
- regular expression pattern describing valid macro names
macro_pattern: '[0-9A-Z_]+'
_help_global_var_pattern:
- regular expression pattern describing valid names for
- variables with global (cache) scope
global_var_pattern: '[A-Z][0-9A-Z_]+'
_help_internal_var_pattern:
- regular expression pattern describing valid names for
- variables with global scope (but internal semantic)
internal_var_pattern: _[A-Z][0-9A-Z_]+
_help_local_var_pattern:
- regular expression pattern describing valid names for
- variables with local scope
local_var_pattern: '[a-z][a-z0-9_]+'
_help_private_var_pattern:
- regular expression pattern describing valid names for
- privatedirectory variables
private_var_pattern: _[0-9a-z_]+
_help_public_var_pattern:
- regular expression pattern describing valid names for public
- directory variables
public_var_pattern: '[A-Z][0-9A-Z_]+'
_help_argument_var_pattern:
- regular expression pattern describing valid names for
- function/macro arguments and loop variables.
argument_var_pattern: '[a-z][a-z0-9_]+'
_help_keyword_pattern:
- regular expression pattern describing valid names for
- keywords used in functions or macros
keyword_pattern: '[A-Z][0-9A-Z_]+'
_help_max_conditionals_custom_parser:
- In the heuristic for C0201, how many conditionals to match
- within a loop in before considering the loop a parser.
max_conditionals_custom_parser: 2
_help_min_statement_spacing:
- Require at least this many newlines between statements
min_statement_spacing: 1
_help_max_statement_spacing:
- Require no more than this many newlines between statements
max_statement_spacing: 2
max_returns: 6
max_branches: 12
max_arguments: 5
max_localvars: 15
max_statements: 50
_help_encode: Options affecting file encoding
encode:
_help_emit_byteorder_mark:
- If true, emit the unicode byte-order mark (BOM) at the start
- of the file
emit_byteorder_mark: false
_help_input_encoding:
- Specify the encoding of the input file. Defaults to utf-8
input_encoding: utf-8
_help_output_encoding:
- Specify the encoding of the output file. Defaults to utf-8.
- Note that cmake only claims to support utf-8 so be careful
- when using anything else
output_encoding: utf-8
_help_misc: Miscellaneous configurations options.
misc:
_help_per_command:
- A dictionary containing any per-command configuration
- overrides. Currently only `command_case` is supported.
per_command: {}

26
.github/ISSUE_TEMPLATE/tensorrtx-issue-template.md vendored Normal file
View File

@ -0,0 +1,26 @@
---
name: tensorrtx issue template
about: To understand your issue better
title: ''
labels: ''
assignees: ''
---
## Env
- GPU, e.g. V100, RTX2080, TX2, Xavier NX, Nano, etc.
- OS, e.g. Ubuntu16.04, Win10, etc.
- Cuda version
- TensorRT version
## About this repo
- which branch/tag/commit are you using?
- which model? yolov5, retinaface?
## Your problem
- what is your command? e.g. `sudo ./yolov5 -s`
- what's your output?
- what output do you expect?

17
.github/stale.yml vendored Normal file
View File

@ -0,0 +1,17 @@
# Number of days of inactivity before an issue becomes stale
daysUntilStale: 60
# Number of days of inactivity before a stale issue is closed
daysUntilClose: 7
# Issues with these labels will never be considered stale
exemptLabels:
- pinned
- security
# Label to use when marking an issue as stale
staleLabel: wontfix
# Comment to post when marking an issue as stale. Set to `false` to disable
markComment: >
This issue has been automatically marked as stale because it has not had
recent activity. It will be closed if no further activity occurs. Thank you
for your contributions.
# Comment to post when closing a stale issue. Set to `false` to disable
closeComment: false

20
.github/workflows/pre-commit.yml vendored Normal file
View File

@ -0,0 +1,20 @@
name: pre-commit
on:
pull_request:
branches:
- master
- trt10
jobs:
pre-commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
with:
# grab the history of the PR
fetch-depth: 0
- uses: actions/setup-python@v3
- uses: pre-commit/action@v3.0.1
with:
extra_args: --from-ref ${{ github.event.pull_request.base.sha }} --to-ref ${{ github.event.pull_request.head.sha }}

77
.gitignore vendored Normal file
View File

@ -0,0 +1,77 @@
build
*.wts
*.engine
*/*.ppm
*idea*
.vscode/*
!.vscode/settings.json
!.vscode/tasks.json
!.vscode/launch.json
!.vscode/extensions.json
!.vscode/*.code-snippets
# Local History for Visual Studio Code
.history/
# Built Visual Studio Code Extensions
*.vsix
.vscode/*
!.vscode/settings.json
!.vscode/tasks.json
!.vscode/launch.json
!.vscode/extensions.json
!.vscode/*.code-snippets
# Local History for Visual Studio Code
.history/
# Built Visual Studio Code Extensions
*.vsix
# Prerequisites
*.d
# Compiled Object files
*.slo
*.lo
*.o
*.obj
# Precompiled Headers
*.gch
*.pch
# Compiled Dynamic libraries
*.so
*.dylib
*.dll
# Fortran module files
*.mod
*.smod
# Compiled Static libraries
*.lai
*.la
*.a
*.lib
# Executables
*.exe
*.out
*.app
CMakeLists.txt.user
CMakeCache.txt
CMakeFiles
CMakeScripts
Testing
Makefile
cmake_install.cmake
install_manifest.txt
compile_commands.json
CTestTestfile.cmake
_deps
CMakeUserPresets.json

27
.pre-commit-config.yaml Normal file
View File

@ -0,0 +1,27 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0
hooks:
- id: check-merge-conflict
- id: check-symlinks
- id: end-of-file-fixer
- id: trailing-whitespace
- id: check-added-large-files
- repo: https://github.com/pre-commit/mirrors-clang-format
rev: v14.0.6
hooks:
- id: clang-format
types_or: [c++, c, cuda]
- repo: https://github.com/PyCQA/flake8
rev: 7.0.0
hooks:
- id: flake8
args: [--max-line-length=120]
- repo: https://github.com/cheshirekow/cmake-format-precommit
rev: v0.6.13
hooks:
- id: cmake-format
additional_dependencies: [pyyaml]
args: [--in-place, -c, .cmake-format.yaml]
types: [file]
files: (\.cmake|CMakeLists.txt)(.in)?$

21
LICENSE Normal file
View File

@ -0,0 +1,21 @@
MIT License
Copyright (c) 2019-2020 Wang Xinyu
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.

177
README.md Normal file
View File

@ -0,0 +1,177 @@
# TensorRTx
TensorRTx aims to implement popular deep learning networks with TensorRT network definition API.
Why don't we use a parser (ONNX parser, UFF parser, caffe parser, etc), but use complex APIs to build a network from scratch? I have summarized the advantages in the following aspects.
- **Flexible**, easy to modify the network, add/delete a layer or input/output tensor, replace a layer, merge layers, integrate preprocessing and postprocessing into network, etc.
- **Debuggable**, construct the entire network in an incremental development manner, easy to get middle layer results.
- **Educational**, learn about the network structure during this development, rather than treating everything as a black box.
The basic workflow of TensorRTx is:
1. Get the trained models from pytorch, mxnet or tensorflow, etc. Some pytorch models can be found in my repo [pytorchx](https://github.com/wang-xinyu/pytorchx), the remaining are from popular open-source repos.
2. Export the weights to a plain text file -- [.wts file](./tutorials/getting_started.md#the-wts-content-format).
3. Load weights in TensorRT, define the network, build a TensorRT engine.
4. Load the TensorRT engine and run inference.
## News
- `22 Oct 2024`. [lindsayshuo](https://github.com/lindsayshuo): YOLOv8-obb
- `18 Oct 2024`. [zgjja](https://github.com/zgjja): Rafactor docker image.
- `11 Oct 2024`. [mpj1234](https://github.com/mpj1234): YOLO11
- `9 Oct 2024`. [Phoenix8215](https://github.com/Phoenix8215): GhostNet V1 and V2.
- `21 Aug 2024`. [Lemonononon](https://github.com/Lemonononon): real-esrgan-general-x4v3
- `29 Jul 2024`. [mpj1234](https://github.com/mpj1234): Check the YOLOv5, YOLOv8 & YOLOv10 in TensorRT 10.x API, branch → [trt10](https://github.com/wang-xinyu/tensorrtx/tree/trt10)
- `29 Jul 2024`. [mpj1234](https://github.com/mpj1234): YOLOv10
- `21 Jun 2024`. [WuxinrongY](https://github.com/WuxinrongY): YOLOv9-T, YOLOv9-S, YOLOv9-M
- `28 Apr 2024`. [lindsayshuo](https://github.com/lindsayshuo): YOLOv8-pose
- `22 Apr 2024`. [B1SH0PP](https://github.com/B1SH0PP): EfficientAd: Accurate Visual Anomaly Detection at Millisecond-Level Latencies.
- `18 Apr 2024`. [lindsayshuo](https://github.com/lindsayshuo): YOLOv8-p2
- `12 Mar 2024`. [lindsayshuo](https://github.com/lindsayshuo): YOLOv8-cls
- `11 Mar 2024`. [WuxinrongY](https://github.com/WuxinrongY): YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
- `7 Mar 2024`. [AadeIT](https://github.com/AadeIT): CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes
- `17 Oct 2023`. [Rex-LK](https://github.com/Rex-LK): YOLOv8-Seg
## Tutorials
- [How to make contribution](./tutorials/contribution.md)
- [Install the dependencies.](./tutorials/install.md)
- [A guide for quickly getting started, taking lenet5 as a demo.](./tutorials/getting_started.md)
- [The .wts file content format](./tutorials/getting_started.md#the-wts-content-format)
- [Frequently Asked Questions (FAQ)](./tutorials/faq.md)
- [Migrating from TensorRT 4 to 7](./tutorials/migrating_from_tensorrt_4_to_7.md)
- [How to implement multi-GPU processing, taking YOLOv4 as example](./tutorials/multi_GPU_processing.md)
- [Check if Your GPU support FP16/INT8](./tutorials/check_fp16_int8_support.md)
- [How to Compile and Run on Windows](./tutorials/run_on_windows.md)
- [Deploy YOLOv4 with Triton Inference Server](https://github.com/isarsoft/yolov4-triton-tensorrt)
- [From pytorch to trt step by step, hrnet as example(Chinese)](./tutorials/from_pytorch_to_trt_stepbystep_hrnet.md)
## Test Environment
1. TensorRT 7.x
2. TensorRT 8.x(Some of the models support 8.x)
## How to run
Each folder has a readme inside, which explains how to run the models inside.
## Models
Following models are implemented.
|Name | Description |
|-|-|
|[mlp](./mlp) | the very basic model for starters, properly documented |
|[lenet](./lenet) | the simplest, as a "hello world" of this project |
|[alexnet](./alexnet)| easy to implement, all layers are supported in tensorrt |
|[googlenet](./googlenet)| GoogLeNet (Inception v1) |
|[inception](./inception)| Inception v3, v4 |
|[mnasnet](./mnasnet)| MNASNet with depth multiplier of 0.5 from the paper |
|[mobilenet](./mobilenet)| MobileNet v2, v3-small, v3-large |
|[resnet](./resnet)| resnet-18, resnet-50 and resnext50-32x4d are implemented |
|[senet](./senet)| se-resnet50 |
|[shufflenet](./shufflenetv2)| ShuffleNet v2 with 0.5x output channels |
|[squeezenet](./squeezenet)| SqueezeNet 1.1 model |
|[vgg](./vgg)| VGG 11-layer model |
|[yolov3-tiny](./yolov3-tiny)| weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|[yolov3](./yolov3)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|[yolov3-spp](./yolov3-spp)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|[yolov4](./yolov4)| CSPDarknet53, weights from [AlexeyAB/darknet](https://github.com/AlexeyAB/darknet#pre-trained-models), pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|[yolov5](./yolov5)| yolov5 v1.0-v7.0 of [ultralytics/yolov5](https://github.com/ultralytics/yolov5), detection, classification and instance segmentation |
|[yolov7](./yolov7)| yolov7 v0.1, pytorch implementation from [WongKinYiu/yolov7](https://github.com/WongKinYiu/yolov7) |
|[yolov8](./yolov8)| yolov8, pytorch implementation from [ultralytics](https://github.com/ultralytics/ultralytics) |
|[yolov9](./yolov9)| The Pytorch implementation is [WongKinYiu/yolov9](https://github.com/WongKinYiu/yolov9). |
|[yolov10](./yolov10)| The Pytorch implementation is [THU-MIG/yolov10](https://github.com/THU-MIG/yolov10). |
|[yolo11](./yolo11)| The Pytorch implementation is [ultralytics](https://github.com/ultralytics/ultralytics). |
|[yolop](./yolop)| yolop, pytorch implementation from [hustvl/YOLOP](https://github.com/hustvl/YOLOP) |
|[retinaface](./retinaface)| resnet50 and mobilnet0.25, weights from [biubug6/Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface) |
|[arcface](./arcface)| LResNet50E-IR, LResNet100E-IR and MobileFaceNet, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface) |
|[retinafaceAntiCov](./retinafaceAntiCov)| mobilenet0.25, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface), retinaface anti-COVID-19, detect face and mask attribute |
|[dbnet](./dbnet)| Scene Text Detection, weights from [BaofengZan/DBNet.pytorch](https://github.com/BaofengZan/DBNet.pytorch) |
|[crnn](./crnn)| pytorch implementation from [meijieru/crnn.pytorch](https://github.com/meijieru/crnn.pytorch) |
|[ufld](./ufld)| pytorch implementation from [Ultra-Fast-Lane-Detection](https://github.com/cfzd/Ultra-Fast-Lane-Detection), ECCV2020 |
|[hrnet](./hrnet)| hrnet-image-classification and hrnet-semantic-segmentation, pytorch implementation from [HRNet-Image-Classification](https://github.com/HRNet/HRNet-Image-Classification) and [HRNet-Semantic-Segmentation](https://github.com/HRNet/HRNet-Semantic-Segmentation) |
|[psenet](./psenet)| PSENet Text Detection, tensorflow implementation from [liuheng92/tensorflow_PSENet](https://github.com/liuheng92/tensorflow_PSENet) |
|[ibnnet](./ibnnet)| IBN-Net, pytorch implementation from [XingangPan/IBN-Net](https://github.com/XingangPan/IBN-Net), ECCV2018 |
|[unet](./unet)| U-Net, pytorch implementation from [milesial/Pytorch-UNet](https://github.com/milesial/Pytorch-UNet) |
|[repvgg](./repvgg)| RepVGG, pytorch implementation from [DingXiaoH/RepVGG](https://github.com/DingXiaoH/RepVGG) |
|[lprnet](./lprnet)| LPRNet, pytorch implementation from [xuexingyu24/License_Plate_Detection_Pytorch](https://github.com/xuexingyu24/License_Plate_Detection_Pytorch) |
|[refinedet](./refinedet)| RefineDet, pytorch implementation from [luuuyi/RefineDet.PyTorch](https://github.com/luuuyi/RefineDet.PyTorch) |
|[densenet](./densenet)| DenseNet-121, from torchvision.models |
|[rcnn](./rcnn)| FasterRCNN and MaskRCNN, model from [detectron2](https://github.com/facebookresearch/detectron2) |
|[tsm](./tsm)| TSM: Temporal Shift Module for Efficient Video Understanding, ICCV2019 |
|[scaled-yolov4](./scaled-yolov4)| yolov4-csp, pytorch from [WongKinYiu/ScaledYOLOv4](https://github.com/WongKinYiu/ScaledYOLOv4) |
|[centernet](./centernet)| CenterNet DLA-34, pytorch from [xingyizhou/CenterNet](https://github.com/xingyizhou/CenterNet) |
|[efficientnet](./efficientnet)| EfficientNet b0-b8 and l2, pytorch from [lukemelas/EfficientNet-PyTorch](https://github.com/lukemelas/EfficientNet-PyTorch) |
|[detr](./detr)| DE⫶TR, pytorch from [facebookresearch/detr](https://github.com/facebookresearch/detr) |
|[swin-transformer](./swin-transformer)| Swin Transformer - Semantic Segmentation, only support Swin-T. The Pytorch implementation is [microsoft/Swin-Transformer](https://github.com/microsoft/Swin-Transformer.git) |
|[real-esrgan](./real-esrgan)| Real-ESRGAN. The Pytorch implementation is [real-esrgan](https://github.com/xinntao/Real-ESRGAN) |
|[superpoint](./superpoint)| SuperPoint. The Pytorch model is from [magicleap/SuperPointPretrainedNetwork](https://github.com/magicleap/SuperPointPretrainedNetwork) |
|[csrnet](./csrnet)| CSRNet. The Pytorch implementation is [leeyeehoo/CSRNet-pytorch](https://github.com/leeyeehoo/CSRNet-pytorch) |
|[EfficientAd](./efficient_ad)| EfficientAd: Accurate Visual Anomaly Detection at Millisecond-Level Latencies. From [anomalib](https://github.com/openvinotoolkit/anomalib) |
## Model Zoo
The .wts files can be downloaded from model zoo for quick evaluation. But it is recommended to convert .wts from pytorch/mxnet/tensorflow model, so that you can retrain your own model.
[GoogleDrive](https://drive.google.com/drive/folders/1Ri0IDa5OChtcA3zjqRTW57uG6TnfN4Do?usp=sharing) | [BaiduPan](https://pan.baidu.com/s/19s6hO8esU7-TtZEXN7G3OA) pwd: uvv2
## Tricky Operations
Some tricky operations encountered in these models, already solved, but might have better solutions.
|Name | Description |
|-|-|
|BatchNorm| Implement by a scale layer, used in resnet, googlenet, mobilenet, etc. |
|MaxPool2d(ceil_mode=True)| use a padding layer before maxpool to solve ceil_mode=True, see googlenet. |
|average pool with padding| use setAverageCountExcludesPadding() when necessary, see inception. |
|relu6| use `Relu6(x) = Relu(x) - Relu(x-6)`, see mobilenet. |
|torch.chunk()| implement the 'chunk(2, dim=C)' by tensorrt plugin, see shufflenet. |
|channel shuffle| use two shuffle layers to implement `channel_shuffle`, see shufflenet. |
|adaptive pool| use fixed input dimension, and use regular average pooling, see shufflenet. |
|leaky relu| I wrote a leaky relu plugin, but PRelu in `NvInferPlugin.h` can be used, see yolov3 in branch `trt4`. |
|yolo layer v1| yolo layer is implemented as a plugin, see yolov3 in branch `trt4`. |
|yolo layer v2| three yolo layers implemented in one plugin, see yolov3-spp. |
|upsample| replaced by a deconvolution layer, see yolov3. |
|hsigmoid| hard sigmoid is implemented as a plugin, hsigmoid and hswish are used in mobilenetv3 |
|retinaface output decode| implement a plugin to decode bbox, confidence and landmarks, see retinaface. |
|mish| mish activation is implemented as a plugin, mish is used in yolov4 |
|prelu| mxnet's prelu activation with trainable gamma is implemented as a plugin, used in arcface |
|HardSwish| hard_swish = x * hard_sigmoid, used in yolov5 v3.0 |
|LSTM| Implemented pytorch nn.LSTM() with tensorrt api |
## Speed Benchmark
| Models | Device | BatchSize | Mode | Input Shape(HxW) | FPS |
|-|-|:-:|:-:|:-:|:-:|
| YOLOv3-tiny | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 333 |
| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 39.2 |
| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | INT8 | 608x608 | 71.4 |
| YOLOv3-spp(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 38.5 |
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 35.7 |
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 4 | FP32 | 608x608 | 40.9 |
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 8 | FP32 | 608x608 | 41.3 |
| YOLOv5-s v3.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 142 |
| YOLOv5-s v3.0 | Xeon E5-2620/GTX1080 | 4 | FP32 | 608x608 | 173 |
| YOLOv5-s v3.0 | Xeon E5-2620/GTX1080 | 8 | FP32 | 608x608 | 190 |
| YOLOv5-m v3.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 71 |
| YOLOv5-l v3.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 43 |
| YOLOv5-x v3.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 29 |
| YOLOv5-s v4.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 142 |
| YOLOv5-m v4.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 71 |
| YOLOv5-l v4.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 40 |
| YOLOv5-x v4.0 | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 27 |
| RetinaFace(resnet50) | Xeon E5-2620/GTX1080 | 1 | FP32 | 480x640 | 90 |
| RetinaFace(resnet50) | Xeon E5-2620/GTX1080 | 1 | INT8 | 480x640 | 204 |
| RetinaFace(mobilenet0.25) | Xeon E5-2620/GTX1080 | 1 | FP32 | 480x640 | 417 |
| ArcFace(LResNet50E-IR) | Xeon E5-2620/GTX1080 | 1 | FP32 | 112x112 | 333 |
| CRNN | Xeon E5-2620/GTX1080 | 1 | FP32 | 32x100 | 1000 |
Help wanted, if you got speed results, please add an issue or PR.
## Acknowledgments & Contact
Any contributions, questions and discussions are welcomed, contact me by following info.
E-mail: wangxinyu_es@163.com
WeChat ID: wangxinyu0375 (可加我微信进tensorrtx交流群，**备注：tensorrtx**)

25
alexnet/CMakeLists.txt Normal file
View File

@ -0,0 +1,25 @@
cmake_minimum_required(VERSION 2.6)
project(alexnet)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
add_executable(alexnet ${PROJECT_SOURCE_DIR}/alex.cpp)
target_link_libraries(alexnet nvinfer)
target_link_libraries(alexnet cudart)
add_definitions(-O2 -pthread)

33
alexnet/README.md Normal file
View File

@ -0,0 +1,33 @@
# alexnet
AlexNet model architecture from the "One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
For the details, you can refer to [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet)
This alexnet is just several `conv-relu-pool` blocks followed by several `fc-relu`, nothing special. All layers can be implemented by tensorrt api, including `addConvolution`, `addActivation`, `addPooling`, `addFullyConnected`.
```
// 1. generate alexnet.wts from [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet)
// 2. put alexnet.wts into tensorrtx/alexnet
// 3. build and run
cd tensorrtx/alexnet
mkdir build
cd build
cmake ..
make
sudo ./alexnet -s // serialize model to plan file i.e. 'alexnet.engine'
sudo ./alexnet -d // deserialize plan file and run inference
// 4. see if the output is same as pytorchx/alexnet
```

297
alexnet/alex.cpp Normal file
View File

@ -0,0 +1,297 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <fstream>
#include <map>
#include <chrono>
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 224;
static const int INPUT_W = 224;
static const int OUTPUT_SIZE = 1000;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
using namespace nvinfer1;
static Logger gLogger;
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)
{
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape { 1, 1, 32, 32 } with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../alexnet.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{11, 11}, weightMap["features.0.weight"], weightMap["features.0.bias"]);
assert(conv1);
conv1->setStrideNd(DimsHW{4, 4});
conv1->setPaddingNd(DimsHW{2, 2});
// Add activation layer using the ReLU algorithm.
IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// Add max pooling layer with stride of 2x2 and kernel size of 2x2.
IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
assert(pool1);
pool1->setStrideNd(DimsHW{2, 2});
IConvolutionLayer* conv2 = network->addConvolutionNd(*pool1->getOutput(0), 192, DimsHW{5, 5}, weightMap["features.3.weight"], weightMap["features.3.bias"]);
assert(conv2);
conv2->setPaddingNd(DimsHW{2, 2});
IActivationLayer* relu2 = network->addActivation(*conv2->getOutput(0), ActivationType::kRELU);
assert(relu2);
IPoolingLayer* pool2 = network->addPoolingNd(*relu2->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
assert(pool2);
pool2->setStrideNd(DimsHW{2, 2});
IConvolutionLayer* conv3 = network->addConvolutionNd(*pool2->getOutput(0), 384, DimsHW{3, 3}, weightMap["features.6.weight"], weightMap["features.6.bias"]);
assert(conv3);
conv3->setPaddingNd(DimsHW{1, 1});
IActivationLayer* relu3 = network->addActivation(*conv3->getOutput(0), ActivationType::kRELU);
assert(relu3);
IConvolutionLayer* conv4 = network->addConvolutionNd(*relu3->getOutput(0), 256, DimsHW{3, 3}, weightMap["features.8.weight"], weightMap["features.8.bias"]);
assert(conv4);
conv4->setPaddingNd(DimsHW{1, 1});
IActivationLayer* relu4 = network->addActivation(*conv4->getOutput(0), ActivationType::kRELU);
assert(relu4);
IConvolutionLayer* conv5 = network->addConvolutionNd(*relu4->getOutput(0), 256, DimsHW{3, 3}, weightMap["features.10.weight"], weightMap["features.10.bias"]);
assert(conv5);
conv5->setPaddingNd(DimsHW{1, 1});
IActivationLayer* relu5 = network->addActivation(*conv5->getOutput(0), ActivationType::kRELU);
assert(relu5);
IPoolingLayer* pool3 = network->addPoolingNd(*relu5->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
assert(pool3);
pool3->setStrideNd(DimsHW{2, 2});
IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool3->getOutput(0), 4096, weightMap["classifier.1.weight"], weightMap["classifier.1.bias"]);
assert(fc1);
IActivationLayer* relu6 = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU);
assert(relu6);
IFullyConnectedLayer* fc2 = network->addFullyConnected(*relu6->getOutput(0), 4096, weightMap["classifier.4.weight"], weightMap["classifier.4.bias"]);
assert(fc2);
IActivationLayer* relu7 = network->addActivation(*fc2->getOutput(0), ActivationType::kRELU);
assert(relu7);
IFullyConnectedLayer* fc3 = network->addFullyConnected(*relu7->getOutput(0), 1000, weightMap["classifier.6.weight"], weightMap["classifier.6.bias"]);
assert(fc3);
fc3->getOutput(0)->setName(OUTPUT_BLOB_NAME);
std::cout << "set name out" << std::endl;
network->markOutput(*fc3->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(1 << 20);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "build out" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
{
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv)
{
if (argc != 2) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./alexnet -s // serialize model to plan file" << std::endl;
std::cerr << "./alexnet -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(1, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("alexnet.engine", std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 1;
} else if (std::string(argv[1]) == "-d") {
std::ifstream file("alexnet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
return -1;
}
// Subtract mean from image
float data[3 * INPUT_H * INPUT_W];
for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
data[i] = 1;
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
// Run inference
float prob[OUTPUT_SIZE];
for (int i = 0; i < 100; i++) {
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, 1);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
// Print histogram of the output distribution
std::cout << "\nOutput:\n\n";
for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
{
std::cout << prob[i] << ", ";
if (i % 10 == 0) std::cout << i / 10 << std::endl;
}
std::cout << std::endl;
return 0;
}

259
alexnet/alexnet.py Normal file
View File

@ -0,0 +1,259 @@
import os
import sys
import struct
import argparse
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
BATCH_SIZE = 1
INPUT_H = 224
INPUT_W = 224
OUTPUT_SIZE = 1000
INPUT_BLOB_NAME = "data"
OUTPUT_BLOB_NAME = "prob"
WEIGHT_PATH = "./alexnet.wts"
ENGINE_PATH = "./alexnet.engine"
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def load_weights(file):
print(f"Loading weights: {file}")
assert os.path.exists(file), 'Unable to load weight file.'
weight_map = {}
with open(file, "r") as f:
lines = [line.strip() for line in f]
count = int(lines[0])
assert count == len(lines) - 1
for i in range(1, count + 1):
splits = lines[i].split(" ")
name = splits[0]
cur_count = int(splits[1])
assert cur_count + 2 == len(splits)
values = []
for j in range(2, len(splits)):
# hex string to bytes to float
values.append(struct.unpack(">f", bytes.fromhex(splits[j])))
weight_map[name] = np.array(values, dtype=np.float32)
return weight_map
def create_engine(max_batch_size, builder, config, dt):
weight_map = load_weights(WEIGHT_PATH)
network = builder.create_network()
data = network.add_input(INPUT_BLOB_NAME, dt, (3, INPUT_H, INPUT_W))
assert data
conv1 = network.add_convolution(input=data,
num_output_maps=64,
kernel_shape=(11, 11),
kernel=weight_map["features.0.weight"],
bias=weight_map["features.0.bias"])
assert conv1
conv1.stride = (4, 4)
conv1.padding = (2, 2)
relu1 = network.add_activation(conv1.get_output(0), type=trt.ActivationType.RELU)
assert relu1
pool1 = network.add_pooling(input=relu1.get_output(0),
type=trt.PoolingType.MAX,
window_size=trt.DimsHW(3, 3))
assert pool1
pool1.stride_nd = (2, 2)
conv2 = network.add_convolution(input=pool1.get_output(0),
num_output_maps=192,
kernel_shape=(5, 5),
kernel=weight_map["features.3.weight"],
bias=weight_map["features.3.bias"])
assert conv2
conv2.padding = (2, 2)
relu2 = network.add_activation(conv2.get_output(0), type=trt.ActivationType.RELU)
assert relu2
pool2 = network.add_pooling(input=relu2.get_output(0),
type=trt.PoolingType.MAX,
window_size=trt.DimsHW(3, 3))
assert pool2
pool2.stride_nd = (2, 2)
conv3 = network.add_convolution(input=pool2.get_output(0),
num_output_maps=384,
kernel_shape=(3, 3),
kernel=weight_map["features.6.weight"],
bias=weight_map["features.6.bias"])
assert conv3
conv3.padding = (1, 1)
relu3 = network.add_activation(conv3.get_output(0), type=trt.ActivationType.RELU)
assert relu3
conv4 = network.add_convolution(input=relu3.get_output(0),
num_output_maps=256,
kernel_shape=(3, 3),
kernel=weight_map["features.8.weight"],
bias=weight_map["features.8.bias"])
assert conv4
conv4.padding = (1, 1)
relu4 = network.add_activation(conv4.get_output(0), type=trt.ActivationType.RELU)
assert relu4
conv5 = network.add_convolution(input=relu4.get_output(0),
num_output_maps=256,
kernel_shape=(3, 3),
kernel=weight_map["features.10.weight"],
bias=weight_map["features.10.bias"])
assert conv5
conv5.padding = (1, 1)
relu5 = network.add_activation(conv5.get_output(0), type=trt.ActivationType.RELU)
assert relu5
pool3 = network.add_pooling(input=relu5.get_output(0),
type=trt.PoolingType.MAX,
window_size=trt.DimsHW(3, 3))
assert pool3
pool3.stride_nd = (2, 2)
fc1 = network.add_fully_connected(input=pool3.get_output(0),
num_outputs=4096,
kernel=weight_map["classifier.1.weight"],
bias=weight_map["classifier.1.bias"])
assert fc1
relu6 = network.add_activation(fc1.get_output(0), type=trt.ActivationType.RELU)
assert relu6
fc2 = network.add_fully_connected(input=relu6.get_output(0),
num_outputs=4096,
kernel=weight_map["classifier.4.weight"],
bias=weight_map["classifier.4.bias"])
assert fc2
relu7 = network.add_activation(fc2.get_output(0), type=trt.ActivationType.RELU)
assert relu7
fc3 = network.add_fully_connected(input=relu7.get_output(0),
num_outputs=1000,
kernel=weight_map["classifier.6.weight"],
bias=weight_map["classifier.6.bias"])
assert fc3
fc3.get_output(0).name = OUTPUT_BLOB_NAME
network.mark_output(fc3.get_output(0))
# Build Engine
builder.max_batch_size = max_batch_size
builder.max_workspace_size = 1 << 20
engine = builder.build_engine(network, config)
del network
del weight_map
return engine
def API_to_model(max_batch_size):
builder = trt.Builder(TRT_LOGGER)
config = builder.create_builder_config()
engine = create_engine(max_batch_size, builder, config, trt.float32)
assert engine
with open(ENGINE_PATH, "wb") as f:
f.write(engine.serialize())
del engine
del builder
del config
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-s", action='store_true')
parser.add_argument("-d", action='store_true')
args = parser.parse_args()
if not (args.s ^ args.d):
print(
"arguments not right!\n"
"python alexnet.py -s # serialize model to plan file\n"
"python alexnet.py -d # deserialize plan file and run inference"
)
sys.exit()
if args.s:
API_to_model(BATCH_SIZE)
else:
runtime = trt.Runtime(TRT_LOGGER)
assert runtime
with open(ENGINE_PATH, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
assert engine
context = engine.create_execution_context()
assert context
data = np.ones((BATCH_SIZE * 3 * INPUT_H * INPUT_W), dtype=np.float32)
inputs, outputs, bindings, stream = allocate_buffers(engine)
inputs[0].host = data
trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
print(f'Output: \n{trt_outputs[0][:10]}\n{trt_outputs[0][-10:]}')

503
alexnet/logging.h Normal file
View File

@ -0,0 +1,503 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

51
arcface/CMakeLists.txt Normal file
View File

@ -0,0 +1,51 @@
cmake_minimum_required(VERSION 2.6)
project(arcface)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
message("embed_platform on")
include_directories(/usr/local/cuda/targets/aarch64-linux/include)
link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
else()
message("embed_platform off")
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
endif()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
cuda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/prelu.cu)
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(arcface-r50 ${PROJECT_SOURCE_DIR}/arcface-r50.cpp)
target_link_libraries(arcface-r50 nvinfer)
target_link_libraries(arcface-r50 cudart)
target_link_libraries(arcface-r50 myplugins)
target_link_libraries(arcface-r50 ${OpenCV_LIBS})
add_executable(arcface-mobilefacenet ${PROJECT_SOURCE_DIR}/arcface-mobilefacenet.cpp)
target_link_libraries(arcface-mobilefacenet nvinfer)
target_link_libraries(arcface-mobilefacenet cudart)
target_link_libraries(arcface-mobilefacenet myplugins)
target_link_libraries(arcface-mobilefacenet ${OpenCV_LIBS})
add_executable(arcface-r100 ${PROJECT_SOURCE_DIR}/arcface-r100.cpp)
target_link_libraries(arcface-r100 nvinfer)
target_link_libraries(arcface-r100 cudart)
target_link_libraries(arcface-r100 myplugins)
target_link_libraries(arcface-r100 ${OpenCV_LIBS})
add_definitions(-O2 -pthread)

69
arcface/README.md Normal file
View File

@ -0,0 +1,69 @@
# arcface
### TensortRT 8
The mxnet implementation is from [deepinsight/insightface.](https://github.com/deepinsight/insightface)
**Updated Pretrained Weights:** ArcFace-R100 [Insight Face Google Drive](https://drive.google.com/file/d/1Hc5zUfBATaXUgcU2haUNa7dcaZSw95h2/view)
---
**Previous Pre-trained models:** The pretrained models are from [LResNet50E-IR,ArcFace@ms1m-refine-v1](https://github.com/deepinsight/insightface/wiki/Model-Zoo#32-lresnet50e-irarcfacems1m-refine-v1), [LResNet100E-IR,ArcFace@ms1m-refine-v2](https://github.com/deepinsight/insightface/wiki/Model-Zoo#31-lresnet100e-irarcfacems1m-refine-v2) and [MobileFaceNet,ArcFace@ms1m-refine-v1](https://github.com/deepinsight/insightface/wiki/Model-Zoo#34-mobilefacenetarcfacems1m-refine-v1)
---
The two input images used in this project are joey0.ppm and joey1.ppm, download them from [Google Drive.](https://drive.google.com/drive/folders/1ctqpkRCRKyBZRCNwo9Uq4eUoMRLtFq1e). The input image is 112x112, and generated from `get_input()` in `insightface/deploy/face_model.py`, which is cropped and aligned face image.
<p align="center">
<img src="https://user-images.githubusercontent.com/15235574/83122953-f45f8d80-a106-11ea-84b0-4f6ff91b5924.jpg">
</p>
## Config
- FP16/FP32 can be selected by the macro `USE_FP16` in arcface-r50/r100/mobilefacenet.cpp
- GPU id can be selected by the macro `DEVICE` in arcface-r50/r100/mobilefacenet.cpp
## Run
1.Generate .wts file from mxnet implementation of pretrained model. The following example described how to generate arcface-r100.wts from mxnet implementation of LResNet100E-IR,ArcFace@ms1m-refine-v1.
```
git clone https://github.com/deepinsight/insightface
cd insightface
git checkout 3866cd77a6896c934b51ed39e9651b791d78bb57
cd deploy
// copy tensorrtx/arcface/gen_wts.py to here(insightface/deploy)
// download model-r100-ii.zip and unzip here(insightface/deploy)
python gen_wts.py
// a file 'arcface-r100.wts' will be generated.
// the master branch of insightface should work, if not, you can checkout 94ad870abb3203d6f31b049b70dd080dc8f33fca
// arcface-r50.wts/arcface-mobilefacenet.wts can be generated in similar way from mxnet implementation of LResNet50E-IR,ArcFace@ms1m-refine-v1/MobileFaceNet,ArcFace@ms1m-refine-v1 pretrained model.
```
2.Put .wts file into tensorrtx/arcface, build and run
```
cd tensorrtx/arcface
// download joey0.ppm and joey1.ppm, and put here(tensorrtx/arcface)
mkdir build
cd build
cmake ..
make
sudo ./arcface-r100 -s // serialize model to plan file i.e. 'arcface-r100.engine'
sudo ./arcface-r100 -d // deserialize plan file and run inference
or
sudo ./arcface-r50 -s // serialize model to plan file i.e. 'arcface-r50.engine'
sudo ./arcface-r50 -d // deserialize plan file and run inference
or
sudo ./arcface-mobilefacenet -s // serialize model to plan file i.e. 'arcface-mobilefacenet.engine'
sudo ./arcface-mobilefacenet -d // deserialize plan file and run inference
```
3.Check the output log, latency and similarity score.
## More Information
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)

451
arcface/arcface-mobilefacenet.cpp Normal file
View File

@ -0,0 +1,451 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include <dirent.h>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
//#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define BATCH_SIZE 1 // currently, only support BATCH=1
using namespace nvinfer1;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 112;
static const int INPUT_W = 112;
static const int OUTPUT_SIZE = 128;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
float *beta = (float*)weightMap[lname + "_beta"].values;
float *mean = (float*)weightMap[lname + "_moving_mean"].values;
float *var = (float*)weightMap[lname + "_moving_var"].values;
int len = weightMap[lname + "_moving_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* addPRelu(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
int len = weightMap[lname + "_gamma"].count;
float *scval_1 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
float *scval_2 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval_1[i] = -1.0;
scval_2[i] = -gamma[i];
}
Weights scale_1{ DataType::kFLOAT, scval_1, len };
Weights scale_2{ DataType::kFLOAT, scval_2, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = 0.0;
}
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
auto relu1 = network->addActivation(input, ActivationType::kRELU);
assert(relu1);
IScaleLayer* scale1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale_1, power);
assert(scale1);
auto relu2 = network->addActivation(*scale1->getOutput(0), ActivationType::kRELU);
assert(relu2);
IScaleLayer* scale2 = network->addScale(*relu2->getOutput(0), ScaleMode::kCHANNEL, shift, scale_2, power);
assert(scale2);
IElementWiseLayer* ew1 = network->addElementWise(*relu1->getOutput(0), *scale2->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
return ew1;
}
ILayer* conv_bn_relu(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, int oup, int k = 3, int p = 1, int s = 2, int groups=1) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(input, oup, DimsHW{k, k}, weightMap[lname + "_conv2d_weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
conv1->setNbGroups(groups);
auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "_batchnorm", 1e-3);
assert(bn1);
auto act1 = addPRelu(network, weightMap, *bn1->getOutput(0), lname + "_relu");
assert(act1);
return act1;
}
ILayer* conv_bn(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, int oup, int k = 3, int p = 1, int s = 1, int groups=1) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(input, oup, DimsHW{k, k}, weightMap[lname + "_conv2d_weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
conv1->setNbGroups(groups);
auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "_batchnorm", 1e-3);
assert(bn1);
return bn1;
}
ILayer* DepthWise(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, int inp, int oup, int groups, int s) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(input, groups, DimsHW{1, 1}, weightMap[lname + "_conv_sep_conv2d_weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{1, 1});
conv1->setPaddingNd(DimsHW{0, 0});
conv1->setNbGroups(1);
auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "_conv_sep_batchnorm", 1e-3);
assert(bn1);
auto act1 = addPRelu(network, weightMap, *bn1->getOutput(0), lname + "_conv_sep_relu");
assert(act1);
IConvolutionLayer* conv2 = network->addConvolutionNd(*act1->getOutput(0), groups, DimsHW{3, 3}, weightMap[lname + "_conv_dw_conv2d_weight"], emptywts);
assert(conv2);
conv2->setStrideNd(DimsHW{s, s});
conv2->setPaddingNd(DimsHW{1, 1});
conv2->setNbGroups(groups);
auto bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "_conv_dw_batchnorm", 1e-3);
assert(bn2);
auto act2 = addPRelu(network, weightMap, *bn2->getOutput(0), lname + "_conv_dw_relu");
assert(act2);
IConvolutionLayer* conv3 = network->addConvolutionNd(*act2->getOutput(0), oup, DimsHW{1, 1}, weightMap[lname + "_conv_proj_conv2d_weight"], emptywts);
assert(conv3);
conv3->setStrideNd(DimsHW{1, 1});
conv3->setPaddingNd(DimsHW{0, 0});
conv3->setNbGroups(1);
auto bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + "_conv_proj_batchnorm", 1e-3);
assert(bn3);
return bn3;
}
ILayer* DWResidual(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, int inp, int oup, int groups, int s) {
auto dw1 = DepthWise(network, weightMap, input, lname, inp, oup, groups, s);
IElementWiseLayer* ew1;
ew1 = network->addElementWise(input, *dw1->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
return ew1;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../arcface-mobilefacenet.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
auto conv_1 = conv_bn_relu(network, weightMap, *data, "conv_1", 64, 3, 1, 2);
auto conv_2_dw = conv_bn_relu(network, weightMap, *conv_1->getOutput(0), "conv_2_dw", 64, 3, 1, 1, 64);
auto conv_23 = DepthWise(network, weightMap, *conv_2_dw->getOutput(0), "dconv_23", 64, 64, 128, 2);
auto res_3_block0 = DWResidual(network, weightMap, *conv_23->getOutput(0), "res_3_block0", 64, 64, 128, 1);
auto res_3_block1 = DWResidual(network, weightMap, *res_3_block0->getOutput(0), "res_3_block1", 64, 64, 128, 1);
auto res_3_block2 = DWResidual(network, weightMap, *res_3_block1->getOutput(0), "res_3_block2", 64, 64, 128, 1);
auto res_3_block3 = DWResidual(network, weightMap, *res_3_block2->getOutput(0), "res_3_block3", 64, 64, 128, 1);
auto conv_34 = DepthWise(network, weightMap, *res_3_block3->getOutput(0), "dconv_34", 64, 128, 256, 2);
auto res_4_block0 = DWResidual(network, weightMap, *conv_34->getOutput(0), "res_4_block0", 128, 128, 256, 1);
auto res_4_block1 = DWResidual(network, weightMap, *res_4_block0->getOutput(0), "res_4_block1", 128, 128, 256, 1);
auto res_4_block2 = DWResidual(network, weightMap, *res_4_block1->getOutput(0), "res_4_block2", 128, 128, 256, 1);
auto res_4_block3 = DWResidual(network, weightMap, *res_4_block2->getOutput(0), "res_4_block3", 128, 128, 256, 1);
auto res_4_block4 = DWResidual(network, weightMap, *res_4_block3->getOutput(0), "res_4_block4", 128, 128, 256, 1);
auto res_4_block5 = DWResidual(network, weightMap, *res_4_block4->getOutput(0), "res_4_block5", 128, 128, 256, 1);
auto conv_45 = DepthWise(network, weightMap, *res_4_block5->getOutput(0), "dconv_45", 128, 128, 512, 2);
auto res_5_block0 = DWResidual(network, weightMap, *conv_45->getOutput(0), "res_5_block0", 128, 128, 256, 1);
auto res_5_block1 = DWResidual(network, weightMap, *res_5_block0->getOutput(0), "res_5_block1", 128, 128, 256, 1);
auto conv_6_sep = conv_bn_relu(network, weightMap, *res_5_block1->getOutput(0), "conv_6sep", 512, 1, 0, 1);
auto conv_6dw7_7 = conv_bn(network, weightMap, *conv_6_sep->getOutput(0), "conv_6dw7_7", 512, 7, 0, 1, 512);
IFullyConnectedLayer* fc1 = network->addFullyConnected(*conv_6dw7_7->getOutput(0), 128, weightMap["fc1_weight"], weightMap["pre_fc1_bias"]);
assert(fc1);
auto bn1 = addBatchNorm2d(network, weightMap, *fc1->getOutput(0), "fc1", 2e-5);
assert(bn1);
bn1->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*bn1->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(BATCH_SIZE, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("arcface-mobilefacenet.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (argc == 2 && std::string(argv[1]) == "-d") {
std::ifstream file("arcface-mobilefacenet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./arcface-mobilefacenet -s // serialize model to plan file" << std::endl;
std::cerr << "./arcface-mobilefacenet -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
// data[i] = 1.0;
static float prob[BATCH_SIZE * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
cv::Mat img = cv::imread("../joey0.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out(128, 1, CV_32FC1, prob);
cv::Mat out_norm;
cv::normalize(out, out_norm);
img = cv::imread("../joey1.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out1(1, 128, CV_32FC1, prob);
cv::Mat out_norm1;
cv::normalize(out1, out_norm1);
cv::Mat res = out_norm1 * out_norm;
std::cout << "similarity score: " << *(float*)res.data << std::endl;
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
//Print histogram of the output distribution
//std::cout << "\nOutput:\n\n";
//for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
//{
// std::cout << p_out_norm[i] << ", ";
// if (i % 10 == 0) std::cout << i / 10 << std::endl;
//}
//std::cout << std::endl;
return 0;
}

455
arcface/arcface-r100.cpp Normal file
View File

@ -0,0 +1,455 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include <dirent.h>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
//#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define BATCH_SIZE 1 // currently, only support BATCH=1
using namespace nvinfer1;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 112;
static const int INPUT_W = 112;
static const int OUTPUT_SIZE = 512;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
float *beta = (float*)weightMap[lname + "_beta"].values;
float *mean = (float*)weightMap[lname + "_moving_mean"].values;
float *var = (float*)weightMap[lname + "_moving_var"].values;
int len = weightMap[lname + "_moving_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* addPRelu(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
int len = weightMap[lname + "_gamma"].count;
float *scval_1 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
float *scval_2 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval_1[i] = -1.0;
scval_2[i] = -gamma[i];
}
Weights scale_1{ DataType::kFLOAT, scval_1, len };
Weights scale_2{ DataType::kFLOAT, scval_2, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = 0.0;
}
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
auto relu1 = network->addActivation(input, ActivationType::kRELU);
assert(relu1);
IScaleLayer* scale1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale_1, power);
assert(scale1);
auto relu2 = network->addActivation(*scale1->getOutput(0), ActivationType::kRELU);
assert(relu2);
IScaleLayer* scale2 = network->addScale(*relu2->getOutput(0), ScaleMode::kCHANNEL, shift, scale_2, power);
assert(scale2);
IElementWiseLayer* ew1 = network->addElementWise(*relu1->getOutput(0), *scale2->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
return ew1;
}
ILayer* resUnit(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int num_filters, int s, bool dim_match, std::string lname) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
auto bn1 = addBatchNorm2d(network, weightMap, input, lname + "_bn1", 2e-5);
IConvolutionLayer* conv1 = network->addConvolutionNd(*bn1->getOutput(0), num_filters, DimsHW{3, 3}, weightMap[lname + "_conv1_weight"], emptywts);
assert(conv1);
conv1->setPaddingNd(DimsHW{1, 1});
auto bn2 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "_bn2", 2e-5);
auto act1 = addPRelu(network, weightMap, *bn2->getOutput(0), lname + "_relu1");
IConvolutionLayer* conv2 = network->addConvolutionNd(*act1->getOutput(0), num_filters, DimsHW{3, 3}, weightMap[lname + "_conv2_weight"], emptywts);
assert(conv2);
conv2->setStrideNd(DimsHW{s, s});
conv2->setPaddingNd(DimsHW{1, 1});
auto bn3 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "_bn3", 2e-5);
IElementWiseLayer* ew1;
if (dim_match) {
ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM);
} else {
IConvolutionLayer* conv1sc = network->addConvolutionNd(input, num_filters, DimsHW{1, 1}, weightMap[lname + "_conv1sc_weight"], emptywts);
assert(conv1sc);
conv1sc->setStrideNd(DimsHW{s, s});
auto bn1sc = addBatchNorm2d(network, weightMap, *conv1sc->getOutput(0), lname + "_sc", 2e-5);
ew1 = network->addElementWise(*bn1sc->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
}
assert(ew1);
return ew1;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../arcface-r100.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv0 = network->addConvolutionNd(*data, 64, DimsHW{3, 3}, weightMap["conv0_weight"], emptywts);
assert(conv0);
conv0->setPaddingNd(DimsHW{1, 1});
auto bn0 = addBatchNorm2d(network, weightMap, *conv0->getOutput(0), "bn0", 2e-5);
auto relu0 = addPRelu(network, weightMap, *bn0->getOutput(0), "relu0");
auto s1u1 = resUnit(network, weightMap, *relu0->getOutput(0), 64, 2, false, "stage1_unit1");
auto s1u2 = resUnit(network, weightMap, *s1u1->getOutput(0), 64, 1, true, "stage1_unit2");
auto s1u3 = resUnit(network, weightMap, *s1u2->getOutput(0), 64, 1, true, "stage1_unit3");
auto s2u1 = resUnit(network, weightMap, *s1u3->getOutput(0), 128, 2, false, "stage2_unit1");
auto s2u2 = resUnit(network, weightMap, *s2u1->getOutput(0), 128, 1, true, "stage2_unit2");
auto s2u3 = resUnit(network, weightMap, *s2u2->getOutput(0), 128, 1, true, "stage2_unit3");
auto s2u4 = resUnit(network, weightMap, *s2u3->getOutput(0), 128, 1, true, "stage2_unit4");
auto s2u5 = resUnit(network, weightMap, *s2u4->getOutput(0), 128, 1, true, "stage2_unit5");
auto s2u6 = resUnit(network, weightMap, *s2u5->getOutput(0), 128, 1, true, "stage2_unit6");
auto s2u7 = resUnit(network, weightMap, *s2u6->getOutput(0), 128, 1, true, "stage2_unit7");
auto s2u8 = resUnit(network, weightMap, *s2u7->getOutput(0), 128, 1, true, "stage2_unit8");
auto s2u9 = resUnit(network, weightMap, *s2u8->getOutput(0), 128, 1, true, "stage2_unit9");
auto s2u10 = resUnit(network, weightMap, *s2u9->getOutput(0), 128, 1, true, "stage2_unit10");
auto s2u11 = resUnit(network, weightMap, *s2u10->getOutput(0), 128, 1, true, "stage2_unit11");
auto s2u12 = resUnit(network, weightMap, *s2u11->getOutput(0), 128, 1, true, "stage2_unit12");
auto s2u13 = resUnit(network, weightMap, *s2u12->getOutput(0), 128, 1, true, "stage2_unit13");
auto s3u1 = resUnit(network, weightMap, *s2u13->getOutput(0), 256, 2, false, "stage3_unit1");
auto s3u2 = resUnit(network, weightMap, *s3u1->getOutput(0), 256, 1, true, "stage3_unit2");
auto s3u3 = resUnit(network, weightMap, *s3u2->getOutput(0), 256, 1, true, "stage3_unit3");
auto s3u4 = resUnit(network, weightMap, *s3u3->getOutput(0), 256, 1, true, "stage3_unit4");
auto s3u5 = resUnit(network, weightMap, *s3u4->getOutput(0), 256, 1, true, "stage3_unit5");
auto s3u6 = resUnit(network, weightMap, *s3u5->getOutput(0), 256, 1, true, "stage3_unit6");
auto s3u7 = resUnit(network, weightMap, *s3u6->getOutput(0), 256, 1, true, "stage3_unit7");
auto s3u8 = resUnit(network, weightMap, *s3u7->getOutput(0), 256, 1, true, "stage3_unit8");
auto s3u9 = resUnit(network, weightMap, *s3u8->getOutput(0), 256, 1, true, "stage3_unit9");
auto s3u10 = resUnit(network, weightMap, *s3u9->getOutput(0), 256, 1, true, "stage3_unit10");
auto s3u11 = resUnit(network, weightMap, *s3u10->getOutput(0), 256, 1, true, "stage3_unit11");
auto s3u12 = resUnit(network, weightMap, *s3u11->getOutput(0), 256, 1, true, "stage3_unit12");
auto s3u13 = resUnit(network, weightMap, *s3u12->getOutput(0), 256, 1, true, "stage3_unit13");
auto s3u14 = resUnit(network, weightMap, *s3u13->getOutput(0), 256, 1, true, "stage3_unit14");
auto s3u15 = resUnit(network, weightMap, *s3u14->getOutput(0), 256, 1, true, "stage3_unit15");
auto s3u16 = resUnit(network, weightMap, *s3u15->getOutput(0), 256, 1, true, "stage3_unit16");
auto s3u17 = resUnit(network, weightMap, *s3u16->getOutput(0), 256, 1, true, "stage3_unit17");
auto s3u18 = resUnit(network, weightMap, *s3u17->getOutput(0), 256, 1, true, "stage3_unit18");
auto s3u19 = resUnit(network, weightMap, *s3u18->getOutput(0), 256, 1, true, "stage3_unit19");
auto s3u20 = resUnit(network, weightMap, *s3u19->getOutput(0), 256, 1, true, "stage3_unit20");
auto s3u21 = resUnit(network, weightMap, *s3u20->getOutput(0), 256, 1, true, "stage3_unit21");
auto s3u22 = resUnit(network, weightMap, *s3u21->getOutput(0), 256, 1, true, "stage3_unit22");
auto s3u23 = resUnit(network, weightMap, *s3u22->getOutput(0), 256, 1, true, "stage3_unit23");
auto s3u24 = resUnit(network, weightMap, *s3u23->getOutput(0), 256, 1, true, "stage3_unit24");
auto s3u25 = resUnit(network, weightMap, *s3u24->getOutput(0), 256, 1, true, "stage3_unit25");
auto s3u26 = resUnit(network, weightMap, *s3u25->getOutput(0), 256, 1, true, "stage3_unit26");
auto s3u27 = resUnit(network, weightMap, *s3u26->getOutput(0), 256, 1, true, "stage3_unit27");
auto s3u28 = resUnit(network, weightMap, *s3u27->getOutput(0), 256, 1, true, "stage3_unit28");
auto s3u29 = resUnit(network, weightMap, *s3u28->getOutput(0), 256, 1, true, "stage3_unit29");
auto s3u30 = resUnit(network, weightMap, *s3u29->getOutput(0), 256, 1, true, "stage3_unit30");
auto s4u1 = resUnit(network, weightMap, *s3u30->getOutput(0), 512, 2, false, "stage4_unit1");
auto s4u2 = resUnit(network, weightMap, *s4u1->getOutput(0), 512, 1, true, "stage4_unit2");
auto s4u3 = resUnit(network, weightMap, *s4u2->getOutput(0), 512, 1, true, "stage4_unit3");
auto bn1 = addBatchNorm2d(network, weightMap, *s4u3->getOutput(0), "bn1", 2e-5);
IFullyConnectedLayer* fc1 = network->addFullyConnected(*bn1->getOutput(0), 512, weightMap["pre_fc1_weight"], weightMap["pre_fc1_bias"]);
assert(fc1);
auto bn2 = addBatchNorm2d(network, weightMap, *fc1->getOutput(0), "fc1", 2e-5);
bn2->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*bn2->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(256, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("arcface-r100.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (argc == 2 && std::string(argv[1]) == "-d") {
std::ifstream file("arcface-r100.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./arcface-r100 -s // serialize model to plan file" << std::endl;
std::cerr << "./arcface-r100 -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
// data[i] = 1.0;
static float prob[BATCH_SIZE * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
cv::Mat img = cv::imread("../joey0.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out(512, 1, CV_32FC1, prob);
cv::Mat out_norm;
cv::normalize(out, out_norm);
img = cv::imread("../joey1.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out1(1, 512, CV_32FC1, prob);
cv::Mat out_norm1;
cv::normalize(out1, out_norm1);
cv::Mat res = out_norm1 * out_norm;
std::cout << "similarity score: " << *(float*)res.data << std::endl;
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
//Print histogram of the output distribution
//std::cout << "\nOutput:\n\n";
//for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
//{
// std::cout << p_out_norm[i] << ", ";
// if (i % 10 == 0) std::cout << i / 10 << std::endl;
//}
//std::cout << std::endl;
return 0;
}

426
arcface/arcface-r50.cpp Normal file
View File

@ -0,0 +1,426 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include <dirent.h>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
//#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define BATCH_SIZE 1 // currently, only support BATCH=1
using namespace nvinfer1;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 112;
static const int INPUT_W = 112;
static const int OUTPUT_SIZE = 512;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
float *beta = (float*)weightMap[lname + "_beta"].values;
float *mean = (float*)weightMap[lname + "_moving_mean"].values;
float *var = (float*)weightMap[lname + "_moving_var"].values;
int len = weightMap[lname + "_moving_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* addPRelu(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname) {
float *gamma = (float*)weightMap[lname + "_gamma"].values;
int len = weightMap[lname + "_gamma"].count;
float *scval_1 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
float *scval_2 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval_1[i] = -1.0;
scval_2[i] = -gamma[i];
}
Weights scale_1{ DataType::kFLOAT, scval_1, len };
Weights scale_2{ DataType::kFLOAT, scval_2, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = 0.0;
}
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
auto relu1 = network->addActivation(input, ActivationType::kRELU);
assert(relu1);
IScaleLayer* scale1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale_1, power);
assert(scale1);
auto relu2 = network->addActivation(*scale1->getOutput(0), ActivationType::kRELU);
assert(relu2);
IScaleLayer* scale2 = network->addScale(*relu2->getOutput(0), ScaleMode::kCHANNEL, shift, scale_2, power);
assert(scale2);
IElementWiseLayer* ew1 = network->addElementWise(*relu1->getOutput(0), *scale2->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
return ew1;
}
ILayer* resUnit(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int num_filters, int s, bool dim_match, std::string lname) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
auto bn1 = addBatchNorm2d(network, weightMap, input, lname + "_bn1", 2e-5);
IConvolutionLayer* conv1 = network->addConvolutionNd(*bn1->getOutput(0), num_filters, DimsHW{3, 3}, weightMap[lname + "_conv1_weight"], emptywts);
assert(conv1);
conv1->setPaddingNd(DimsHW{1, 1});
auto bn2 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "_bn2", 2e-5);
auto act1 = addPRelu(network, weightMap, *bn2->getOutput(0), lname + "_relu1");
IConvolutionLayer* conv2 = network->addConvolutionNd(*act1->getOutput(0), num_filters, DimsHW{3, 3}, weightMap[lname + "_conv2_weight"], emptywts);
assert(conv2);
conv2->setStrideNd(DimsHW{s, s});
conv2->setPaddingNd(DimsHW{1, 1});
auto bn3 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "_bn3", 2e-5);
IElementWiseLayer* ew1;
if (dim_match) {
ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM);
} else {
IConvolutionLayer* conv1sc = network->addConvolutionNd(input, num_filters, DimsHW{1, 1}, weightMap[lname + "_conv1sc_weight"], emptywts);
assert(conv1sc);
conv1sc->setStrideNd(DimsHW{s, s});
auto bn1sc = addBatchNorm2d(network, weightMap, *conv1sc->getOutput(0), lname + "_sc", 2e-5);
ew1 = network->addElementWise(*bn1sc->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
}
assert(ew1);
return ew1;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../arcface-r50.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv0 = network->addConvolutionNd(*data, 64, DimsHW{3, 3}, weightMap["conv0_weight"], emptywts);
assert(conv0);
conv0->setPaddingNd(DimsHW{1, 1});
auto bn0 = addBatchNorm2d(network, weightMap, *conv0->getOutput(0), "bn0", 2e-5);
auto relu0 = addPRelu(network, weightMap, *bn0->getOutput(0), "relu0");
auto s1u1 = resUnit(network, weightMap, *relu0->getOutput(0), 64, 2, false, "stage1_unit1");
auto s1u2 = resUnit(network, weightMap, *s1u1->getOutput(0), 64, 1, true, "stage1_unit2");
auto s1u3 = resUnit(network, weightMap, *s1u2->getOutput(0), 64, 1, true, "stage1_unit3");
auto s2u1 = resUnit(network, weightMap, *s1u3->getOutput(0), 128, 2, false, "stage2_unit1");
auto s2u2 = resUnit(network, weightMap, *s2u1->getOutput(0), 128, 1, true, "stage2_unit2");
auto s2u3 = resUnit(network, weightMap, *s2u2->getOutput(0), 128, 1, true, "stage2_unit3");
auto s2u4 = resUnit(network, weightMap, *s2u3->getOutput(0), 128, 1, true, "stage2_unit4");
auto s3u1 = resUnit(network, weightMap, *s2u4->getOutput(0), 256, 2, false, "stage3_unit1");
auto s3u2 = resUnit(network, weightMap, *s3u1->getOutput(0), 256, 1, true, "stage3_unit2");
auto s3u3 = resUnit(network, weightMap, *s3u2->getOutput(0), 256, 1, true, "stage3_unit3");
auto s3u4 = resUnit(network, weightMap, *s3u3->getOutput(0), 256, 1, true, "stage3_unit4");
auto s3u5 = resUnit(network, weightMap, *s3u4->getOutput(0), 256, 1, true, "stage3_unit5");
auto s3u6 = resUnit(network, weightMap, *s3u5->getOutput(0), 256, 1, true, "stage3_unit6");
auto s3u7 = resUnit(network, weightMap, *s3u6->getOutput(0), 256, 1, true, "stage3_unit7");
auto s3u8 = resUnit(network, weightMap, *s3u7->getOutput(0), 256, 1, true, "stage3_unit8");
auto s3u9 = resUnit(network, weightMap, *s3u8->getOutput(0), 256, 1, true, "stage3_unit9");
auto s3u10 = resUnit(network, weightMap, *s3u9->getOutput(0), 256, 1, true, "stage3_unit10");
auto s3u11 = resUnit(network, weightMap, *s3u10->getOutput(0), 256, 1, true, "stage3_unit11");
auto s3u12 = resUnit(network, weightMap, *s3u11->getOutput(0), 256, 1, true, "stage3_unit12");
auto s3u13 = resUnit(network, weightMap, *s3u12->getOutput(0), 256, 1, true, "stage3_unit13");
auto s3u14 = resUnit(network, weightMap, *s3u13->getOutput(0), 256, 1, true, "stage3_unit14");
auto s4u1 = resUnit(network, weightMap, *s3u14->getOutput(0), 512, 2, false, "stage4_unit1");
auto s4u2 = resUnit(network, weightMap, *s4u1->getOutput(0), 512, 1, true, "stage4_unit2");
auto s4u3 = resUnit(network, weightMap, *s4u2->getOutput(0), 512, 1, true, "stage4_unit3");
auto bn1 = addBatchNorm2d(network, weightMap, *s4u3->getOutput(0), "bn1", 2e-5);
IFullyConnectedLayer* fc1 = network->addFullyConnected(*bn1->getOutput(0), 512, weightMap["pre_fc1_weight"], weightMap["pre_fc1_bias"]);
assert(fc1);
auto bn2 = addBatchNorm2d(network, weightMap, *fc1->getOutput(0), "fc1", 2e-5);
bn2->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*bn2->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(BATCH_SIZE, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("arcface-r50.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (argc == 2 && std::string(argv[1]) == "-d") {
std::ifstream file("arcface-r50.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./arcface-r50 -s // serialize model to plan file" << std::endl;
std::cerr << "./arcface-r50 -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
// data[i] = 1.0;
static float prob[BATCH_SIZE * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
cv::Mat img = cv::imread("../joey0.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out(512, 1, CV_32FC1, prob);
cv::Mat out_norm;
cv::normalize(out, out_norm);
img = cv::imread("../joey1.ppm");
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<cv::Vec3b>(i)[2] - 127.5) * 0.0078125;
data[i + INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[1] - 127.5) * 0.0078125;
data[i + 2 * INPUT_H * INPUT_W] = ((float)img.at<cv::Vec3b>(i)[0] - 127.5) * 0.0078125;
}
// Run inference
start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
cv::Mat out1(1, 512, CV_32FC1, prob);
cv::Mat out_norm1;
cv::normalize(out1, out_norm1);
cv::Mat res = out_norm1 * out_norm;
std::cout << "similarity score: " << *(float*)res.data << std::endl;
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
//Print histogram of the output distribution
//std::cout << "\nOutput:\n\n";
//for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
//{
// std::cout << p_out_norm[i] << ", ";
// if (i % 10 == 0) std::cout << i / 10 << std::endl;
//}
//std::cout << std::endl;
return 0;
}

37
arcface/gen_wts.py Normal file
View File

@ -0,0 +1,37 @@
import struct
import sys
import argparse
import face_model
import cv2
import numpy as np
parser = argparse.ArgumentParser(description='face model test')
# general
parser.add_argument('--image-size', default='112,112', help='')
parser.add_argument('--model', default='model-r100-ii/model,0', help='path to load model.')
parser.add_argument('--ga-model', default='', help='path to load model.')
parser.add_argument('--gpu', default=0, type=int, help='gpu id')
parser.add_argument('--det', default=0, type=int, help='mtcnn option, 1 means using R+O, 0 means detect from begining')
parser.add_argument('--flip', default=0, type=int, help='whether do lr flip aug')
parser.add_argument('--threshold', default=1.24, type=float, help='ver dist threshold')
args = parser.parse_args()
model = face_model.FaceModel(args)
f = open('arcface-r100.wts', 'w')
f.write('{}\n'.format(len(model.model.get_params()[0].keys()) + len(model.model.get_params()[1].keys())))
for k, v in model.model.get_params()[0].items():
vr = v.reshape(-1).asnumpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')
for k, v in model.model.get_params()[1].items():
vr = v.reshape(-1).asnumpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')

510
arcface/logging.h Normal file
View File

@ -0,0 +1,510 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
#include "macros.h"
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#else
#define TRT_NOEXCEPT
#endif
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) TRT_NOEXCEPT override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

12
arcface/macros.h Normal file
View File

@ -0,0 +1,12 @@
#ifndef __MACROS_H
#define __MACROS_H
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif
#endif // __MACROS_H

210
arcface/prelu.cu Normal file
View File

@ -0,0 +1,210 @@
#include <cmath>
#include <stdio.h>
#include <cassert>
#include <iostream>
#include "prelu.h"
namespace nvinfer1
{
PReluPlugin::PReluPlugin(const std::vector<float>& gamma) : gamma_(gamma)
{
}
PReluPlugin::~PReluPlugin()
{
}
// create the plugin at runtime from a byte stream
PReluPlugin::PReluPlugin(const void* data, size_t length)
{
char *p = (char*)data;
input_size_ = reinterpret_cast<const int*>(p)[0];
p += sizeof(int);
gamma_.assign((float*)p, (float*)p + (length - sizeof(int)) / sizeof(float));
}
void PReluPlugin::serialize(void* buffer) const TRT_NOEXCEPT
{
*reinterpret_cast<int*>(buffer) = input_size_;
char *p = reinterpret_cast<char*>(buffer);
p += sizeof(int);
memcpy(p, gamma_.data(), gamma_.size() * sizeof(float));
}
size_t PReluPlugin::getSerializationSize() const TRT_NOEXCEPT
{
return sizeof(input_size_) + gamma_.size() * sizeof(float);
}
int PReluPlugin::initialize() TRT_NOEXCEPT
{
return 0;
}
Dims PReluPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT
{
assert(nbInputDims == 1);
assert(index == 0);
input_size_ = inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2];
// Output dimensions
return Dims3(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]);
}
// Set plugin namespace
void PReluPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT
{
mPluginNamespace = pluginNamespace;
}
const char* PReluPlugin::getPluginNamespace() const TRT_NOEXCEPT
{
return mPluginNamespace;
}
// Return the DataType of the plugin output at the requested index
DataType PReluPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT
{
return DataType::kFLOAT;
}
// Return true if output tensor is broadcast across a batch.
bool PReluPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT
{
return false;
}
// Return true if plugin can use input that is broadcast across batch without replication.
bool PReluPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT
{
return false;
}
void PReluPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT
{
}
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void PReluPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT
{
}
// Detach the plugin object from its execution context.
void PReluPlugin::detachFromContext() TRT_NOEXCEPT {}
const char* PReluPlugin::getPluginType() const TRT_NOEXCEPT
{
return "PRelu_TRT";
}
const char* PReluPlugin::getPluginVersion() const TRT_NOEXCEPT
{
return "1";
}
void PReluPlugin::destroy() TRT_NOEXCEPT
{
delete this;
}
// Clone the plugin
IPluginV2IOExt* PReluPlugin::clone() const TRT_NOEXCEPT
{
PReluPlugin *p = new PReluPlugin(gamma_);
p->input_size_ = input_size_;
p->setPluginNamespace(mPluginNamespace);
return p;
}
__global__ void prelu_kernel(const float *input, float *output, int num_elem, int input_size, int fm_size, const float* gamma) {
int idx = threadIdx.x + blockDim.x * blockIdx.x;
if (idx >= num_elem) return;
if (input[idx] >= 0.0f) {
output[idx] = input[idx];
return;
}
int c = (idx % input_size) / fm_size;
output[idx] = input[idx] * gamma[c];
}
void PReluPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) {
int block_size = thread_count_;
int grid_size = (input_size_ * batchSize + block_size - 1) / block_size;
void *dev_gamma;
assert(cudaMalloc(&dev_gamma, sizeof(float) * gamma_.size()) == cudaSuccess);
assert(cudaMemcpy(dev_gamma, gamma_.data(), sizeof(float) * gamma_.size(), cudaMemcpyHostToDevice) == cudaSuccess);
prelu_kernel<<<grid_size, block_size>>>(inputs[0], output, input_size_ * batchSize, input_size_, input_size_ / gamma_.size(), (const float*)dev_gamma);
assert(cudaFree(dev_gamma) == cudaSuccess);
}
int PReluPlugin::enqueue(int batchSize, const void*const * inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT
{
//assert(batchSize == 1);
//GPU
//CUDA_CHECK(cudaStreamSynchronize(stream));
forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize);
return 0;
}
PluginFieldCollection PReluPluginCreator::mFC{};
std::vector<PluginField> PReluPluginCreator::mPluginAttributes;
PReluPluginCreator::PReluPluginCreator()
{
mPluginAttributes.emplace_back(PluginField("gamma", nullptr, PluginFieldType::kFLOAT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
const char* PReluPluginCreator::getPluginName() const TRT_NOEXCEPT
{
return "PRelu_TRT";
}
const char* PReluPluginCreator::getPluginVersion() const TRT_NOEXCEPT
{
return "1";
}
const PluginFieldCollection* PReluPluginCreator::getFieldNames() TRT_NOEXCEPT
{
return &mFC;
}
IPluginV2IOExt* PReluPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT
{
std::vector<float> gamma;
const PluginField* fields = fc->fields;
for (int i = 0; i < fc->nbFields; ++i) {
const char* attrName = fields[i].name;
if (!strcmp(attrName, "gamma")) {
assert(fields[i].type == PluginFieldType::kFLOAT32);
int size = fields[i].length;
gamma.reserve(size);
const auto* w = static_cast<const float*>(fields[i].data);
for (int j = 0; j < size; j++)
{
gamma.push_back(*w);
w++;
}
}
}
PReluPlugin* obj = new PReluPlugin(gamma);
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}
IPluginV2IOExt* PReluPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT
{
// This object will be deleted when the network is destroyed, which will
// call PReluPlugin::destroy()
PReluPlugin* obj = new PReluPlugin(serialData, serialLength);
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}
}

108
arcface/prelu.h Normal file
View File

@ -0,0 +1,108 @@
#ifndef _PRELU_PLUGIN_H
#define _PRELU_PLUGIN_H
#include <string>
#include <vector>
#include "NvInfer.h"
#include "macros.h"
namespace nvinfer1
{
class PReluPlugin: public IPluginV2IOExt
{
public:
PReluPlugin(const std::vector<float>& gamma);
PReluPlugin(const void* data, size_t length);
~PReluPlugin();
int getNbOutputs() const TRT_NOEXCEPT override
{
return 1;
}
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override;
int initialize() TRT_NOEXCEPT override;
virtual void terminate() TRT_NOEXCEPT override {};
virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0;}
virtual int enqueue(int batchSize, const void*const * inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override;
virtual size_t getSerializationSize() const TRT_NOEXCEPT override;
virtual void serialize(void* buffer) const TRT_NOEXCEPT override;
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override {
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
}
const char* getPluginType() const TRT_NOEXCEPT override;
const char* getPluginVersion() const TRT_NOEXCEPT override;
void destroy() TRT_NOEXCEPT override;
IPluginV2IOExt* clone() const TRT_NOEXCEPT override;
void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override;
const char* getPluginNamespace() const TRT_NOEXCEPT override;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override;
void attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override;
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override;
void detachFromContext() TRT_NOEXCEPT override;
int input_size_;
private:
void forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize = 1);
int thread_count_ = 256;
std::vector<float> gamma_;
const char* mPluginNamespace;
};
class PReluPluginCreator : public IPluginCreator
{
public:
PReluPluginCreator();
~PReluPluginCreator() override = default;
const char* getPluginName() const TRT_NOEXCEPT override;
const char* getPluginVersion() const TRT_NOEXCEPT override;
const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override;
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override;
void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override
{
mNamespace = libNamespace;
}
const char* getPluginNamespace() const TRT_NOEXCEPT override
{
return mNamespace.c_str();
}
private:
std::string mNamespace;
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
};
};
#endif

29
centernet/README.md Normal file
View File

@ -0,0 +1,29 @@
# CenterNet
This is the trt implementation of detection model [ctdet_coco_dla_2x](https://drive.google.com/open?id=1pl_-ael8wERdUREEnaIfqOV_VF2bEVRT) from [xingyizhou/CenterNet](https://github.com/xingyizhou/CenterNet) official work.
## How to Run
1. Follow [NVIDIA/TensorRT](https://github.com/NVIDIA/TensorRT) tutorial to build TensorRT7
2. Copy folder `dcnv2Plugin` to `TensorRT/plugin` and edit `InferPlugin.cpp` and `CMakeLists.txt`
3. Rebuild to install custom plugin
4. Use `tensorrt-7.2.3.4-cp36-none-linux_x86_64.whl` in TensorRT OSS to update your python-tensorrt
5. Run `python centernet.py -m ${PTH_PATH} -s` to create trt engine
## Sample
```
// Download ctdet_coco_dla_2x.pth and transfer it into trt engine first
// Download the test img from https://raw.githubusercontent.com/tensorflow/models/master/research/deeplab/g3doc/img/image2.jpg or choose your own one
cd sample
python test.py ${ENGINE_PATH} ${IMG_PATH}
```
![trt_out](https://user-images.githubusercontent.com/47047345/119128637-7a878900-ba68-11eb-91ff-5dcc10f01b77.jpg)
## TODO
Integrate the post process with trt engine to make it more easier to use.

333
centernet/centernet.py Normal file
View File

@ -0,0 +1,333 @@
import numpy as np
import tensorrt as trt
import torch
from sample import common
import argparse
import time
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
for plugin_creator in PLUGIN_CREATORS:
if plugin_creator.name == 'DCNv2_TRT':
dcnCreator = plugin_creator
class ModelData(object):
INPUT_NAME = "data"
INPUT_SHAPE = (3, 512, 512)
OUTPUT_NAME = "prob"
DTYPE = trt.float16
class Centernet_dla34(object):
def __init__(self, weights) -> None:
super().__init__()
self.weights = weights
self.levels = [1, 1, 1, 2, 2, 1]
self.channels = [16, 32, 64, 128, 256, 512]
self.down_ratio = 4
self.last_level = 5
self.engine = self.build_engine()
def add_batchnorm_2d(self, input_tensor, parent):
gamma = self.weights[parent + '.weight'].numpy()
beta = self.weights[parent + '.bias'].numpy()
mean = self.weights[parent + '.running_mean'].numpy()
var = self.weights[parent + '.running_var'].numpy()
eps = 1e-5
scale = gamma / np.sqrt(var + eps)
shift = beta - mean * gamma / np.sqrt(var + eps)
power = np.ones_like(scale)
return self.network.add_scale(input=input_tensor.get_output(0), mode=trt.ScaleMode.CHANNEL, shift=shift, scale=scale, power=power)
def add_basic_block(self, input_tensor, out_channels, residual=None, stride=1, dilation=1, parent=''):
conv1_w = self.weights[parent + '.conv1.weight'].numpy()
conv1 = self.network.add_convolution(input=input_tensor.get_output(
0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv1_w)
conv1.stride = (stride, stride)
conv1.padding = (dilation, dilation)
conv1.dilation = (dilation, dilation)
bn1 = self.add_batchnorm_2d(conv1, parent + '.bn1')
ac1 = self.network.add_activation(
input=bn1.get_output(0), type=trt.ActivationType.RELU)
conv2_w = self.weights[parent + '.conv2.weight'].numpy()
conv2 = self.network.add_convolution(input=ac1.get_output(
0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv2_w)
conv2.padding = (dilation, dilation)
conv2.dilation = (dilation, dilation)
out = self.add_batchnorm_2d(conv2, parent + '.bn2')
if residual is None:
out = self.network.add_elementwise(input_tensor.get_output(
0), out.get_output(0), trt.ElementWiseOperation.SUM)
else:
out = self.network.add_elementwise(residual.get_output(
0), out.get_output(0), trt.ElementWiseOperation.SUM)
return self.network.add_activation(input=out.get_output(0), type=trt.ActivationType.RELU)
def add_level(self, input_tensor, out_channels, stride=1, dilation=1, parent=''):
conv1_w = self.weights[parent + '.0.weight'].numpy()
conv1 = self.network.add_convolution(input=input_tensor.get_output(
0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv1_w)
conv1.stride = (stride, stride)
conv1.padding = (dilation, dilation)
conv1.dilation = (dilation, dilation)
bn1 = self.add_batchnorm_2d(conv1, parent + '.1')
ac1 = self.network.add_activation(
input=bn1.get_output(0), type=trt.ActivationType.RELU)
return ac1
def add_root(self, input_tensors: list, out_channels, kernel_size=1, residual=False, parent=''):
ct = self.network.add_concatenation(
[x.get_output(0) for x in input_tensors])
conv_w = self.weights[parent + '.conv.weight'].numpy()
conv = self.network.add_convolution(input=ct.get_output(
0), num_output_maps=out_channels, kernel_shape=(1, 1), kernel=conv_w)
conv.padding = ((kernel_size - 1) // 2, (kernel_size - 1) // 2)
bn1 = self.add_batchnorm_2d(conv, parent + '.bn')
out = self.network.add_activation(
input=bn1.get_output(0), type=trt.ActivationType.RELU)
if residual:
out = self.network.add_elementwise(input_tensors[0].get_output(
0), out.get_output(0), trt.ElementWiseOperation.SUM)
return self.network.add_activation(input=out.get_output(0), type=trt.ActivationType.RELU)
def add_tree(self, input_tensor, level, out_channels, residual=None, children=None, stride=1, level_root=False, parent=''):
children = [] if children is None else children
if stride > 1:
bottom = self.network.add_pooling(input_tensor.get_output(
0), trt.PoolingType.MAX, (stride, stride))
bottom.stride = (stride, stride)
else:
bottom = input_tensor
if input_tensor.get_output(0).shape[0] != out_channels:
project_conv1_w = self.weights[parent +
'.project.0.weight'].numpy()
project_conv1 = self.network.add_convolution(input=bottom.get_output(
0), num_output_maps=out_channels, kernel_shape=(1, 1), kernel=project_conv1_w)
residual = self.add_batchnorm_2d(
project_conv1, parent + '.project.1')
else:
residual = bottom
if level_root:
children.append(bottom)
if level == 1:
tree1 = self.add_basic_block(
input_tensor, out_channels, residual, stride, parent=parent+'.tree1')
tree2 = self.add_basic_block(
tree1, out_channels, parent=parent+'.tree2')
return self.add_root([tree2, tree1]+children, out_channels, parent=parent+'.root')
else:
tree1 = self.add_tree(input_tensor, level-1, out_channels,
residual, stride=stride, parent=parent+'.tree1')
children.append(tree1)
return self.add_tree(tree1, level-1, out_channels, children=children, parent=parent+'.tree2')
def add_base(self, input_tensor, parent):
base_conv1_w = self.weights[parent+'.base_layer.0.weight'].numpy()
base_conv1 = self.network.add_convolution(
input=input_tensor, num_output_maps=self.channels[0], kernel_shape=(7, 7), kernel=base_conv1_w)
base_conv1.padding = (3, 3)
base_bn1 = self.add_batchnorm_2d(base_conv1, parent+'.base_layer.1')
base_ac1 = self.network.add_activation(
input=base_bn1.get_output(0), type=trt.ActivationType.RELU)
level0 = self.add_level(
base_ac1, self.channels[0], parent=parent+'.level0')
level1 = self.add_level(
level0, self.channels[1], 2, parent=parent+'.level1')
level2 = self.add_tree(
level1, self.levels[2], self.channels[2], stride=2, level_root=False, parent=parent+'.level2')
level3 = self.add_tree(
level2, self.levels[3], self.channels[3], stride=2, level_root=True, parent=parent+'.level3')
level4 = self.add_tree(
level3, self.levels[4], self.channels[4], stride=2, level_root=True, parent=parent+'.level4')
level5 = self.add_tree(
level4, self.levels[5], self.channels[5], stride=2, level_root=True, parent=parent+'.level5')
return [level0, level1, level2, level3, level4, level5]
def add_deform_conv(self, input_tensor, out_channels, kernel=3, stride=1, padding=1, dilation=1, deformable_group=1, parent=''):
conv_offset_mask_w = self.weights[parent +
'.conv.conv_offset_mask.weight'].numpy()
conv_offset_mask_b = self.weights[parent +
'.conv.conv_offset_mask.bias'].numpy()
conv_offset_mask = self.network.add_convolution(input=input_tensor.get_output(0),
num_output_maps=deformable_group*3*kernel*kernel,
kernel_shape=(
kernel, kernel),
kernel=conv_offset_mask_w,
bias=conv_offset_mask_b)
conv_offset_mask.stride = (stride, stride)
conv_offset_mask.padding = (padding, padding)
out_channels = trt.PluginField("out_channels", np.array(
[out_channels], dtype=np.int32), trt.PluginFieldType.INT32)
kernel = trt.PluginField("kernel", np.array(
[kernel], dtype=np.int32), trt.PluginFieldType.INT32)
deformable_group = trt.PluginField("deformable_group", np.array(
[deformable_group], dtype=np.int32), trt.PluginFieldType.INT32)
dilation = trt.PluginField("dilation", np.array(
[dilation], dtype=np.int32), trt.PluginFieldType.INT32)
padding = trt.PluginField("padding", np.array(
[padding], dtype=np.int32), trt.PluginFieldType.INT32)
stride = trt.PluginField("stride", np.array(
[stride], dtype=np.int32), trt.PluginFieldType.INT32)
weight = trt.PluginField(
"weight", self.weights[parent + '.conv.weight'].numpy(), trt.PluginFieldType.FLOAT32)
bias = trt.PluginField(
"bias", self.weights[parent + '.conv.bias'].numpy(), trt.PluginFieldType.FLOAT32)
field_collection = trt.PluginFieldCollection(
[out_channels, kernel, deformable_group, dilation, padding, stride, weight, bias])
DCN = dcnCreator.create_plugin(
name='DCNv2_TRT', field_collection=field_collection)
sigmoid_conv_offset_mask = self.network.add_activation(
input=conv_offset_mask.get_output(0), type=trt.ActivationType.SIGMOID)
dcn = self.network.add_plugin_v2(inputs=[input_tensor.get_output(
0), conv_offset_mask.get_output(0), sigmoid_conv_offset_mask.get_output(0)], plugin=DCN)
bn = self.add_batchnorm_2d(dcn, parent+'.actf.0')
return self.network.add_activation(input=bn.get_output(0), type=trt.ActivationType.RELU)
def add_ida_up(self, input_tensors, out_channels, up_f, startp, parent):
for i in range(startp + 1, len(input_tensors)):
proj = self.add_deform_conv(
input_tensors[i], out_channels, parent=parent+'.proj_%d' % (i-startp))
f = up_f[i-startp]
up_w = self.weights[parent + '.up_%d.weight' % (i-startp)].numpy()
up = self.network.add_deconvolution(
proj.get_output(0), out_channels, (f*2, f*2), up_w)
up.stride = (f, f)
up.padding = (f//2, f//2)
up.num_groups = out_channels
node = self.network.add_elementwise(
input_tensors[i-1].get_output(0), up.get_output(0), trt.ElementWiseOperation.SUM)
input_tensors[i] = self.add_deform_conv(
node, out_channels, parent=parent+'.node_%d' % (i-startp))
return input_tensors
def add_dla_up(self, input_tensors, first_level, parent):
channels = self.channels[first_level:]
scales = [2 ** i for i in range(len(self.channels[first_level:]))]
scales = np.array(scales, dtype=int)
out = [input_tensors[-1]]
for i in range(len(channels) - 1):
j = -i - 2
input_tensors = self.add_ida_up(
input_tensors, channels[j], scales[j:] // scales[j], len(input_tensors) - i - 2, parent+'.ida_%d' % i)
out.insert(0, input_tensors[-1])
scales[j + 1:] = scales[j]
channels[j + 1:] = [channels[j] for _ in channels[j + 1:]]
return out
def add_head(self, input_tensor, out_channels, head, head_conv=256, final_kernal=1):
conv1_w = self.weights[head+'.0.weight'].numpy()
conv1_b = self.weights[head+'.0.bias'].numpy()
conv1 = self.network.add_convolution(
input_tensor.get_output(0), head_conv, (3, 3), conv1_w, conv1_b)
conv1.padding = (1, 1)
ac1 = self.network.add_activation(
input=conv1.get_output(0), type=trt.ActivationType.RELU)
conv2_w = self.weights[head + '.2.weight'].numpy()
conv2_b = self.weights[head+'.2.bias'].numpy()
conv2 = self.network.add_convolution(ac1.get_output(
0), out_channels, (final_kernal, final_kernal), conv2_w, conv2_b)
return conv2
def populate_network(self):
# Configure the network layers based on the self.weights provided.
input_tensor = self.network.add_input(
name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)
y = self.add_base(input_tensor, 'module.base')
first_level = int(np.log2(self.down_ratio))
last_level = self.last_level
dla_up = self.add_dla_up(y, first_level, 'module.dla_up')
ida_up = self.add_ida_up(dla_up[:last_level-first_level], self.channels[first_level], [
2 ** i for i in range(last_level - first_level)], 0, 'module.ida_up')
hm = self.add_head(ida_up[-1], 80, 'module.hm')
wh = self.add_head(ida_up[-1], 2, 'module.wh')
reg = self.add_head(ida_up[-1], 2, 'module.reg')
hm.get_output(0).name = 'hm'
wh.get_output(0).name = 'wh'
reg.get_output(0).name = 'reg'
self.network.mark_output(tensor=hm.get_output(0))
self.network.mark_output(tensor=wh.get_output(0))
self.network.mark_output(tensor=reg.get_output(0))
def build_engine(self):
# For more information on TRT basics, refer to the introductory samples.
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
self.network = network
builder.max_workspace_size = common.GiB(1)
builder.max_batch_size = 1
# Populate the network using self.weights from the PyTorch model.
self.populate_network()
# Build and return an engine.
return builder.build_cuda_engine(self.network)
def load_random_test_case(pagelocked_buffer):
# Select an image at random to be the test case.
img = np.random.randn(1, 3, 512, 512).astype(np.float32)
# Copy to the pagelocked input buffer
np.copyto(pagelocked_buffer, img.ravel())
return img
def main(args):
# Get the PyTorch weights
weights = torch.load(args.model, map_location={
'cuda:0': 'cpu'})['state_dict']
# Do inference with TensorRT.
with Centernet_dla34(weights).engine as engine:
if args.save_engine:
with open('centernet.engine', "wb") as f:
f.write(engine.serialize())
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
with engine.create_execution_context() as context:
img = load_random_test_case(pagelocked_buffer=inputs[0].host)
# For more information on performing inference, refer to the introductory samples.
# The common.do_inference function will return a list of outputs - we only have one in this case.
t = time.time()
[hm, wh, reg] = common.do_inference(
context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=1)
t = time.time() - t
print('output: hm:%f, wh:%f, reg:%f' %
(hm.mean(), wh.mean(), reg.mean()))
print(t)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='CenterNet dla34 ctdet')
parser.add_argument('--model', '-m', type=str,
default='./ctdet_coco_dla_2x.pth', help='path of pytorch .pth')
parser.add_argument('--save_engine', '-s',
action='store_true', help='if save trt engine')
args = parser.parse_args()
main(args)

21
centernet/dcnv2Plugin/CMakeLists.txt Normal file
View File

@ -0,0 +1,21 @@
#
# Copyright (c) 2020, 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)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
file(GLOB CU_SRCS *.cu)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} ${CU_SRCS})
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE)

399
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu Normal file
View File

@ -0,0 +1,399 @@
#include "dcn_v2_im2col_cuda.h"
#include <cstdio>
#include <algorithm>
#include <cstring>
#define CUDA_KERNEL_LOOP(i, n) \
for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
i < (n); \
i += blockDim.x * gridDim.x)
const int CUDA_NUM_THREADS = 512;
//inline int GET_BLOCKS(const int N)
//{
// return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
//}
dim3 GET_BLOCKS(uint n)
{
uint k = (n - 1) /CUDA_NUM_THREADS + 1;
uint x = k ;
uint y = 1 ;
if (x > 65535 )
{
x = ceil(sqrt(x));
y = (n - 1 )/(x*CUDA_NUM_THREADS) + 1;
}
dim3 d = {x,y,1} ;
return d;
}
__device__ float dmcn_im2col_bilinear(const float *bottom_data, const int data_width,
const int height, const int width, float h, float w)
{
int h_low = floor(h);
int w_low = floor(w);
int h_high = h_low + 1;
int w_high = w_low + 1;
float lh = h - h_low;
float lw = w - w_low;
float hh = 1 - lh, hw = 1 - lw;
float v1 = 0;
if (h_low >= 0 && w_low >= 0)
v1 = bottom_data[h_low * data_width + w_low];
float v2 = 0;
if (h_low >= 0 && w_high <= width - 1)
v2 = bottom_data[h_low * data_width + w_high];
float v3 = 0;
if (h_high <= height - 1 && w_low >= 0)
v3 = bottom_data[h_high * data_width + w_low];
float v4 = 0;
if (h_high <= height - 1 && w_high <= width - 1)
v4 = bottom_data[h_high * data_width + w_high];
float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
return val;
}
__device__ float dmcn_get_gradient_weight(float argmax_h, float argmax_w,
const int h, const int w, const int height, const int width)
{
if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
{
//empty
return 0;
}
int argmax_h_low = floor(argmax_h);
int argmax_w_low = floor(argmax_w);
int argmax_h_high = argmax_h_low + 1;
int argmax_w_high = argmax_w_low + 1;
float weight = 0;
if (h == argmax_h_low && w == argmax_w_low)
weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);
if (h == argmax_h_low && w == argmax_w_high)
weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);
if (h == argmax_h_high && w == argmax_w_low)
weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);
if (h == argmax_h_high && w == argmax_w_high)
weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);
return weight;
}
__device__ float dmcn_get_coordinate_weight(float argmax_h, float argmax_w,
const int height, const int width, const float *im_data,
const int data_width, const int bp_dir)
{
if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
{
//empty
return 0;
}
int argmax_h_low = floor(argmax_h);
int argmax_w_low = floor(argmax_w);
int argmax_h_high = argmax_h_low + 1;
int argmax_w_high = argmax_w_low + 1;
float weight = 0;
if (bp_dir == 0)
{
if (argmax_h_low >= 0 && argmax_w_low >= 0)
weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];
if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];
if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];
if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];
}
else if (bp_dir == 1)
{
if (argmax_h_low >= 0 && argmax_w_low >= 0)
weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];
if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];
if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];
if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];
}
return weight;
}
__global__ void modulated_deformable_im2col_gpu_kernel(const int n,
const float *data_im, const float *data_offset, const float *data_mask,
const int height, const int width, const int kernel_h, const int kernel_w,
const int pad_h, const int pad_w,
const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int channel_per_deformable_group,
const int batch_size, const int num_channels, const int deformable_group,
const int height_col, const int width_col,
float *data_col)
{
CUDA_KERNEL_LOOP(index, n)
{
// index index of output matrix
const int w_col = index % width_col;
const int h_col = (index / width_col) % height_col;
const int b_col = (index / width_col / height_col) % batch_size;
const int c_im = (index / width_col / height_col) / batch_size;
const int c_col = c_im * kernel_h * kernel_w;
// compute deformable group index
const int deformable_group_index = c_im / channel_per_deformable_group;
const int h_in = h_col * stride_h - pad_h;
const int w_in = w_col * stride_w - pad_w;
float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
//const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
for (int i = 0; i < kernel_h; ++i)
{
for (int j = 0; j < kernel_w; ++j)
{
const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
const float offset_h = data_offset_ptr[data_offset_h_ptr];
const float offset_w = data_offset_ptr[data_offset_w_ptr];
const float mask = data_mask_ptr[data_mask_hw_ptr];
float val = static_cast<float>(0);
const float h_im = h_in + i * dilation_h + offset_h;
const float w_im = w_in + j * dilation_w + offset_w;
//if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
{
//const float map_h = i * dilation_h + offset_h;
//const float map_w = j * dilation_w + offset_w;
//const int cur_height = height - h_in;
//const int cur_width = width - w_in;
//val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
}
*data_col_ptr = val * mask;
data_col_ptr += batch_size * height_col * width_col;
//data_col_ptr += height_col * width_col;
}
}
}
}
__global__ void modulated_deformable_col2im_gpu_kernel(const int n,
const float *data_col, const float *data_offset, const float *data_mask,
const int channels, const int height, const int width,
const int kernel_h, const int kernel_w,
const int pad_h, const int pad_w,
const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int channel_per_deformable_group,
const int batch_size, const int deformable_group,
const int height_col, const int width_col,
float *grad_im)
{
CUDA_KERNEL_LOOP(index, n)
{
const int j = (index / width_col / height_col / batch_size) % kernel_w;
const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;
const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;
// compute the start and end of the output
const int deformable_group_index = c / channel_per_deformable_group;
int w_out = index % width_col;
int h_out = (index / width_col) % height_col;
int b = (index / width_col / height_col) % batch_size;
int w_in = w_out * stride_w - pad_w;
int h_in = h_out * stride_h - pad_h;
const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;
const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;
const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out;
const float offset_h = data_offset_ptr[data_offset_h_ptr];
const float offset_w = data_offset_ptr[data_offset_w_ptr];
const float mask = data_mask_ptr[data_mask_hw_ptr];
const float cur_inv_h_data = h_in + i * dilation_h + offset_h;
const float cur_inv_w_data = w_in + j * dilation_w + offset_w;
const float cur_top_grad = data_col[index] * mask;
const int cur_h = (int)cur_inv_h_data;
const int cur_w = (int)cur_inv_w_data;
for (int dy = -2; dy <= 2; dy++)
{
for (int dx = -2; dx <= 2; dx++)
{
if (cur_h + dy >= 0 && cur_h + dy < height &&
cur_w + dx >= 0 && cur_w + dx < width &&
abs(cur_inv_h_data - (cur_h + dy)) < 1 &&
abs(cur_inv_w_data - (cur_w + dx)) < 1)
{
int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;
float weight = dmcn_get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);
atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);
}
}
}
}
}
__global__ void modulated_deformable_col2im_coord_gpu_kernel(const int n,
const float *data_col, const float *data_im,
const float *data_offset, const float *data_mask,
const int channels, const int height, const int width,
const int kernel_h, const int kernel_w,
const int pad_h, const int pad_w,
const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int channel_per_deformable_group,
const int batch_size, const int offset_channels, const int deformable_group,
const int height_col, const int width_col,
float *grad_offset, float *grad_mask)
{
CUDA_KERNEL_LOOP(index, n)
{
float val = 0, mval = 0;
int w = index % width_col;
int h = (index / width_col) % height_col;
int c = (index / width_col / height_col) % offset_channels;
int b = (index / width_col / height_col) / offset_channels;
// compute the start and end of the output
const int deformable_group_index = c / (2 * kernel_h * kernel_w);
const int col_step = kernel_h * kernel_w;
int cnt = 0;
const float *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col;
const float *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width;
const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;
for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)
{
const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;
const int bp_dir = offset_c % 2;
int j = (col_pos / width_col / height_col / batch_size) % kernel_w;
int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;
int w_out = col_pos % width_col;
int h_out = (col_pos / width_col) % height_col;
int w_in = w_out * stride_w - pad_w;
int h_in = h_out * stride_h - pad_h;
const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);
const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);
const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out);
const float offset_h = data_offset_ptr[data_offset_h_ptr];
const float offset_w = data_offset_ptr[data_offset_w_ptr];
const float mask = data_mask_ptr[data_mask_hw_ptr];
float inv_h = h_in + i * dilation_h + offset_h;
float inv_w = w_in + j * dilation_w + offset_w;
if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)
{
inv_h = inv_w = -2;
}
else
{
mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w);
}
const float weight = dmcn_get_coordinate_weight(
inv_h, inv_w,
height, width, data_im_ptr + cnt * height * width, width, bp_dir);
val += weight * data_col_ptr[col_pos] * mask;
cnt += 1;
}
// KERNEL_ASSIGN(grad_offset[index], offset_req, val);
grad_offset[index] = val;
if (offset_c % 2 == 0)
// KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval);
grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval;
}
}
void modulated_deformable_im2col_cuda(cudaStream_t stream,
const float* data_im, const float* data_offset, const float* data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group, float* data_col) {
// num_axes should be smaller than block size
const int channel_per_deformable_group = channels / deformable_group;
const int num_kernels = channels * batch_size * height_col * width_col;
modulated_deformable_im2col_gpu_kernel
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kenerl_w,
pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,
batch_size, channels, deformable_group, height_col, width_col, data_col);
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess)
{
printf("error in modulated_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
}
}
void modulated_deformable_col2im_cuda(cudaStream_t stream,
const float* data_col, const float* data_offset, const float* data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kernel_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group, float* grad_im){
const int channel_per_deformable_group = channels / deformable_group;
const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col;
modulated_deformable_col2im_gpu_kernel
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
num_kernels, data_col, data_offset, data_mask, channels, height_im, width_im,
kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w,
dilation_h, dilation_w, channel_per_deformable_group,
batch_size, deformable_group, height_col, width_col, grad_im);
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess)
{
printf("error in modulated_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
}
}
void modulated_deformable_col2im_coord_cuda(cudaStream_t stream,
const float* data_col, const float* data_im, const float* data_offset, const float* data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kernel_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group,
float* grad_offset, float* grad_mask) {
const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group;
const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group;
modulated_deformable_col2im_coord_gpu_kernel
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
num_kernels, data_col, data_im, data_offset, data_mask, channels, height_im, width_im,
kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
dilation_h, dilation_w, channel_per_deformable_group,
batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col,
grad_offset, grad_mask);
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess)
{
printf("error in modulated_deformable_col2im_coord_cuda: %s\n", cudaGetErrorString(err));
}
}

100
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h Normal file
View File

@ -0,0 +1,100 @@
/*!
******************* BEGIN Caffe Copyright Notice and Disclaimer ****************
*
* COPYRIGHT
*
* All contributions by the University of California:
* Copyright (c) 2014-2017 The Regents of the University of California (Regents)
* All rights reserved.
*
* All other contributions:
* Copyright (c) 2014-2017, the respective contributors
* All rights reserved.
*
* Caffe uses a shared copyright model: each contributor holds copyright over
* their contributions to Caffe. The project versioning records all such
* contribution and copyright details. If a contributor wants to further mark
* their specific copyright on a particular contribution, they should indicate
* their copyright solely in the commit message of the change when it is
* committed.
*
* LICENSE
*
* 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 COPYRIGHT HOLDERS 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 COPYRIGHT OWNER 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.
*
* CONTRIBUTION AGREEMENT
*
* By contributing to the BVLC/caffe repository through pull-request, comment,
* or otherwise, the contributor releases their content to the
* license and copyright terms herein.
*
***************** END Caffe Copyright Notice and Disclaimer ********************
*
* Copyright (c) 2018 Microsoft
* Licensed under The MIT License [see LICENSE for details]
* \file modulated_deformable_im2col.h
* \brief Function definitions of converting an image to
* column matrix based on kernel, padding, dilation, and offset.
* These functions are mainly used in deformable convolution operators.
* \ref: https://arxiv.org/abs/1811.11168
* \author Yuwen Xiong, Haozhi Qi, Jifeng Dai, Xizhou Zhu, Han Hu
*/
/***************** Adapted by Charles Shang *********************/
#ifndef DCN_V2_IM2COL_CUDA
#define DCN_V2_IM2COL_CUDA
// #ifdef __cplusplus
// extern "C"
// {
// #endif
void modulated_deformable_im2col_cuda(cudaStream_t stream,
const float *data_im, const float *data_offset, const float *data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group, float *data_col);
void modulated_deformable_col2im_cuda(cudaStream_t stream,
const float *data_col, const float *data_offset, const float *data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group, float *grad_im);
void modulated_deformable_col2im_coord_cuda(cudaStream_t stream,
const float *data_col, const float *data_im, const float *data_offset, const float *data_mask,
const int batch_size, const int channels, const int height_im, const int width_im,
const int height_col, const int width_col, const int kernel_h, const int kenerl_w,
const int pad_h, const int pad_w, const int stride_h, const int stride_w,
const int dilation_h, const int dilation_w,
const int deformable_group,
float *grad_offset, float *grad_mask);
// #ifdef __cplusplus
// }
// #endif
#endif

402
centernet/dcnv2Plugin/dcnv2Plugin.cpp Normal file
View File

@ -0,0 +1,402 @@
#include "dcnv2Plugin.h"
#include <iostream>
using namespace nvinfer1;
using nvinfer1::plugin::DeformableConvolutionalLayer;
using nvinfer1::plugin::DCNv2PluginCreator;
namespace
{
const char* DCNv2_PLUGIN_VERSION{"1"};
const char* DCNv2_PLUGIN_NAME{"DCNv2_TRT"};
} // namespace
#define CHECK_CUDA(call) \
do \
{ \
cudaError_t status = call; \
if (status != cudaSuccess) \
{ \
return status; \
} \
} while (0)
PluginFieldCollection DCNv2PluginCreator::mFC{};
std::vector<PluginField> DCNv2PluginCreator::mPluginAttributes;
// Parameterized constructor
DeformableConvolutionalLayer::DeformableConvolutionalLayer(
int out_channels,
int kernel,
int deformable_group,
int dilation,
int padding,
int stride,
const Weights* weight, const Weights* bias):
out_channels(out_channels),kernel_size(kernel),deformable_group(deformable_group),
dilation(dilation),padding(padding),stride(stride){
mWeight = copyToDevice(weight[0].values, weight[0].count);
mBias = copyToDevice(bias[0].values, bias[0].count);
}
DeformableConvolutionalLayer::DeformableConvolutionalLayer(const void* buffer, size_t length)
{
const char* d = static_cast<const char*>(buffer);
const char* a = d;
in_channels = read<int>(d);
height = read<int>(d);
width = read<int>(d);
height_out = read<int>(d);
width_out = read<int>(d);
out_channels = read<int>(d);
kernel_size = read<int>(d);
deformable_group = read<int>(d);
dilation = read<int>(d);
padding = read<int>(d);
stride = read<int>(d);
int count = read<int>(d);
mWeight = deserializeToDevice(d, count);
count = read<int>(d);
mBias = deserializeToDevice(d, count);
ASSERT(d == a + length);
}
int DeformableConvolutionalLayer::getNbOutputs() const
{
// Plugin layer has 2 outputs
return 1;
}
int DeformableConvolutionalLayer::initialize()
{
size_t oneSize = height_out * width_out * sizeof(float);
std::vector<float> one_((int)oneSize, 1.0f);
CHECK_CUDA(cudaMalloc((void**)&mOne, oneSize));
CHECK_CUDA(cudaMalloc((void**)&mColumn, in_channels * kernel_size * kernel_size * oneSize));
CHECK_CUDA(cudaMemcpy(mOne, one_.data(), oneSize, cudaMemcpyHostToDevice));
return STATUS_SUCCESS;
}
Dims DeformableConvolutionalLayer::getOutputDimensions(int index, const Dims* inputs, int nbInputs)
{
ASSERT(index == 0);
ASSERT(nbInputs == 3);
in_channels = inputs[0].d[0];
height = inputs[0].d[1];
width = inputs[0].d[2];
height_out = (inputs[0].d[1] + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1;
width_out = (inputs[0].d[2] + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1;
return Dims3(out_channels, height_out, width_out);
}
size_t DeformableConvolutionalLayer::getWorkspaceSize(int maxBatchSize) const
{
return 0;
}
int DeformableConvolutionalLayer::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)
{
const float* input = static_cast<const float *>(inputs[0]);
const float* offset = static_cast<const float *>(inputs[1]);
const float* offset_mask = static_cast<const float *>(inputs[2]);
const float* mask = offset_mask + deformable_group * 2 * kernel_size * kernel_size * height * width;
float * output = static_cast<float *>(outputs[0]);
float alpha{1}, beta{0};
// Do Bias first:
// M,N,K are dims of matrix A and B
// (see http://docs.nvidia.com/cuda/cublas/#cublas-lt-t-gt-gemm)
// (N x 1) (1 x M)
int m_ = out_channels;
int n_ = height_out * width_out;
int k_ = 1;
cublasSgemm(mCublas, CUBLAS_OP_T, CUBLAS_OP_N, n_, m_, k_, &alpha,
mOne, k_,
static_cast<const float *>(mBias.values), k_, &beta,
output, n_);
modulated_deformable_im2col_cuda(stream, input, offset, mask,
1, in_channels, height, width,
height_out, width_out, kernel_size, kernel_size,
padding, padding, stride, stride, dilation, dilation,
deformable_group, mColumn);
//(k * m) x (m * n)
// Y = WC
int m = out_channels;
int n = height_out * width_out;
int k = in_channels * kernel_size * kernel_size;
cublasSgemm(mCublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha,
mColumn, n,
static_cast<const float *>(mWeight.values), k, &alpha,
output, n);
return 0;
}
size_t DeformableConvolutionalLayer::getSerializationSize() const
{
return sizeof(int) * 13 + (mWeight.count + mBias.count) * sizeof(float);
}
void DeformableConvolutionalLayer::serialize(void* buffer) const
{
char *d = reinterpret_cast<char*>(buffer), *a = d;
write(d, in_channels);
write(d, height);
write(d, width);
write(d, height_out);
write(d, width_out);
write(d, out_channels);
write(d, kernel_size);
write(d, deformable_group);
write(d, dilation);
write(d, padding);
write(d, stride);
write(d, (int) mWeight.count);
serializeFromDevice(d, mWeight);
write(d, (int) mBias.count);
serializeFromDevice(d, mBias);
ASSERT(d == a + getSerializationSize());
}
bool DeformableConvolutionalLayer::supportsFormat(DataType type, PluginFormat format) const
{
return (type == DataType::kFLOAT && format == PluginFormat::kNCHW);
}
Weights DeformableConvolutionalLayer::copyToDevice(const void* hostData, size_t count)
{
void* deviceData;
CUASSERT(cudaMalloc(&deviceData, count * sizeof(float)));
CUASSERT(cudaMemcpy(deviceData, hostData, count * sizeof(float), cudaMemcpyHostToDevice));
return Weights{DataType::kFLOAT, deviceData, int64_t(count)};
}
void DeformableConvolutionalLayer::serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const
{
CUASSERT(cudaMemcpy(hostBuffer, deviceWeights.values, deviceWeights.count * sizeof(float), cudaMemcpyDeviceToHost));
hostBuffer += deviceWeights.count * sizeof(float);
}
Weights DeformableConvolutionalLayer::deserializeToDevice(const char*& hostBuffer, size_t count)
{
Weights w = copyToDevice(hostBuffer, count);
hostBuffer += count * sizeof(float);
return w;
}
const char* DeformableConvolutionalLayer::getPluginType() const
{
return DCNv2_PLUGIN_NAME;
}
const char* DeformableConvolutionalLayer::getPluginVersion() const
{
return DCNv2_PLUGIN_VERSION;
}
void DeformableConvolutionalLayer::terminate() {
if (mOne)
{
cudaFree(mOne);
mOne = nullptr;
}
if (mColumn)
{
cudaFree(mColumn);
mColumn = nullptr;
}
}
void DeformableConvolutionalLayer::destroy()
{
delete this;
}
IPluginV2Ext* DeformableConvolutionalLayer::clone() const
{
IPluginV2Ext* plugin = new DeformableConvolutionalLayer(*this);
plugin->setPluginNamespace(mPluginNamespace.c_str());
return plugin;
}
// Set plugin namespace
void DeformableConvolutionalLayer::setPluginNamespace(const char* pluginNamespace)
{
mPluginNamespace = pluginNamespace;
}
const char* DeformableConvolutionalLayer::getPluginNamespace() const
{
return mPluginNamespace.c_str();
}
// Return the DataType of the plugin output at the requested index.
DataType DeformableConvolutionalLayer::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
{
// Only DataType::kFLOAT is acceptable by the plugin layer
return DataType::kFLOAT;
}
// Return true if output tensor is broadcast across a batch.
bool DeformableConvolutionalLayer::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
{
return false;
}
// Return true if plugin can use input that is broadcast across batch without replication.
bool DeformableConvolutionalLayer::canBroadcastInputAcrossBatch(int inputIndex) const
{
return false;
}
// Configure the layer with input and output data types.
// inutDims: input Dimensions for the plugin layer
// nInputs : Number of inputs to the plugin layer
// outputDims: output Dimensions from the plugin layer
// nOutputs: number of outputs from the plugin layer
// type: DataType configuration for the plugin layer
// format: format NCHW, NHWC etc
// maxbatchSize: maximum batch size for the plugin layer
void DeformableConvolutionalLayer::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)
{
ASSERT(*inputTypes == DataType::kFLOAT && floatFormat == PluginFormat::kNCHW);
}
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void DeformableConvolutionalLayer::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
{
mCublas = cublasContext;
}
// Detach the plugin object from its execution context.
void DeformableConvolutionalLayer::detachFromContext() {}
DCNv2PluginCreator::DCNv2PluginCreator()
{
mPluginAttributes.emplace_back(PluginField("out_channels", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("kernel", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("deformable_group", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("dilation", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("padding", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("stride", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("weight", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
const char* DCNv2PluginCreator::getPluginName() const
{
return DCNv2_PLUGIN_NAME;
}
const char* DCNv2PluginCreator::getPluginVersion() const
{
return DCNv2_PLUGIN_VERSION;
}
const PluginFieldCollection* DCNv2PluginCreator::getFieldNames()
{
return &mFC;
}
IPluginV2Ext* DCNv2PluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
std::vector<float> weight;
std::vector<float> bias;
int out_channels, kernel, deformable_group, padding, stride, dilation;
const PluginField* fields = fc->fields;
for (int i = 0; i < fc->nbFields; ++i)
{
const char* attrName = fields[i].name;
if (!strcmp(attrName, "out_channels"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
out_channels = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "kernel"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
kernel = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "deformable_group"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
deformable_group = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "dilation"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
dilation = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "stride"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
stride = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "padding"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
padding = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "weight"))
{
ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
int size = fields[i].length;
weight.reserve(size);
const auto* w = static_cast<const float*>(fields[i].data);
for (int j = 0; j < size; j++)
{
weight.push_back(*w);
w++;
}
}
else if (!strcmp(attrName, "bias"))
{
ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
int size = fields[i].length;
bias.reserve(size);
const auto* w = static_cast<const float*>(fields[i].data);
for (int j = 0; j < size; j++)
{
bias.push_back(*w);
w++;
}
}
}
Weights mWeight{DataType::kFLOAT, weight.data(), (int64_t) weight.size()};
Weights mBias{DataType::kFLOAT, bias.data(), (int64_t) bias.size()};
DeformableConvolutionalLayer* obj = new DeformableConvolutionalLayer(out_channels,
kernel,
deformable_group,
dilation,
padding,
stride,
&mWeight, &mBias);
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}
IPluginV2Ext* DCNv2PluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
{
// This object will be deleted when the network is destroyed, which will
// call Normalize::destroy()
DeformableConvolutionalLayer* obj = new DeformableConvolutionalLayer(serialData, serialLength);
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}

152
centernet/dcnv2Plugin/dcnv2Plugin.h Normal file
View File

@ -0,0 +1,152 @@
/*
* Copyright (c) 2020, 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_DCNV2_PLUGIN_H
#define TRT_DCNV2_PLUGIN_H
#include "kernel.h"
#include "plugin.h"
#include "dcn_v2_im2col_cuda.h"
#include "serialize.hpp"
#include <cudnn.h>
#include <vector>
#include <cublas_v2.h>
#include <cuda.h>
#include <string>
#include <vector>
using namespace nvinfer1::plugin;
namespace nvinfer1
{
namespace plugin
{
class DeformableConvolutionalLayer : public IPluginV2Ext
{
public:
DeformableConvolutionalLayer(int out_channels,
int kernel,
int deformable_group,
int dilation,
int padding,
int stride,
const Weights* weight, const Weights* bias);
DeformableConvolutionalLayer(const void* buffer, size_t length);
~DeformableConvolutionalLayer() override = default;
int getNbOutputs() const override;
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputs) 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;
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;
void setPluginNamespace(const char* pluginNamespace) override;
const char* getPluginNamespace() const override;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
bool canBroadcastInputAcrossBatch(int inputIndex) const override;
void attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) 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;
void detachFromContext() override;
private:
Weights copyToDevice(const void* hostData, size_t count);
void serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const;
Weights deserializeToDevice(const char*& hostBuffer, size_t count);
std::string mPluginNamespace;
int in_channels{};
int height_out{};
int width_out{};
int height{};
int width{};
int out_channels{};
int kernel_size{};
int deformable_group{};
int dilation{};
int padding{};
int stride{};
Weights mWeight{};
Weights mBias{};
float* mOne;
float* mColumn;
cublasHandle_t mCublas;
};
class DCNv2PluginCreator : public BaseCreator
{
public:
DCNv2PluginCreator();
~DCNv2PluginCreator() 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;
// Parameters for DeformableConvolutionalLayer
static std::vector<PluginField> mPluginAttributes;
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_DCNv2_PLUGIN_H

238
centernet/sample/common.py Normal file
View File

@ -0,0 +1,238 @@
#
# Copyright 1993-2020 NVIDIA Corporation. All rights reserved.
#
# NOTICE TO LICENSEE:
#
# This source code and/or documentation ("Licensed Deliverables") are
# subject to NVIDIA intellectual property rights under U.S. and
# international Copyright laws.
#
# These Licensed Deliverables contained herein is PROPRIETARY and
# CONFIDENTIAL to NVIDIA and is being provided under the terms and
# conditions of a form of NVIDIA software license agreement by and
# between NVIDIA and Licensee ("License Agreement") or electronically
# accepted by Licensee. Notwithstanding any terms or conditions to
# the contrary in the License Agreement, reproduction or disclosure
# of the Licensed Deliverables to any third party without the express
# written consent of NVIDIA is prohibited.
#
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
# SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS
# PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
# NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
# DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
# NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
# SPECIAL, INDIRECT, INCIDENTAL, 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 THESE LICENSED DELIVERABLES.
#
# U.S. Government End Users. These Licensed Deliverables are a
# "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
# 1995), consisting of "commercial computer software" and "commercial
# computer software documentation" as such terms are used in 48
# C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
# only as a commercial end item. Consistent with 48 C.F.R.12.212 and
# 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
# U.S. Government End Users acquire the Licensed Deliverables with
# only those rights set forth herein.
#
# Any use of the Licensed Deliverables in individual and commercial
# software must include, in the user documentation and internal
# comments to the code, the above Disclaimer and U.S. Government End
# Users Notice.
#
from itertools import chain
import argparse
import os
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import tensorrt as trt
try:
# Sometimes python2 does not understand FileNotFoundError
FileNotFoundError
except NameError:
FileNotFoundError = IOError
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
def GiB(val):
return val * 1 << 30
def add_help(description):
parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
args, _ = parser.parse_known_args()
def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""):
'''
Parses sample arguments.
Args:
description (str): Description of the sample.
subfolder (str): The subfolder containing data relevant to this sample
find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.
Returns:
str: Path of data directory.
'''
# Standard command-line arguments for all samples.
kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("-d", "--datadir", help="Location of the TensorRT sample data directory, and any additional data directories.", action="append", default=[kDEFAULT_DATA_ROOT])
args, _ = parser.parse_known_args()
def get_data_path(data_dir):
# If the subfolder exists, append it to the path, otherwise use the provided path as-is.
data_path = os.path.join(data_dir, subfolder)
if not os.path.exists(data_path):
if data_dir != kDEFAULT_DATA_ROOT:
print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")
data_path = data_dir
# Make sure data directory exists.
if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:
print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(data_path))
return data_path
data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
return data_paths, locate_files(data_paths, find_files, err_msg)
def locate_files(data_paths, filenames, err_msg=""):
"""
Locates the specified files in the specified data directories.
If a file exists in multiple data directories, the first directory is used.
Args:
data_paths (List[str]): The data directories.
filename (List[str]): The names of the files to find.
Returns:
List[str]: The absolute paths of the files.
Raises:
FileNotFoundError if a file could not be located.
"""
found_files = [None] * len(filenames)
for data_path in data_paths:
# Find all requested files.
for index, (found, filename) in enumerate(zip(found_files, filenames)):
if not found:
file_path = os.path.abspath(os.path.join(data_path, filename))
if os.path.exists(file_path):
found_files[index] = file_path
# Check that all files were found
for f, filename in zip(found_files, filenames):
if not f or not os.path.exists(f):
raise FileNotFoundError("Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg))
return found_files
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
# This function is generalized for multiple inputs/outputs for full dimension networks.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference_v2(context, bindings, inputs, outputs, stream):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
# `retry_call` and `retry` are used to wrap the function we want to try multiple times
def retry_call(func, args=[], kwargs={}, n_retries=3):
"""Wrap a function to retry it several times.
Args:
func: function to call
args (List): args parsed to func
kwargs (Dict): kwargs parsed to func
n_retries (int): maximum times of tries
"""
for i_try in range(n_retries):
try:
func(*args, **kwargs)
break
except:
if i_try == n_retries - 1:
raise
print("retry...")
# Usage: @retry(n_retries)
def retry(n_retries=3):
"""Wrap a function to retry it several times. Decorator version of `retry_call`.
Args:
n_retries (int): maximum times of tries
Usage:
@retry(n_retries)
def func(...):
pass
"""
def wrapper(func):
def _wrapper(*args, **kwargs):
retry_call(func, args, kwargs, n_retries)
return _wrapper
return wrapper

134
centernet/sample/test.py Normal file
View File

@ -0,0 +1,134 @@
import cv2 as cv
import numpy as np
import tensorrt as trt
import common
import torch
import time
from sys import argv
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
def _gather_feat(feat, ind, mask=None):
dim = feat.size(2)
ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
feat = feat.gather(1, ind)
if mask is not None:
mask = mask.unsqueeze(2).expand_as(feat)
feat = feat[mask]
feat = feat.view(-1, dim)
return feat
def _transpose_and_gather_feat(feat, ind):
feat = feat.permute(0, 2, 3, 1).contiguous()
feat = feat.view(feat.size(0), -1, feat.size(3))
feat = _gather_feat(feat, ind)
return feat
def pre_process(image):
long_size = max(image.shape)
img = np.zeros((long_size, long_size, 3))
img[:image.shape[0], :img.shape[1], :] = image[:]
img = cv.resize(img, (512,512))
inp_image = ((img / 255. - 0.5) / 0.5).astype(np.float32)
images = inp_image.transpose(2, 0, 1)
return images, long_size/512
def _nms(heat, kernel=3):
pad = (kernel - 1) // 2
hmax = torch.nn.functional.max_pool2d(
heat, (kernel, kernel), stride=1, padding=pad)
keep = (hmax == heat).float()
return heat * keep
def _topk(scores, K=40):
batch, cat, height, width = scores.size()
topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)
topk_inds = topk_inds % (height * width)
topk_ys = (topk_inds.true_divide(width)).int().float()
topk_xs = (topk_inds % width).int().float()
topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
topk_clses = (topk_ind.true_divide(K)).int()
topk_inds = _gather_feat(
topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)
return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
def ctdet_decode(heat, wh, reg=None, cat_spec_wh=False, K=100):
batch, cat, height, width = heat.size()
heat = torch.sigmoid(heat)
# perform nms on heatmaps
heat = _nms(heat)
scores, inds, clses, ys, xs = _topk(heat, K=K)
if reg is not None:
reg = _transpose_and_gather_feat(reg, inds)
reg = reg.view(batch, K, 2)
xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
else:
xs = xs.view(batch, K, 1) + 0.5
ys = ys.view(batch, K, 1) + 0.5
wh = _transpose_and_gather_feat(wh, inds)
if cat_spec_wh:
wh = wh.view(batch, K, cat, 2)
clses_ind = clses.view(batch, K, 1, 1).expand(batch, K, 1, 2).long()
wh = wh.gather(2, clses_ind).view(batch, K, 2)
else:
wh = wh.view(batch, K, 2)
clses = clses.view(batch, K, 1).float()
scores = scores.view(batch, K, 1)
bboxes = torch.cat([xs - wh[..., 0:1] / 2,
ys - wh[..., 1:2] / 2,
xs + wh[..., 0:1] / 2,
ys + wh[..., 1:2] / 2], dim=2)
detections = torch.cat([bboxes, scores, clses], dim=2)
return detections
if __name__ == '__main__':
try:
engine_path = argv[1]
img_path = argv[2]
except:
print('engine path and image path are needed!')
exit()
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine:
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
with engine.create_execution_context() as context:
img = cv.imread('test.jpg')
dis = img.copy()
img, s = pre_process(img)
# Copy to the pagelocked input buffer
np.copyto(inputs[0].host, img.ravel())
[hm, wh, reg] = common.do_inference(
context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=1)
[dets] = ctdet_decode(torch.from_numpy(hm.reshape(1, 80, 128, 128)), torch.from_numpy(
wh.reshape(1, 2, 128, 128)), torch.from_numpy(reg.reshape(1, 2, 128, 128)))
for i in dets:
if i[-2] > 0.5:
i[:4] *= 4*s
cv.rectangle(dis, (int(i[0]), int(
i[1])), (int(i[2]), int(i[3])), 255, 1)
cv.putText(dis, '%d' %
int(i[-1]), (int(i[0]), int(i[1])), 1, 1, 255)
cv.imwrite('trt_out.jpg', dis)

33
crnn/CMakeLists.txt Normal file
View File

@ -0,0 +1,33 @@
cmake_minimum_required(VERSION 2.6)
project(crnn)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
message("embed_platform on")
include_directories(/usr/local/cuda/targets/aarch64-linux/include)
link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
else()
message("embed_platform off")
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
endif()
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(crnn ${PROJECT_SOURCE_DIR}/crnn.cpp)
target_link_libraries(crnn nvinfer)
target_link_libraries(crnn cudart)
target_link_libraries(crnn ${OpenCV_LIBS})
add_definitions(-O2 -pthread)

44
crnn/README.md Normal file
View File

@ -0,0 +1,44 @@
# crnn
The Pytorch implementation is [meijieru/crnn.pytorch](https://github.com/meijieru/crnn.pytorch).
## How to Run
```
1. generate crnn.wts from pytorch
git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/meijieru/crnn.pytorch.git
// download its weights 'crnn.pth'
// copy tensorrtx/crnn/genwts.py into crnn.pytorch/
// go to crnn.pytorch/
python genwts.py
// a file 'crnn.wts' will be generated.
2. build tensorrtx/crnn and run
// put crnn.wts into tensorrtx/crnn
// go to tensorrtx/crnn
mkdir build
cd build
cmake ..
make
sudo ./crnn -s // serialize model to plan file i.e. 'crnn.engine'
// copy crnn.pytorch/data/demo.png here
sudo ./crnn -d // deserialize plan file and run inference
3. check the output as follows:
raw: a-----v--a-i-l-a-bb-l-e---
sim: available
```
## More Information
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
## Acknowledgment
Thanks for the donation for this crnn tensorrt implementation from @雍.

418
crnn/crnn.cpp Normal file
View File

@ -0,0 +1,418 @@
#include <iostream>
#include <chrono>
#include <map>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define BATCH_SIZE 1
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 32;
static const int INPUT_W = 100;
static const int OUTPUT_SIZE = 26 * 37;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
const int ks[] = {3, 3, 3, 3, 3, 3, 2};
const int ps[] = {1, 1, 1, 1, 1, 1, 0};
const int ss[] = {1, 1, 1, 1, 1, 1, 1};
const int nm[] = {64, 128, 256, 256, 512, 512, 512};
const std::string alphabet = "-0123456789abcdefghijklmnopqrstuvwxyz";
using namespace nvinfer1;
std::string strDecode(std::vector<int>& preds, bool raw) {
std::string str;
if (raw) {
for (auto v: preds) {
str.push_back(alphabet[v]);
}
} else {
for (size_t i = 0; i < preds.size(); i++) {
if (preds[i] == 0 || (i > 0 && preds[i - 1] == preds[i])) continue;
str.push_back(alphabet[preds[i]]);
}
}
return str;
}
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + ".weight"].values;
float *beta = (float*)weightMap[lname + ".bias"].values;
float *mean = (float*)weightMap[lname + ".running_mean"].values;
float *var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* convRelu(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int i, bool use_bn = false) {
int nOut = nm[i];
IConvolutionLayer* conv = network->addConvolutionNd(input, nOut, DimsHW{ks[i], ks[i]}, weightMap["cnn.conv" + std::to_string(i) + ".weight"], weightMap["cnn.conv" + std::to_string(i) + ".bias"]);
assert(conv);
conv->setStrideNd(DimsHW{ss[i], ss[i]});
conv->setPaddingNd(DimsHW{ps[i], ps[i]});
ILayer *tmp = conv;
if (use_bn) {
tmp = addBatchNorm2d(network, weightMap, *conv->getOutput(0), "cnn.batchnorm" + std::to_string(i), 1e-5);
}
auto relu = network->addActivation(*tmp->getOutput(0), ActivationType::kRELU);
assert(relu);
return relu;
}
void splitLstmWeights(std::map<std::string, Weights>& weightMap, std::string lname) {
int weight_size = weightMap[lname].count;
for (int i = 0; i < 4; i++) {
Weights wt{DataType::kFLOAT, nullptr, 0};
wt.count = weight_size / 4;
float *val = reinterpret_cast<float*>(malloc(sizeof(float) * wt.count));
memcpy(val, (float*)weightMap[lname].values + wt.count * i, sizeof(float) * wt.count);
wt.values = val;
weightMap[lname + std::to_string(i)] = wt;
}
}
ILayer* addLSTM(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int nHidden, std::string lname) {
splitLstmWeights(weightMap, lname + ".weight_ih_l0");
splitLstmWeights(weightMap, lname + ".weight_hh_l0");
splitLstmWeights(weightMap, lname + ".bias_ih_l0");
splitLstmWeights(weightMap, lname + ".bias_hh_l0");
splitLstmWeights(weightMap, lname + ".weight_ih_l0_reverse");
splitLstmWeights(weightMap, lname + ".weight_hh_l0_reverse");
splitLstmWeights(weightMap, lname + ".bias_ih_l0_reverse");
splitLstmWeights(weightMap, lname + ".bias_hh_l0_reverse");
Dims dims = input.getDimensions();
std::cout << "lstm input shape: " << dims.nbDims << " [" << dims.d[0] << " " << dims.d[1] << " " << dims.d[2] << "]"<< std::endl;
auto lstm = network->addRNNv2(input, 1, nHidden, dims.d[1], RNNOperation::kLSTM);
lstm->setDirection(RNNDirection::kBIDIRECTION);
lstm->setWeightsForGate(0, RNNGateType::kINPUT, true, weightMap[lname + ".weight_ih_l00"]);
lstm->setWeightsForGate(0, RNNGateType::kFORGET, true, weightMap[lname + ".weight_ih_l01"]);
lstm->setWeightsForGate(0, RNNGateType::kCELL, true, weightMap[lname + ".weight_ih_l02"]);
lstm->setWeightsForGate(0, RNNGateType::kOUTPUT, true, weightMap[lname + ".weight_ih_l03"]);
lstm->setWeightsForGate(0, RNNGateType::kINPUT, false, weightMap[lname + ".weight_hh_l00"]);
lstm->setWeightsForGate(0, RNNGateType::kFORGET, false, weightMap[lname + ".weight_hh_l01"]);
lstm->setWeightsForGate(0, RNNGateType::kCELL, false, weightMap[lname + ".weight_hh_l02"]);
lstm->setWeightsForGate(0, RNNGateType::kOUTPUT, false, weightMap[lname + ".weight_hh_l03"]);
lstm->setBiasForGate(0, RNNGateType::kINPUT, true, weightMap[lname + ".bias_ih_l00"]);
lstm->setBiasForGate(0, RNNGateType::kFORGET, true, weightMap[lname + ".bias_ih_l01"]);
lstm->setBiasForGate(0, RNNGateType::kCELL, true, weightMap[lname + ".bias_ih_l02"]);
lstm->setBiasForGate(0, RNNGateType::kOUTPUT, true, weightMap[lname + ".bias_ih_l03"]);
lstm->setBiasForGate(0, RNNGateType::kINPUT, false, weightMap[lname + ".bias_hh_l00"]);
lstm->setBiasForGate(0, RNNGateType::kFORGET, false, weightMap[lname + ".bias_hh_l01"]);
lstm->setBiasForGate(0, RNNGateType::kCELL, false, weightMap[lname + ".bias_hh_l02"]);
lstm->setBiasForGate(0, RNNGateType::kOUTPUT, false, weightMap[lname + ".bias_hh_l03"]);
lstm->setWeightsForGate(1, RNNGateType::kINPUT, true, weightMap[lname + ".weight_ih_l0_reverse0"]);
lstm->setWeightsForGate(1, RNNGateType::kFORGET, true, weightMap[lname + ".weight_ih_l0_reverse1"]);
lstm->setWeightsForGate(1, RNNGateType::kCELL, true, weightMap[lname + ".weight_ih_l0_reverse2"]);
lstm->setWeightsForGate(1, RNNGateType::kOUTPUT, true, weightMap[lname + ".weight_ih_l0_reverse3"]);
lstm->setWeightsForGate(1, RNNGateType::kINPUT, false, weightMap[lname + ".weight_hh_l0_reverse0"]);
lstm->setWeightsForGate(1, RNNGateType::kFORGET, false, weightMap[lname + ".weight_hh_l0_reverse1"]);
lstm->setWeightsForGate(1, RNNGateType::kCELL, false, weightMap[lname + ".weight_hh_l0_reverse2"]);
lstm->setWeightsForGate(1, RNNGateType::kOUTPUT, false, weightMap[lname + ".weight_hh_l0_reverse3"]);
lstm->setBiasForGate(1, RNNGateType::kINPUT, true, weightMap[lname + ".bias_ih_l0_reverse0"]);
lstm->setBiasForGate(1, RNNGateType::kFORGET, true, weightMap[lname + ".bias_ih_l0_reverse1"]);
lstm->setBiasForGate(1, RNNGateType::kCELL, true, weightMap[lname + ".bias_ih_l0_reverse2"]);
lstm->setBiasForGate(1, RNNGateType::kOUTPUT, true, weightMap[lname + ".bias_ih_l0_reverse3"]);
lstm->setBiasForGate(1, RNNGateType::kINPUT, false, weightMap[lname + ".bias_hh_l0_reverse0"]);
lstm->setBiasForGate(1, RNNGateType::kFORGET, false, weightMap[lname + ".bias_hh_l0_reverse1"]);
lstm->setBiasForGate(1, RNNGateType::kCELL, false, weightMap[lname + ".bias_hh_l0_reverse2"]);
lstm->setBiasForGate(1, RNNGateType::kOUTPUT, false, weightMap[lname + ".bias_hh_l0_reverse3"]);
return lstm;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {C, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../crnn.wts");
// cnn
auto x = convRelu(network, weightMap, *data, 0);
auto p = network->addPoolingNd(*x->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
p->setStrideNd(DimsHW{2, 2});
x = convRelu(network, weightMap, *p->getOutput(0), 1);
p = network->addPoolingNd(*x->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
p->setStrideNd(DimsHW{2, 2});
x = convRelu(network, weightMap, *p->getOutput(0), 2, true);
x = convRelu(network, weightMap, *x->getOutput(0), 3);
p = network->addPoolingNd(*x->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
p->setStrideNd(DimsHW{2, 1});
p->setPaddingNd(DimsHW{0, 1});
x = convRelu(network, weightMap, *p->getOutput(0), 4, true);
x = convRelu(network, weightMap, *x->getOutput(0), 5);
p = network->addPoolingNd(*x->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
p->setStrideNd(DimsHW{2, 1});
p->setPaddingNd(DimsHW{0, 1});
x = convRelu(network, weightMap, *p->getOutput(0), 6, true);
auto sfl = network->addShuffle(*x->getOutput(0));
sfl->setFirstTranspose(Permutation{1, 2, 0});
// rnn
auto lstm0 = addLSTM(network, weightMap, *sfl->getOutput(0), 256, "rnn.0.rnn");
auto sfl0 = network->addShuffle(*lstm0->getOutput(0));
sfl0->setReshapeDimensions(Dims4{26, 1, 1, 512});
auto fc0 = network->addFullyConnected(*sfl0->getOutput(0), 256, weightMap["rnn.0.embedding.weight"], weightMap["rnn.0.embedding.bias"]);
sfl = network->addShuffle(*fc0->getOutput(0));
sfl->setFirstTranspose(Permutation{2, 3, 0, 1});
sfl->setReshapeDimensions(Dims3{1, 26, 256});
auto lstm1 = addLSTM(network, weightMap, *sfl->getOutput(0), 256, "rnn.1.rnn");
auto sfl1 = network->addShuffle(*lstm1->getOutput(0));
sfl1->setReshapeDimensions(Dims4{26, 1, 1, 512});
auto fc1 = network->addFullyConnected(*sfl1->getOutput(0), 37, weightMap["rnn.1.embedding.weight"], weightMap["rnn.1.embedding.bias"]);
Dims dims = fc1->getOutput(0)->getDimensions();
std::cout << "fc1 shape " << dims.d[0] << " " << dims.d[1] << " " << dims.d[2] << std::endl;
fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*fc1->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) {
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 1 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(BATCH_SIZE, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("crnn.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (argc == 2 && std::string(argv[1]) == "-d") {
std::ifstream file("crnn.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./crnn -s // serialize model to plan file" << std::endl;
std::cerr << "./crnn -d ../samples // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
static float data[BATCH_SIZE * 1 * INPUT_H * INPUT_W];
//for (int i = 0; i < 1 * INPUT_H * INPUT_W; i++)
// data[i] = 1.0;
static float prob[BATCH_SIZE * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
assert(engine->getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);
assert(inputIndex == 0);
assert(outputIndex == 1);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * 1 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
cv::Mat img = cv::imread("demo.png");
if (img.empty()) {
std::cerr << "demo.png not found !!!" << std::endl;
return -1;
}
cv::cvtColor(img, img, CV_BGR2GRAY);
cv::resize(img, img, cv::Size(INPUT_W, INPUT_H));
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
data[i] = ((float)img.at<uchar>(i) / 255.0 - 0.5) * 2.0;
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, stream, buffers, data, prob, BATCH_SIZE);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
std::vector<int> preds;
for (int i = 0; i < 26; i++) {
int maxj = 0;
for (int j = 1; j < 37; j++) {
if (prob[37 * i + j] > prob[37 * i + maxj]) maxj = j;
}
preds.push_back(maxj);
}
std::cout << "raw: " << strDecode(preds, true) << std::endl;
std::cout << "sim: " << strDecode(preds, false) << std::endl;
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
// Print histogram of the output distribution
//std::cout << "\nOutput:\n\n";
//for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
//{
// std::cout << prob[i] << ", ";
// if (i % 10 == 0) std::cout << std::endl;
//}
//std::cout << std::endl;
return 0;
}

35
crnn/genwts.py Normal file
View File

@ -0,0 +1,35 @@
import torch
from torch.autograd import Variable
import utils
import models.crnn as crnn
import struct
model_path = './data/crnn.pth'
model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict(torch.load(model_path))
image = torch.ones(1, 1, 32, 100)
if torch.cuda.is_available():
image = image.cuda()
model.eval()
print(model)
print('image shape ', image.shape)
preds = model(image)
f = open("crnn.wts", 'w')
f.write("{}\n".format(len(model.state_dict().keys())))
for k,v in model.state_dict().items():
print('key: ', k)
print('value: ', v.shape)
vr = v.reshape(-1).cpu().numpy()
f.write("{} {}".format(k, len(vr)))
for vv in vr:
f.write(" ")
f.write(struct.pack(">f", float(vv)).hex())
f.write("\n")

503
crnn/logging.h Normal file
View File

@ -0,0 +1,503 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

26
csrnet/CMakeLists.txt Normal file
View File

@ -0,0 +1,26 @@
cmake_minimum_required(VERSION 3.10)
project(csrnet)
add_definitions(-std=c++11)
add_definitions(-DAPI_EXPORTS)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
# cuda
include_directories(/usr/local/cuda/targets/x86_64-linux/include )
link_directories(/usr/local/cuda/targets/x86_64-linux/lib)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
# opencv
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
include_directories(${PROJECT_SOURCE_DIR}/)
add_executable(csrnet csrnet.cpp)
target_link_libraries(csrnet nvinfer cudart ${OpenCV_LIBS})

58
csrnet/README.md Normal file
View File

@ -0,0 +1,58 @@
# csrnet
The Pytorch implementation is [leeyeehoo/CSRNet-pytorch](https://github.com/leeyeehoo/CSRNet-pytorch).
This repo is a TensorRT implementation of CSRNet.
paper : [CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes](https://arxiv.org/abs/1802.10062)
Dev environment:
- Ubuntu 22.04
- TensorRT 8.6
- OpenCV 4.5.4
- CMake 3.24
- GPU Driver 535.113.01
- CUDA 12.2
- RTX3080
# how to run
```bash
1. generate csrnet engine
git clone https://github.com/leeyeehoo/CSRNet-pytorch.git
git clone https://github.com/wang-xinyu/tensorrtx.git
// copy gen_wts.py to CSRNet-pytorch
// generate wts file
python gen_wts.py
// csrnet wts will be generated in CSRNet-pytorch
2. build csrnet.engine
// mv CSRNet-pytorch/csrnet.engine to tensorrtx/csrnet
mv CSRNet-pytorch/csrnet.wts tensorrtx/csrnet
// build
mkdir build
cmake ..
make
sudo ./csrnet -s ./csrnet.wts
Loading weights: ./csrnet.wts
build engine successfully : ./csrnet.engine
// download images https://github.com/wang-xinyu/tensorrtx/assets/46584679/46bc4def-e573-44ae-996d-5d68927c78ff and copy to images
sudo ./csrnet -d ./images
// output e.g
// enqueueV2 time: 0.0323869s
// detect time:44ms
// people num :22.9101 write_path: ../images/data.jpg
```
# result
inference people num: 22.9101
<p align="center">
<img src= https://raw.githubusercontent.com/wang-xinyu/tensorrtx/dbf857d25f77bf64113fc99a745ccf4973bdd44e/Density_Plot.jpg>
</p>

16
csrnet/config.h Normal file
View File

@ -0,0 +1,16 @@
#pragma once
const static char *kInputTensorName = "data";
const static char *kOutputTensorName = "prob";
const static char *kEngineFile = "./csrnet.engine";
const static int kBatchSize = 1;
const static int MAX_INPUT_SIZE = 1440; // 32x
const static int MIN_INPUT_SIZE = 608;
const static int OPT_INPUT_W = 1152;
const static int OPT_INPUT_H = 640;
constexpr static int kMaxInputImageSize = MAX_INPUT_SIZE * MAX_INPUT_SIZE * 3;
constexpr static int kMaxOutputProbSize =
(MAX_INPUT_SIZE * MAX_INPUT_SIZE) >> 6;

536
csrnet/csrnet.cpp Normal file
View File

@ -0,0 +1,536 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include <chrono>
#include <config.h>
#include <cstring>
#include <dirent.h>
#include <fstream>
#include <iostream>
#include <logging.h>
#include <map>
#include <numeric>
#include <opencv2/opencv.hpp>
#include <vector>
using namespace nvinfer1;
#define CHECK(status) \
do { \
auto ret = (status); \
if (ret != 0) { \
std::cerr << "Cuda failure: " << ret << std::endl; \
abort(); \
} \
} while (0)
static Logger gLogger;
static char *kWTSFile = "";
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--) {
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
// clang-format off
/*
CSRNet(
(frontend): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace=True)
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(inplace=True)
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(inplace=True)
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace=True)
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace=True)
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace=True)
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace=True)
)
(backend): Sequential(
(0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2),
dilation=(2, 2)) (1): ReLU(inplace=True) (2): Conv2d(512, 512,
kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2)) (3):
ReLU(inplace=True) (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1,
1), padding=(2, 2), dilation=(2, 2)) (5): ReLU(inplace=True) (6):
Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2),
dilation=(2, 2)) (7): ReLU(inplace=True) (8): Conv2d(256, 128,
kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2)) (9):
ReLU(inplace=True) (10): Conv2d(128, 64, kernel_size=(3, 3), stride=(1,
1), padding=(2, 2), dilation=(2, 2)) (11): ReLU(inplace=True)
)
(output_layer): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
)
*/
// clang-format on
void doInference(IExecutionContext &context, float *input, float *output,
int input_h, int input_w) {
const ICudaEngine &engine = context.getEngine();
uint64_t input_size = 3 * input_h * input_w * sizeof(float);
uint64_t output_size = ((input_h * input_w) >> 6) * sizeof(float);
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void *buffers[2];
// In order to bind the buffers, we need to know the names of the input and
// output tensors. Note that indices are guaranteed to be less than
// IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(kInputTensorName);
const int outputIndex = engine.getBindingIndex(kOutputTensorName);
context.setBindingDimensions(inputIndex, Dims4(1, 3, input_h, input_w));
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], input_size));
CHECK(cudaMalloc(&buffers[outputIndex], output_size));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA
// output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, input_size,
cudaMemcpyHostToDevice, stream));
auto t1 = std::chrono::high_resolution_clock::now();
context.enqueueV2(buffers, stream, nullptr);
std::cout << "enqueueV2 time: "
<< std::chrono::duration<float>(
std::chrono::high_resolution_clock::now() - t1)
.count()
<< "s" << std::endl;
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size,
cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder,
IBuilderConfig *config, DataType dt) {
// INetworkDefinition *network = builder->createNetworkV2(0U);
const auto explicitBatch =
1U << static_cast<uint32_t>(
NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
ITensor *data = network->addInput(kInputTensorName, dt, Dims4{1, 3, -1, -1});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights(kWTSFile);
IConvolutionLayer *conv1 = network->addConvolutionNd(
*data, 64, DimsHW{3, 3}, weightMap["frontend.0.weight"],
weightMap["frontend.0.bias"]);
assert(conv1);
conv1->setStrideNd(DimsHW{1, 1});
conv1->setPaddingNd(DimsHW{1, 1});
IActivationLayer *relu1 =
network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(relu1);
auto conv2 = network->addConvolutionNd(*relu1->getOutput(0), 64, DimsHW{3, 3},
weightMap["frontend.2.weight"],
weightMap["frontend.2.bias"]);
assert(conv2);
conv2->setStrideNd(DimsHW{1, 1});
conv2->setPaddingNd(DimsHW{1, 1});
auto relu2 =
network->addActivation(*conv2->getOutput(0), ActivationType::kRELU);
assert(relu2);
auto pool1 = network->addPoolingNd(*relu2->getOutput(0), PoolingType::kMAX,
DimsHW{2, 2});
assert(pool1);
pool1->setStrideNd(DimsHW{2, 2});
auto conv3 = network->addConvolutionNd(
*pool1->getOutput(0), 128, DimsHW{3, 3}, weightMap["frontend.5.weight"],
weightMap["frontend.5.bias"]);
assert(conv3);
conv3->setStrideNd(DimsHW{1, 1});
conv3->setPaddingNd(DimsHW{1, 1});
auto relu3 =
network->addActivation(*conv3->getOutput(0), ActivationType::kRELU);
assert(relu3);
auto conv4 = network->addConvolutionNd(
*relu3->getOutput(0), 128, DimsHW{3, 3}, weightMap["frontend.7.weight"],
weightMap["frontend.7.bias"]);
assert(conv4);
conv4->setStrideNd(DimsHW{1, 1});
conv4->setPaddingNd(DimsHW{1, 1});
auto relu4 =
network->addActivation(*conv4->getOutput(0), ActivationType::kRELU);
assert(relu4);
auto pool2 = network->addPoolingNd(*relu4->getOutput(0), PoolingType::kMAX,
DimsHW{2, 2});
assert(pool2);
pool2->setStrideNd(DimsHW{2, 2});
auto conv5 = network->addConvolutionNd(
*pool2->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.10.weight"],
weightMap["frontend.10.bias"]);
assert(conv5);
conv5->setStrideNd(DimsHW{1, 1});
conv5->setPaddingNd(DimsHW{1, 1});
auto relu5 =
network->addActivation(*conv5->getOutput(0), ActivationType::kRELU);
assert(relu5);
auto conv6 = network->addConvolutionNd(
*relu5->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.12.weight"],
weightMap["frontend.12.bias"]);
assert(conv6);
conv6->setStrideNd(DimsHW{1, 1});
conv6->setPaddingNd(DimsHW{1, 1});
auto relu6 =
network->addActivation(*conv6->getOutput(0), ActivationType::kRELU);
assert(relu6);
auto conv7 = network->addConvolutionNd(
*relu6->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.14.weight"],
weightMap["frontend.14.bias"]);
assert(conv7);
conv7->setStrideNd(DimsHW{1, 1});
conv7->setPaddingNd(DimsHW{1, 1});
auto relu7 =
network->addActivation(*conv7->getOutput(0), ActivationType::kRELU);
assert(relu7);
auto pool3 = network->addPoolingNd(*relu7->getOutput(0), PoolingType::kMAX,
DimsHW{2, 2});
assert(pool3);
pool3->setStrideNd(DimsHW{2, 2});
auto conv8 = network->addConvolutionNd(
*pool3->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.17.weight"],
weightMap["frontend.17.bias"]);
assert(conv8);
conv8->setStrideNd(DimsHW{1, 1});
conv8->setPaddingNd(DimsHW{1, 1});
auto relu8 =
network->addActivation(*conv8->getOutput(0), ActivationType::kRELU);
assert(relu8);
auto conv9 = network->addConvolutionNd(
*relu8->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.19.weight"],
weightMap["frontend.19.bias"]);
assert(conv9);
conv9->setStrideNd(DimsHW{1, 1});
conv9->setPaddingNd(DimsHW{1, 1});
auto relu9 =
network->addActivation(*conv9->getOutput(0), ActivationType::kRELU);
assert(relu9);
auto conv10 = network->addConvolutionNd(
*relu9->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.21.weight"],
weightMap["frontend.21.bias"]);
assert(conv10);
conv10->setStrideNd(DimsHW{1, 1});
conv10->setPaddingNd(DimsHW{1, 1});
auto relu10 =
network->addActivation(*conv10->getOutput(0), ActivationType::kRELU);
assert(relu10);
// backend
auto conv11 = network->addConvolutionNd(
*relu10->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.0.weight"],
weightMap["backend.0.bias"]);
assert(conv11);
conv11->setPaddingNd(DimsHW{2, 2});
conv11->setStrideNd(DimsHW{1, 1});
conv11->setDilationNd(DimsHW{2, 2});
auto relu11 =
network->addActivation(*conv11->getOutput(0), ActivationType::kRELU);
assert(relu11);
auto conv12 = network->addConvolutionNd(
*relu11->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.2.weight"],
weightMap["backend.2.bias"]);
assert(conv12);
conv12->setPaddingNd(DimsHW{2, 2});
conv12->setStrideNd(DimsHW{1, 1});
conv12->setDilationNd(DimsHW{2, 2});
auto relu12 =
network->addActivation(*conv12->getOutput(0), ActivationType::kRELU);
assert(relu12);
auto conv13 = network->addConvolutionNd(
*relu12->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.4.weight"],
weightMap["backend.4.bias"]);
assert(conv13);
conv13->setPaddingNd(DimsHW{2, 2});
conv13->setStrideNd(DimsHW{1, 1});
conv13->setDilationNd(DimsHW{2, 2});
auto relu13 =
network->addActivation(*conv13->getOutput(0), ActivationType::kRELU);
assert(relu13);
auto conv14 = network->addConvolutionNd(
*relu13->getOutput(0), 256, DimsHW{3, 3}, weightMap["backend.6.weight"],
weightMap["backend.6.bias"]);
assert(conv14);
conv14->setPaddingNd(DimsHW{2, 2});
conv14->setStrideNd(DimsHW{1, 1});
conv14->setDilationNd(DimsHW{2, 2});
auto relu14 =
network->addActivation(*conv14->getOutput(0), ActivationType::kRELU);
assert(relu14);
auto conv15 = network->addConvolutionNd(
*relu14->getOutput(0), 128, DimsHW{3, 3}, weightMap["backend.8.weight"],
weightMap["backend.8.bias"]);
assert(conv15);
conv15->setPaddingNd(DimsHW{2, 2});
conv15->setStrideNd(DimsHW{1, 1});
conv15->setDilationNd(DimsHW{2, 2});
auto relu15 =
network->addActivation(*conv15->getOutput(0), ActivationType::kRELU);
assert(relu15);
auto conv16 = network->addConvolutionNd(
*relu15->getOutput(0), 64, DimsHW{3, 3}, weightMap["backend.10.weight"],
weightMap["backend.10.bias"]);
assert(conv16);
conv16->setPaddingNd(DimsHW{2, 2});
conv16->setStrideNd(DimsHW{1, 1});
conv16->setDilationNd(DimsHW{2, 2});
auto relu16 =
network->addActivation(*conv16->getOutput(0), ActivationType::kRELU);
assert(relu16);
auto conv17 = network->addConvolutionNd(
*relu16->getOutput(0), 1, DimsHW{1, 1}, weightMap["output_layer.weight"],
weightMap["output_layer.bias"]);
assert(conv17);
conv17->setStrideNd(DimsHW{1, 1});
conv17->getOutput(0)->setName(kOutputTensorName);
network->markOutput(*conv17->getOutput(0));
IOptimizationProfile *profile = builder->createOptimizationProfile();
profile->setDimensions(kInputTensorName, OptProfileSelector::kMIN,
Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
profile->setDimensions(kInputTensorName, OptProfileSelector::kOPT,
Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
profile->setDimensions(kInputTensorName, OptProfileSelector::kMAX,
Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
config->addOptimizationProfile(profile);
builder->setMaxBatchSize(kBatchSize);
config->setMaxWorkspaceSize(16 << 20);
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
printf("build engine successfully : %s\n", kEngineFile);
// Don't need the network any more
network->destroy();
// Release host memory
for (auto &mem : weightMap) {
free((void *)(mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream) {
// Create builder
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an
// engine
ICudaEngine *engine =
createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
config->destroy();
builder->destroy();
}
int read_files_in_dir(const char *p_dir_name,
std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent *p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 && strcmp(p_file->d_name, "..") != 0) {
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
int main(int argc, char **argv) {
if (argc != 3) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./csrnet -s ./csrnet.wts // serialize model to plan file"
<< std::endl;
std::cerr
<< "./csrnet -d ../images // deserialize plan file and run inference"
<< std::endl;
return -1;
}
char *trtModelStream{nullptr};
size_t size{0};
if (std::string(argv[1]) == "-s") {
IHostMemory *modelStream{nullptr};
kWTSFile = argv[2];
APIToModel(kBatchSize, &modelStream);
assert(modelStream != nullptr);
std::ofstream p(kEngineFile, std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char *>(modelStream->data()),
modelStream->size());
modelStream->destroy();
return 1;
} else if (std::string(argv[1]) == "-d") {
std::ifstream file(kEngineFile, std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
return -1;
}
IRuntime *runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext *context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
std::vector<std::string> file_names;
if (read_files_in_dir(argv[2], file_names) < 0) {
std::cout << "read_files_in_dir failed." << std::endl;
return -1;
}
std::vector<float> mean_value{0.406, 0.456, 0.485}; // BGR
std::vector<float> std_value{0.225, 0.224, 0.229};
int fcount = 0;
float *data = new float[kMaxInputImageSize];
float *prob = new float[kMaxOutputProbSize];
for (auto f : file_names) {
fcount++;
cv::Mat src_img = cv::imread(std::string(argv[2]) + "/" + f);
if (src_img.empty())
continue;
int i = 0;
for (int row = 0; row < src_img.rows; ++row) {
uchar *uc_pixel = src_img.data + row * src_img.step;
for (int col = 0; col < src_img.cols; ++col) {
data[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
data[i + src_img.rows * src_img.cols] =
(uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
data[i + 2 * src_img.rows * src_img.cols] =
(uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
uc_pixel += 3;
++i;
}
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, src_img.rows, src_img.cols);
auto end = std::chrono::system_clock::now();
std::cout << "detect time:"
<< std::chrono::duration_cast<std::chrono::milliseconds>(end -
start)
.count()
<< "ms" << std::endl;
float num = std::accumulate(
prob, prob + ((src_img.rows * src_img.cols) >> 6), 0.0f);
cv::Mat densityMap(src_img.rows >> 3, src_img.cols >> 3, CV_32FC1,
(void *)prob);
cv::Mat densityMapScaled;
cv::normalize(densityMap, densityMapScaled, 0, 255, cv::NORM_MINMAX,
CV_8UC1);
cv::Mat densityColorMap;
cv::applyColorMap(densityMapScaled, densityColorMap, cv::COLORMAP_VIRIDIS);
cv::resize(densityColorMap, densityColorMap, src_img.size());
cv::addWeighted(densityColorMap, 0.5, src_img, 0.5, 0, src_img);
// write to jpg
cv::putText(src_img, std::string("people num: ") + std::to_string(num),
cv::Point(10, 50), cv::FONT_HERSHEY_SIMPLEX, 0.5,
cv::Scalar(255, 255, 255), 1);
std::string write_path = std::string(argv[2]) + "result_" + f;
std::cout << "people num :" << num << " write_path: " << write_path
<< std::endl;
cv::imwrite(write_path, src_img);
}
delete[] data;
delete[] prob;
return 0;
}

31
csrnet/gen_wts.py Normal file
View File

@ -0,0 +1,31 @@
from torch.nn.modules import module
from model import CSRNet
import torch
import os
import struct
save_path = os.path.join(os.path.dirname(
__file__), "output", os.path.basename(__file__).split('.')[0])
os.makedirs(save_path, exist_ok=True)
wts_file = os.path.join(save_path, "csrnet.wts")
# load model
model_path = "partBmodel_best.pth.tar"
model = CSRNet()
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['state_dict'])
# save to wts
print(f'Writing into {wts_file}')
with open(wts_file, 'w') as f:
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')

502
csrnet/logging.h Normal file
View File

@ -0,0 +1,502 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include "macros.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf {
public:
LogStreamConsumerBuffer(std::ostream &stream, const std::string &prefix,
bool shouldLog)
: mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
LogStreamConsumerBuffer(LogStreamConsumerBuffer &&other)
: mOutput(other.mOutput) {}
~LogStreamConsumerBuffer() {
// std::streambuf::pbase() gives a pointer to the beginning of the buffered
// part of the output sequence std::streambuf::pptr() gives a pointer to the
// current position of the output sequence if the pointer to the beginning
// is not equal to the pointer to the current position, call putOutput() to
// log the output to the stream
if (pbase() != pptr()) {
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents
// into the stream, resetting the buffer and flushing the stream
virtual int sync() {
putOutput();
return 0;
}
void putOutput() {
if (mShouldLog) {
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm *tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon
<< "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday
<< "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year
<< "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour
<< ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec
<< "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into
// the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
private:
std::ostream &mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before
//! std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase {
public:
LogStreamConsumerBase(std::ostream &stream, const std::string &prefix,
bool shouldLog)
: mBuffer(stream, prefix, shouldLog) {}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when
//! logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the
//! LogStreamConsumerBuffer member field in LogStreamConsumer and then the
//! address of the buffer is passed to std::ostream. This is necessary to
//! prevent the address of an uninitialized buffer from being passed to
//! std::ostream. Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
public:
//! \brief Creates a LogStreamConsumer which logs messages with level
//! severity.
//! Reportable severity determines if the messages are severe enough to be
//! logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity),
severityPrefix(severity),
severity <= reportableSeverity),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(severity <= reportableSeverity), mSeverity(severity) {}
LogStreamConsumer(LogStreamConsumer &&other)
: LogStreamConsumerBase(severityOstream(other.mSeverity),
severityPrefix(other.mSeverity),
other.mShouldLog),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(other.mShouldLog), mSeverity(other.mSeverity) {}
void setReportableSeverity(Severity reportableSeverity) {
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream &severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and
//! samples to log information to the console, and supports logging two types of
//! messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or
//! internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to
//! emitting directly to stdout/stderr is that the logic for controlling the
//! verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results
//! to a file in some standard format (for example, JUnit XML), and providing
//! additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits
//! directly from the nvinfer1::ILogger interface, which is problematic since
//! there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to
//! access the ILogger) we can refactor the class to eliminate the inheritance
//! and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger {
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity) {}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult {
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger
//! associated with this Logger \return The nvinfer1::ILogger associated with
//! this Logger
//!
//! TODO Once all samples are updated to use this method to register the
//! logger with TensorRT, we can eliminate the inheritance of Logger from
//! ILogger
//!
nvinfer1::ILogger &getTRTLogger() { return *this; }
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will
//! eventually go away once we eliminate the inheritance from
//! nvinfer1::ILogger
//!
void log(Severity severity, const char *msg) TRT_NOEXCEPT override {
LogStreamConsumer(mReportableSeverity, severity)
<< "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of
//! this level or higher.
//!
void setReportableSeverity(Severity severity) {
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print
//! test results. The sample must call Logger::defineTest() in order to obtain
//! a TestAtom that can be used with Logger::reportTest{Start,End}().
//!
class TestAtom {
public:
TestAtom(TestAtom &&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string &name, const std::string &cmdline)
: mStarted(started), mName(name), mCmdline(cmdline) {}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting
//! with
//! "TensorRT" and containing dot-separated strings
//! containing the characters [A-Za-z0-9_]. For example,
//! "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string &name,
const std::string &cmdline) {
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an
//! array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string &name, int argc,
char const *const *argv) {
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom &testAtom) {
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of
//! TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom &testAtom, TestResult result) {
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom &testAtom) {
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom &testAtom) {
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom &testAtom) {
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom &testAtom, bool pass) {
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const { return mReportableSeverity; }
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the
//! given severity
//!
static const char *severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message
//! with the given result
//!
static const char *testResultString(TestResult result) {
switch (result) {
case TestResult::kRUNNING:
return "RUNNING";
case TestResult::kPASSED:
return "PASSED";
case TestResult::kFAILED:
return "FAILED";
case TestResult::kWAIVED:
return "WAIVED";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the
//! given severity
//!
static std::ostream &severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom &testAtom, TestResult result) {
severityOstream(Severity::kINFO)
<< "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const *const *argv) {
std::stringstream ss;
for (int i = 0; i < argc; i++) {
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace {
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages
//! of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger &logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages
//! of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger &logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages
//! of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger &logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages
//! of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger &logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages
//! of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger &logger) {
return LogStreamConsumer(logger.getReportableSeverity(),
Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

12
csrnet/macros.h Normal file
View File

@ -0,0 +1,12 @@
#ifndef __MACROS_H
#define __MACROS_H
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif
#endif // __MACROS_H

38
dbnet/CMakeLists.txt Normal file
View File

@ -0,0 +1,38 @@
cmake_minimum_required(VERSION 2.6)
project(dbnet)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
aux_source_directory(. DIRSRCS)
# clipper
include_directories(./ ./clipper)
add_subdirectory(clipper)
add_executable(dbnet ${DIRSRCS})
target_link_libraries(dbnet clipper)
target_link_libraries(dbnet nvinfer)
target_link_libraries(dbnet cudart)
target_link_libraries(dbnet ${OpenCV_LIBS})
add_definitions(-O2 -pthread)

56
dbnet/README.md Normal file
View File

@ -0,0 +1,56 @@
# DBNet
The Pytorch implementation is [DBNet](https://github.com/BaofengZan/DBNet.pytorch).
<p align="center">
<img src="https://user-images.githubusercontent.com/25873202/113959270-1eb8c600-9855-11eb-9c4d-1e6dc8e38a17.jpg">
</p>
## How to Run
* 1. generate `.wts`
Download code and model from [DBNet](https://github.com/BaofengZan/DBNet.pytorch) and config your environments.
Go to file`tools/predict.py`, set `--save_wts` as `True`, then run, the `DBNet.wts` will be generated.
Onnx can also be exported, just need to set `--onnx` as `True`.
* 2. cmake and make
```
mkdir build
cd build
cmake ..
make
cp /your_wts_path/DBNet.wts .
sudo ./dbnet -s // serialize model to plan file i.e. 'DBNet.engine'
sudo ./dbnet -d ./test_imgs // deserialize plan file and run inference, all images in test_imgs folder will be processed.
```
## For windows
https://github.com/BaofengZan/DBNet-TensorRT
## Todo
- [x] 1. In `common.hpp`, the following two functions can be merged.
```c++
ILayer* convBnLeaky(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, bool bias = true)
```
```c++
ILayer* convBnLeaky2(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, bool bias = true)
```
- [x] 2. The postprocess method here should be optimized, which is a little different from pytorch side.
- [x] 3. The input image here is resized to `640 x 640` directly, while the pytorch side is using `letterbox` method.

4
dbnet/clipper/CMakeLists.txt Normal file
View File

@ -0,0 +1,4 @@
cmake_minimum_required(VERSION 2.6)
aux_source_directory(. DIR_CLIPPER_SRCS)
add_library(clipper ${DIR_CLIPPER_SRCS})

4629
dbnet/clipper/clipper.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

406
dbnet/clipper/clipper.hpp Normal file
View File

@ -0,0 +1,406 @@
/*******************************************************************************
* *
* Author : Angus Johnson *
* Version : 6.4.2 *
* Date : 27 February 2017 *
* Website : http://www.angusj.com *
* Copyright : Angus Johnson 2010-2017 *
* *
* License: *
* Use, modification & distribution is subject to Boost Software License Ver 1. *
* http://www.boost.org/LICENSE_1_0.txt *
* *
* Attributions: *
* The code in this library is an extension of Bala Vatti's clipping algorithm: *
* "A generic solution to polygon clipping" *
* Communications of the ACM, Vol 35, Issue 7 (July 1992) pp 56-63. *
* http://portal.acm.org/citation.cfm?id=129906 *
* *
* Computer graphics and geometric modeling: implementation and algorithms *
* By Max K. Agoston *
* Springer; 1 edition (January 4, 2005) *
* http://books.google.com/books?q=vatti+clipping+agoston *
* *
* See also: *
* "Polygon Offsetting by Computing Winding Numbers" *
* Paper no. DETC2005-85513 pp. 565-575 *
* ASME 2005 International Design Engineering Technical Conferences *
* and Computers and Information in Engineering Conference (IDETC/CIE2005) *
* September 24-28, 2005 , Long Beach, California, USA *
* http://www.me.berkeley.edu/~mcmains/pubs/DAC05OffsetPolygon.pdf *
* *
*******************************************************************************/
#ifndef clipper_hpp
#define clipper_hpp
#define CLIPPER_VERSION "6.4.2"
//use_int32: When enabled 32bit ints are used instead of 64bit ints. This
//improve performance but coordinate values are limited to the range +/- 46340
//#define use_int32
//use_xyz: adds a Z member to IntPoint. Adds a minor cost to perfomance.
//#define use_xyz
//use_lines: Enables line clipping. Adds a very minor cost to performance.
#define use_lines
//use_deprecated: Enables temporary support for the obsolete functions
//#define use_deprecated
#include <vector>
#include <list>
#include <set>
#include <stdexcept>
#include <cstring>
#include <cstdlib>
#include <ostream>
#include <functional>
#include <queue>
namespace ClipperLib {
enum ClipType { ctIntersection, ctUnion, ctDifference, ctXor };
enum PolyType { ptSubject, ptClip };
//By far the most widely used winding rules for polygon filling are
//EvenOdd & NonZero (GDI, GDI+, XLib, OpenGL, Cairo, AGG, Quartz, SVG, Gr32)
//Others rules include Positive, Negative and ABS_GTR_EQ_TWO (only in OpenGL)
//see http://glprogramming.com/red/chapter11.html
enum PolyFillType { pftEvenOdd, pftNonZero, pftPositive, pftNegative };
#ifdef use_int32
typedef int cInt;
static cInt const loRange = 0x7FFF;
static cInt const hiRange = 0x7FFF;
#else
typedef signed long long cInt;
static cInt const loRange = 0x3FFFFFFF;
static cInt const hiRange = 0x3FFFFFFFFFFFFFFFLL;
typedef signed long long long64; //used by Int128 class
typedef unsigned long long ulong64;
#endif
struct IntPoint {
cInt X;
cInt Y;
#ifdef use_xyz
cInt Z;
IntPoint(cInt x = 0, cInt y = 0, cInt z = 0): X(x), Y(y), Z(z) {};
#else
IntPoint(cInt x = 0, cInt y = 0): X(x), Y(y) {};
#endif
friend inline bool operator== (const IntPoint& a, const IntPoint& b)
{
return a.X == b.X && a.Y == b.Y;
}
friend inline bool operator!= (const IntPoint& a, const IntPoint& b)
{
return a.X != b.X || a.Y != b.Y;
}
};
//------------------------------------------------------------------------------
typedef std::vector< IntPoint > Path;
typedef std::vector< Path > Paths;
inline Path& operator <<(Path& poly, const IntPoint& p) {poly.push_back(p); return poly;}
inline Paths& operator <<(Paths& polys, const Path& p) {polys.push_back(p); return polys;}
std::ostream& operator <<(std::ostream &s, const IntPoint &p);
std::ostream& operator <<(std::ostream &s, const Path &p);
std::ostream& operator <<(std::ostream &s, const Paths &p);
struct DoublePoint
{
double X;
double Y;
DoublePoint(double x = 0, double y = 0) : X(x), Y(y) {}
DoublePoint(IntPoint ip) : X((double)ip.X), Y((double)ip.Y) {}
};
//------------------------------------------------------------------------------
#ifdef use_xyz
typedef void (*ZFillCallback)(IntPoint& e1bot, IntPoint& e1top, IntPoint& e2bot, IntPoint& e2top, IntPoint& pt);
#endif
enum InitOptions {ioReverseSolution = 1, ioStrictlySimple = 2, ioPreserveCollinear = 4};
enum JoinType {jtSquare, jtRound, jtMiter};
enum EndType {etClosedPolygon, etClosedLine, etOpenButt, etOpenSquare, etOpenRound};
class PolyNode;
typedef std::vector< PolyNode* > PolyNodes;
class PolyNode
{
public:
PolyNode();
virtual ~PolyNode(){};
Path Contour;
PolyNodes Childs;
PolyNode* Parent;
PolyNode* GetNext() const;
bool IsHole() const;
bool IsOpen() const;
int ChildCount() const;
private:
//PolyNode& operator =(PolyNode& other);
unsigned Index; //node index in Parent.Childs
bool m_IsOpen;
JoinType m_jointype;
EndType m_endtype;
PolyNode* GetNextSiblingUp() const;
void AddChild(PolyNode& child);
friend class Clipper; //to access Index
friend class ClipperOffset;
};
class PolyTree: public PolyNode
{
public:
~PolyTree(){ Clear(); };
PolyNode* GetFirst() const;
void Clear();
int Total() const;
private:
//PolyTree& operator =(PolyTree& other);
PolyNodes AllNodes;
friend class Clipper; //to access AllNodes
};
bool Orientation(const Path &poly);
double Area(const Path &poly);
int PointInPolygon(const IntPoint &pt, const Path &path);
void SimplifyPolygon(const Path &in_poly, Paths &out_polys, PolyFillType fillType = pftEvenOdd);
void SimplifyPolygons(const Paths &in_polys, Paths &out_polys, PolyFillType fillType = pftEvenOdd);
void SimplifyPolygons(Paths &polys, PolyFillType fillType = pftEvenOdd);
void CleanPolygon(const Path& in_poly, Path& out_poly, double distance = 1.415);
void CleanPolygon(Path& poly, double distance = 1.415);
void CleanPolygons(const Paths& in_polys, Paths& out_polys, double distance = 1.415);
void CleanPolygons(Paths& polys, double distance = 1.415);
void MinkowskiSum(const Path& pattern, const Path& path, Paths& solution, bool pathIsClosed);
void MinkowskiSum(const Path& pattern, const Paths& paths, Paths& solution, bool pathIsClosed);
void MinkowskiDiff(const Path& poly1, const Path& poly2, Paths& solution);
void PolyTreeToPaths(const PolyTree& polytree, Paths& paths);
void ClosedPathsFromPolyTree(const PolyTree& polytree, Paths& paths);
void OpenPathsFromPolyTree(PolyTree& polytree, Paths& paths);
void ReversePath(Path& p);
void ReversePaths(Paths& p);
struct IntRect { cInt left; cInt top; cInt right; cInt bottom; };
//enums that are used internally ...
enum EdgeSide { esLeft = 1, esRight = 2};
//forward declarations (for stuff used internally) ...
struct TEdge;
struct IntersectNode;
struct LocalMinimum;
struct OutPt;
struct OutRec;
struct Join;
typedef std::vector < OutRec* > PolyOutList;
typedef std::vector < TEdge* > EdgeList;
typedef std::vector < Join* > JoinList;
typedef std::vector < IntersectNode* > IntersectList;
//------------------------------------------------------------------------------
//ClipperBase is the ancestor to the Clipper class. It should not be
//instantiated directly. This class simply abstracts the conversion of sets of
//polygon coordinates into edge objects that are stored in a LocalMinima list.
class ClipperBase
{
public:
ClipperBase();
virtual ~ClipperBase();
virtual bool AddPath(const Path &pg, PolyType PolyTyp, bool Closed);
bool AddPaths(const Paths &ppg, PolyType PolyTyp, bool Closed);
virtual void Clear();
IntRect GetBounds();
bool PreserveCollinear() {return m_PreserveCollinear;};
void PreserveCollinear(bool value) {m_PreserveCollinear = value;};
protected:
void DisposeLocalMinimaList();
TEdge* AddBoundsToLML(TEdge *e, bool IsClosed);
virtual void Reset();
TEdge* ProcessBound(TEdge* E, bool IsClockwise);
void InsertScanbeam(const cInt Y);
bool PopScanbeam(cInt &Y);
bool LocalMinimaPending();
bool PopLocalMinima(cInt Y, const LocalMinimum *&locMin);
OutRec* CreateOutRec();
void DisposeAllOutRecs();
void DisposeOutRec(PolyOutList::size_type index);
void SwapPositionsInAEL(TEdge *edge1, TEdge *edge2);
void DeleteFromAEL(TEdge *e);
void UpdateEdgeIntoAEL(TEdge *&e);
typedef std::vector<LocalMinimum> MinimaList;
MinimaList::iterator m_CurrentLM;
MinimaList m_MinimaList;
bool m_UseFullRange;
EdgeList m_edges;
bool m_PreserveCollinear;
bool m_HasOpenPaths;
PolyOutList m_PolyOuts;
TEdge *m_ActiveEdges;
typedef std::priority_queue<cInt> ScanbeamList;
ScanbeamList m_Scanbeam;
};
//------------------------------------------------------------------------------
class Clipper : public virtual ClipperBase
{
public:
Clipper(int initOptions = 0);
bool Execute(ClipType clipType,
Paths &solution,
PolyFillType fillType = pftEvenOdd);
bool Execute(ClipType clipType,
Paths &solution,
PolyFillType subjFillType,
PolyFillType clipFillType);
bool Execute(ClipType clipType,
PolyTree &polytree,
PolyFillType fillType = pftEvenOdd);
bool Execute(ClipType clipType,
PolyTree &polytree,
PolyFillType subjFillType,
PolyFillType clipFillType);
bool ReverseSolution() { return m_ReverseOutput; };
void ReverseSolution(bool value) {m_ReverseOutput = value;};
bool StrictlySimple() {return m_StrictSimple;};
void StrictlySimple(bool value) {m_StrictSimple = value;};
//set the callback function for z value filling on intersections (otherwise Z is 0)
#ifdef use_xyz
void ZFillFunction(ZFillCallback zFillFunc);
#endif
protected:
virtual bool ExecuteInternal();
private:
JoinList m_Joins;
JoinList m_GhostJoins;
IntersectList m_IntersectList;
ClipType m_ClipType;
typedef std::list<cInt> MaximaList;
MaximaList m_Maxima;
TEdge *m_SortedEdges;
bool m_ExecuteLocked;
PolyFillType m_ClipFillType;
PolyFillType m_SubjFillType;
bool m_ReverseOutput;
bool m_UsingPolyTree;
bool m_StrictSimple;
#ifdef use_xyz
ZFillCallback m_ZFill; //custom callback
#endif
void SetWindingCount(TEdge& edge);
bool IsEvenOddFillType(const TEdge& edge) const;
bool IsEvenOddAltFillType(const TEdge& edge) const;
void InsertLocalMinimaIntoAEL(const cInt botY);
void InsertEdgeIntoAEL(TEdge *edge, TEdge* startEdge);
void AddEdgeToSEL(TEdge *edge);
bool PopEdgeFromSEL(TEdge *&edge);
void CopyAELToSEL();
void DeleteFromSEL(TEdge *e);
void SwapPositionsInSEL(TEdge *edge1, TEdge *edge2);
bool IsContributing(const TEdge& edge) const;
bool IsTopHorz(const cInt XPos);
void DoMaxima(TEdge *e);
void ProcessHorizontals();
void ProcessHorizontal(TEdge *horzEdge);
void AddLocalMaxPoly(TEdge *e1, TEdge *e2, const IntPoint &pt);
OutPt* AddLocalMinPoly(TEdge *e1, TEdge *e2, const IntPoint &pt);
OutRec* GetOutRec(int idx);
void AppendPolygon(TEdge *e1, TEdge *e2);
void IntersectEdges(TEdge *e1, TEdge *e2, IntPoint &pt);
OutPt* AddOutPt(TEdge *e, const IntPoint &pt);
OutPt* GetLastOutPt(TEdge *e);
bool ProcessIntersections(const cInt topY);
void BuildIntersectList(const cInt topY);
void ProcessIntersectList();
void ProcessEdgesAtTopOfScanbeam(const cInt topY);
void BuildResult(Paths& polys);
void BuildResult2(PolyTree& polytree);
void SetHoleState(TEdge *e, OutRec *outrec);
void DisposeIntersectNodes();
bool FixupIntersectionOrder();
void FixupOutPolygon(OutRec &outrec);
void FixupOutPolyline(OutRec &outrec);
bool IsHole(TEdge *e);
bool FindOwnerFromSplitRecs(OutRec &outRec, OutRec *&currOrfl);
void FixHoleLinkage(OutRec &outrec);
void AddJoin(OutPt *op1, OutPt *op2, const IntPoint offPt);
void ClearJoins();
void ClearGhostJoins();
void AddGhostJoin(OutPt *op, const IntPoint offPt);
bool JoinPoints(Join *j, OutRec* outRec1, OutRec* outRec2);
void JoinCommonEdges();
void DoSimplePolygons();
void FixupFirstLefts1(OutRec* OldOutRec, OutRec* NewOutRec);
void FixupFirstLefts2(OutRec* InnerOutRec, OutRec* OuterOutRec);
void FixupFirstLefts3(OutRec* OldOutRec, OutRec* NewOutRec);
#ifdef use_xyz
void SetZ(IntPoint& pt, TEdge& e1, TEdge& e2);
#endif
};
//------------------------------------------------------------------------------
class ClipperOffset
{
public:
ClipperOffset(double miterLimit = 2.0, double roundPrecision = 0.25);
~ClipperOffset();
void AddPath(const Path& path, JoinType joinType, EndType endType);
void AddPaths(const Paths& paths, JoinType joinType, EndType endType);
void Execute(Paths& solution, double delta);
void Execute(PolyTree& solution, double delta);
void Clear();
double MiterLimit;
double ArcTolerance;
private:
Paths m_destPolys;
Path m_srcPoly;
Path m_destPoly;
std::vector<DoublePoint> m_normals;
double m_delta, m_sinA, m_sin, m_cos;
double m_miterLim, m_StepsPerRad;
IntPoint m_lowest;
PolyNode m_polyNodes;
void FixOrientations();
void DoOffset(double delta);
void OffsetPoint(int j, int& k, JoinType jointype);
void DoSquare(int j, int k);
void DoMiter(int j, int k, double r);
void DoRound(int j, int k);
};
//------------------------------------------------------------------------------
class clipperException : public std::exception
{
public:
clipperException(const char* description): m_descr(description) {}
virtual ~clipperException() throw() {}
virtual const char* what() const throw() {return m_descr.c_str();}
private:
std::string m_descr;
};
//------------------------------------------------------------------------------
} //ClipperLib namespace
#endif //clipper_hpp

177
dbnet/common.hpp Normal file
View File

@ -0,0 +1,177 @@
#ifndef DBNET_COMMON_H_
#define DBNET_COMMON_H_
#include <iostream>
#include <fstream>
#include <map>
#include <sstream>
#include <vector>
#include <opencv2/opencv.hpp>
#include "dirent.h"
#include "NvInfer.h"
#include <chrono>
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
using namespace nvinfer1;
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--) {
Weights wt{ DataType::kFLOAT, nullptr, 0 };
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + ".weight"].values;
float *beta = (float*)weightMap[lname + ".bias"].values;
float *mean = (float*)weightMap[lname + ".running_mean"].values;
float *var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{ DataType::kFLOAT, scval, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* convBnLeaky(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, std::string bnname, bool bias = true) {
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
int p = ksize / 2;
IConvolutionLayer* conv1 = nullptr;
if (bias) {
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".weight"], weightMap[lname + ".bias"]);
}
else {
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".weight"], emptywts);
}
assert(conv1);
conv1->setStrideNd(DimsHW{ s, s });
conv1->setPaddingNd(DimsHW{ p, p });
conv1->setNbGroups(g);
//IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-4);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname.substr(0, lname.find_last_of(".")) + bnname, 1e-5);
auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
lr->setAlpha(0.1);
return lr;
}
IActivationLayer* basicBlock(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int stride, std::string lname) {
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ 3, 3 }, weightMap[lname + "conv1.weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{ stride, stride });
conv1->setPaddingNd(DimsHW{ 1, 1 });
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "bn1", 1e-5);
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{ 3, 3 }, weightMap[lname + "conv2.weight"], emptywts);
assert(conv2);
conv2->setPaddingNd(DimsHW{ 1, 1 });
IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "bn2", 1e-5);
IElementWiseLayer* ew1;
if (inch != outch) {
IConvolutionLayer* conv3 = network->addConvolutionNd(input, outch, DimsHW{ 1, 1 }, weightMap[lname + "downsample.0.weight"], emptywts);
assert(conv3);
conv3->setStrideNd(DimsHW{ stride, stride });
IScaleLayer* bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + "downsample.1", 1e-5);
ew1 = network->addElementWise(*bn3->getOutput(0), *bn2->getOutput(0), ElementWiseOperation::kSUM);
}
else {
ew1 = network->addElementWise(input, *bn2->getOutput(0), ElementWiseOperation::kSUM);
}
IActivationLayer* relu2 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU);
assert(relu2);
return relu2;
}
int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
#endif

553
dbnet/dbnet.cpp Normal file
View File

@ -0,0 +1,553 @@
#include <iostream>
#include <chrono>
#include "cuda_runtime_api.h"
#include "logging.h"
#include "common.hpp"
#include <math.h>
#include "clipper.hpp"
#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define EXPANDRATIO 1.5
#define BOX_MINI_SIZE 5
#define SCORE_THRESHOLD 0.3
#define BOX_THRESHOLD 0.7
static const int SHORT_INPUT = 640;
static const int MAX_INPUT_SIZE = 1440; // 32x
static const int MIN_INPUT_SIZE = 608;
static const int OPT_INPUT_W = 1152;
static const int OPT_INPUT_H = 640;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "out";
static Logger gLogger;
cv::RotatedRect expandBox(cv::Point2f temp[], float ratio)
{
ClipperLib::Path path = {
{ClipperLib::cInt(temp[0].x), ClipperLib::cInt(temp[0].y)},
{ClipperLib::cInt(temp[1].x), ClipperLib::cInt(temp[1].y)},
{ClipperLib::cInt(temp[2].x), ClipperLib::cInt(temp[2].y)},
{ClipperLib::cInt(temp[3].x), ClipperLib::cInt(temp[3].y)}};
double area = ClipperLib::Area(path);
double distance;
double length = 0.0;
for (int i = 0; i < 4; i++) {
length = length + sqrtf(powf((temp[i].x - temp[(i + 1) % 4].x), 2) +
powf((temp[i].y - temp[(i + 1) % 4].y), 2));
}
distance = area * ratio / length;
ClipperLib::ClipperOffset offset;
offset.AddPath(path, ClipperLib::JoinType::jtRound,
ClipperLib::EndType::etClosedPolygon);
ClipperLib::Paths paths;
offset.Execute(paths, distance);
std::vector<cv::Point> contour;
for (int i = 0; i < paths[0].size(); i++) {
contour.emplace_back(paths[0][i].X, paths[0][i].Y);
}
offset.Clear();
return cv::minAreaRect(contour);
}
float paddimg(cv::Mat& In_Out_img, int shortsize = 960) {
int w = In_Out_img.cols;
int h = In_Out_img.rows;
float scale = 1.f;
if (w < h) {
scale = (float)shortsize / w;
h = scale * h;
w = shortsize;
}
else {
scale = (float)shortsize / h;
w = scale * w;
h = shortsize;
}
if (h % 32 != 0) {
h = (h / 32 + 1) * 32;
}
if (w % 32 != 0) {
w = (w / 32 + 1) * 32;
}
cv::resize(In_Out_img, In_Out_img, cv::Size(w, h));
return scale;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims4{ 1, 3, -1, -1 });
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("./DBNet.wts");
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
/* ------ Resnet18 backbone------ */
// Add convolution layer with 6 outputs and a 5x5 filter.
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{ 7, 7 }, weightMap["backbone.conv1.weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{ 2, 2 });
conv1->setPaddingNd(DimsHW{ 3, 3 });
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "backbone.bn1", 1e-5);
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
assert(pool1);
pool1->setStrideNd(DimsHW{ 2, 2 });
pool1->setPaddingNd(DimsHW{ 1, 1 });
IActivationLayer* relu2 = basicBlock(network, weightMap, *pool1->getOutput(0), 64, 64, 1, "backbone.layer1.0.");
IActivationLayer* relu3 = basicBlock(network, weightMap, *relu2->getOutput(0), 64, 64, 1, "backbone.layer1.1."); // x2
IActivationLayer* relu4 = basicBlock(network, weightMap, *relu3->getOutput(0), 64, 128, 2, "backbone.layer2.0.");
IActivationLayer* relu5 = basicBlock(network, weightMap, *relu4->getOutput(0), 128, 128, 1, "backbone.layer2.1."); // x3
IActivationLayer* relu6 = basicBlock(network, weightMap, *relu5->getOutput(0), 128, 256, 2, "backbone.layer3.0.");
IActivationLayer* relu7 = basicBlock(network, weightMap, *relu6->getOutput(0), 256, 256, 1, "backbone.layer3.1."); //x4
IActivationLayer* relu8 = basicBlock(network, weightMap, *relu7->getOutput(0), 256, 512, 2, "backbone.layer4.0.");
IActivationLayer* relu9 = basicBlock(network, weightMap, *relu8->getOutput(0), 512, 512, 1, "backbone.layer4.1."); //x5
/* ------- FPN neck ------- */
ILayer* p5 = convBnLeaky(network, weightMap, *relu9->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c5.conv", ".bn"); // k=1 s = 1 p = k/2=1/2=0
ILayer* c4_1 = convBnLeaky(network, weightMap, *relu7->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c4.conv", ".bn");
float *deval = reinterpret_cast<float*>(malloc(sizeof(float) * 64 * 2 * 2));
for (int i = 0; i < 64 * 2 * 2; i++) {
deval[i] = 1.0;
}
Weights deconvwts1{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p4_1 = network->addDeconvolutionNd(*p5->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts1, emptywts);
p4_1->setStrideNd(DimsHW{ 2, 2 });
p4_1->setNbGroups(64);
weightMap["deconv1"] = deconvwts1;
IElementWiseLayer* p4_add = network->addElementWise(*p4_1->getOutput(0), *c4_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p4 = convBnLeaky(network, weightMap, *p4_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p4.conv", ".bn"); // smooth
ILayer* c3_1 = convBnLeaky(network, weightMap, *relu5->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c3.conv", ".bn");
Weights deconvwts2{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_1 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts2, emptywts);
p3_1->setStrideNd(DimsHW{ 2, 2 });
p3_1->setNbGroups(64);
IElementWiseLayer* p3_add = network->addElementWise(*p3_1->getOutput(0), *c3_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p3 = convBnLeaky(network, weightMap, *p3_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p3.conv", ".bn"); // smooth
ILayer* c2_1 = convBnLeaky(network, weightMap, *relu3->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c2.conv", ".bn");
Weights deconvwts3{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p2_1 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts3, emptywts);
p2_1->setStrideNd(DimsHW{ 2, 2 });
p2_1->setNbGroups(64);
IElementWiseLayer* p2_add = network->addElementWise(*p2_1->getOutput(0), *c2_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p2 = convBnLeaky(network, weightMap, *p2_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p2.conv", ".bn"); // smooth
Weights deconvwts4{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_up_p2 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts4, emptywts);
p3_up_p2->setStrideNd(DimsHW{ 2, 2 });
p3_up_p2->setNbGroups(64);
float *deval2 = reinterpret_cast<float*>(malloc(sizeof(float) * 64 * 8 * 8));
for (int i = 0; i < 64 * 8 * 8; i++) {
deval2[i] = 1.0;
}
Weights deconvwts5{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p4_up_p2 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts5, emptywts);
p4_up_p2->setPaddingNd(DimsHW{ 2, 2 });
p4_up_p2->setStrideNd(DimsHW{ 4, 4 });
p4_up_p2->setNbGroups(64);
weightMap["deconv2"] = deconvwts5;
Weights deconvwts6{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p5_up_p2 = network->addDeconvolutionNd(*p5->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts6, emptywts);
p5_up_p2->setStrideNd(DimsHW{ 8, 8 });
p5_up_p2->setNbGroups(64);
// torch.cat([p2, p3, p4, p5], dim=1)
ITensor* inputTensors[] = { p2->getOutput(0), p3_up_p2->getOutput(0), p4_up_p2->getOutput(0), p5_up_p2->getOutput(0) };
IConcatenationLayer* neck_cat = network->addConcatenation(inputTensors, 4);
ILayer* neck_out = convBnLeaky(network, weightMap, *neck_cat->getOutput(0), 256, 3, 1, 1, "neck.conv.0", ".1"); // smooth
assert(neck_out);
ILayer* binarize1 = convBnLeaky(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.binarize.0", ".1"); //
Weights deconvwts7{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup = network->addDeconvolutionNd(*binarize1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts7, emptywts);
binarizeup->setStrideNd(DimsHW{ 2, 2 });
binarizeup->setNbGroups(64);
IScaleLayer* binarizebn1 = addBatchNorm2d(network, weightMap, *binarizeup->getOutput(0), "head.binarize.4", 1e-5);
IActivationLayer* binarizerelu1 = network->addActivation(*binarizebn1->getOutput(0), ActivationType::kRELU);
assert(binarizerelu1);
Weights deconvwts8{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup2 = network->addDeconvolutionNd(*binarizerelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts8, emptywts);
binarizeup2->setStrideNd(DimsHW{ 2, 2 });
binarizeup2->setNbGroups(64);
IConvolutionLayer* binarize3 = network->addConvolutionNd(*binarizeup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.binarize.7.weight"], weightMap["head.binarize.7.bias"]);
assert(binarize3);
binarize3->setStrideNd(DimsHW{ 1, 1 });
binarize3->setPaddingNd(DimsHW{ 1, 1 });
IActivationLayer* binarize4 = network->addActivation(*binarize3->getOutput(0), ActivationType::kSIGMOID);
assert(binarize4);
//threshold_maps = self.thresh(x)
ILayer* thresh1 = convBnLeaky(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.thresh.0", ".1", false); //
Weights deconvwts9{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup = network->addDeconvolutionNd(*thresh1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts9, emptywts);
threshup->setStrideNd(DimsHW{ 2, 2 });
threshup->setNbGroups(64);
IConvolutionLayer* thresh2 = network->addConvolutionNd(*threshup->getOutput(0), 64, DimsHW{ 3, 3 }, weightMap["head.thresh.3.1.weight"], weightMap["head.thresh.3.1.bias"]);
assert(thresh2);
thresh2->setStrideNd(DimsHW{ 1, 1 });
thresh2->setPaddingNd(DimsHW{ 1, 1 });
IScaleLayer* threshbn1 = addBatchNorm2d(network, weightMap, *thresh2->getOutput(0), "head.thresh.4", 1e-5);
IActivationLayer* threshrelu1 = network->addActivation(*threshbn1->getOutput(0), ActivationType::kRELU);
assert(threshrelu1);
Weights deconvwts10{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup2 = network->addDeconvolutionNd(*threshrelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts10, emptywts);
threshup2->setStrideNd(DimsHW{ 2, 2 });
threshup2->setNbGroups(64);
IConvolutionLayer* thresh3 = network->addConvolutionNd(*threshup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.thresh.6.1.weight"], weightMap["head.thresh.6.1.bias"]);
assert(thresh3);
thresh3->setStrideNd(DimsHW{ 1, 1 });
thresh3->setPaddingNd(DimsHW{ 1, 1 });
IActivationLayer* thresh4 = network->addActivation(*thresh3->getOutput(0), ActivationType::kSIGMOID);
assert(thresh4);
ITensor* inputTensors2[] = { binarize4->getOutput(0), thresh4->getOutput(0) };
IConcatenationLayer* head_out = network->addConcatenation(inputTensors2, 2);
// y = F.interpolate(y, size=(H, W))
head_out->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*head_out->getOutput(0));
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMIN, Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kOPT, Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMAX, Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
config->addOptimizationProfile(profile);
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
//ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int h_scale, int w_scale) {
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
context.setBindingDimensions(inputIndex, Dims4(1, 3, h_scale, w_scale));
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], 3 * h_scale * w_scale * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], 2 * h_scale * w_scale * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, 3 * h_scale * w_scale * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], h_scale * w_scale * 2 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
bool get_mini_boxes(cv::RotatedRect& rotated_rect, cv::Point2f rect[],
int min_size)
{
cv::Point2f temp_rect[4];
rotated_rect.points(temp_rect);
for (int i = 0; i < 4; i++) {
for (int j = i + 1; j < 4; j++) {
if (temp_rect[i].x > temp_rect[j].x) {
cv::Point2f temp;
temp = temp_rect[i];
temp_rect[i] = temp_rect[j];
temp_rect[j] = temp;
}
}
}
int index0 = 0;
int index1 = 1;
int index2 = 2;
int index3 = 3;
if (temp_rect[1].y > temp_rect[0].y) {
index0 = 0;
index3 = 1;
} else {
index0 = 1;
index3 = 0;
}
if (temp_rect[3].y > temp_rect[2].y) {
index1 = 2;
index2 = 3;
} else {
index1 = 3;
index2 = 2;
}
rect[0] = temp_rect[index0]; // Left top coordinate
rect[1] = temp_rect[index1]; // Left bottom coordinate
rect[2] = temp_rect[index2]; // Right bottom coordinate
rect[3] = temp_rect[index3]; // Right top coordinate
if (rotated_rect.size.width < min_size ||
rotated_rect.size.height < min_size) {
return false;
} else {
return true;
}
}
float get_box_score(float* map, cv::Point2f rect[], int width, int height,
float threshold)
{
int xmin = width - 1;
int ymin = height - 1;
int xmax = 0;
int ymax = 0;
for (int j = 0; j < 4; j++) {
if (rect[j].x < xmin) {
xmin = rect[j].x;
}
if (rect[j].y < ymin) {
ymin = rect[j].y;
}
if (rect[j].x > xmax) {
xmax = rect[j].x;
}
if (rect[j].y > ymax) {
ymax = rect[j].y;
}
}
float sum = 0;
int num = 0;
for (int i = ymin; i <= ymax; i++) {
for (int j = xmin; j <= xmax; j++) {
if (map[i * width + j] > threshold) {
sum = sum + map[i * width + j];
num++;
}
}
}
return sum / num;
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{ nullptr };
size_t size{ 0 };
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{ nullptr };
APIToModel(1, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("DBNet.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
}
else if (argc == 3 && std::string(argv[1]) == "-d") {
std::ifstream file("DBNet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
}
else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./debnet -s // serialize model to plan file" << std::endl;
std::cerr << "./debnet -d ../samples // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
std::vector<std::string> file_names;
if (read_files_in_dir(argv[2], file_names) < 0) {
std::cout << "read_files_in_dir failed." << std::endl;
return -1;
}
// icdar2015.yaml Hyperparameter
std::vector<float> mean_value{ 0.406, 0.456, 0.485 }; // BGR
std::vector<float> std_value{ 0.225, 0.224, 0.229 };
int fcount = 0;
for (auto f : file_names) {
fcount++;
std::cout << fcount << " " << f << std::endl;
cv::Mat pr_img = cv::imread(std::string(argv[2]) + "/" + f);
cv::Mat src_img = pr_img.clone();
if (pr_img.empty()) continue;
float scale = paddimg(pr_img, SHORT_INPUT); // resize the image
std::cout << "letterbox shape: " << pr_img.cols << ", " << pr_img.rows << std::endl;
if (pr_img.cols < MIN_INPUT_SIZE || pr_img.rows < MIN_INPUT_SIZE) continue;
float* data = new float[3 * pr_img.rows * pr_img.cols];
auto start = std::chrono::system_clock::now();
int i = 0;
for (int row = 0; row < pr_img.rows; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;
for (int col = 0; col < pr_img.cols; ++col) {
data[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
data[i + pr_img.rows * pr_img.cols] = (uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
data[i + 2 * pr_img.rows * pr_img.cols] = (uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
uc_pixel += 3;
++i;
}
}
auto end = std::chrono::system_clock::now();
std::cout << "pre time:"<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
float* prob = new float[pr_img.rows *pr_img.cols * 2];
// Run inference
start = std::chrono::system_clock::now();
doInference(*context, data, prob, pr_img.rows, pr_img.cols);
end = std::chrono::system_clock::now();
std::cout << "detect time:"<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
// prob shape is 2*640*640, get the first one
cv::Mat map = cv::Mat::zeros(cv::Size(pr_img.cols, pr_img.rows), CV_8UC1);
for (int h = 0; h < pr_img.rows; ++h) {
uchar *ptr = map.ptr(h);
for (int w = 0; w < pr_img.cols; ++w) {
ptr[w] = (prob[h * pr_img.cols + w] > 0.3) ? 255 : 0;
}
}
// Extracting minimum circumscribed rectangle
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(map, contours, hierarcy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
std::vector<cv::Rect> boundRect(contours.size());
std::vector<cv::RotatedRect> box(contours.size());
cv::Point2f rect[4];
cv::Point2f order_rect[4];
for (int i = 0; i < contours.size(); i++) {
cv::RotatedRect rotated_rect = cv::minAreaRect(cv::Mat(contours[i]));
if (!get_mini_boxes(rotated_rect, rect, BOX_MINI_SIZE)) {
std::cout << "box too small" << std::endl;
continue;
}
// drop low score boxes
float score = get_box_score(prob, rect, pr_img.cols, pr_img.rows,
SCORE_THRESHOLD);
if (score < BOX_THRESHOLD) {
std::cout << "score too low = " << score << ", threshold = " << BOX_THRESHOLD << std::endl;
continue;
}
// Scaling the predict boxes depend on EXPANDRATIO
cv::RotatedRect expandbox = expandBox(rect, EXPANDRATIO);
expandbox.points(rect);
if (!get_mini_boxes(expandbox, rect, BOX_MINI_SIZE + 2)) {
continue;
}
// Restore the coordinates to the original image
for (int k = 0; k < 4; k++) {
order_rect[k] = rect[k];
order_rect[k].x = int(order_rect[k].x / pr_img.cols * src_img.cols);
order_rect[k].y = int(order_rect[k].y / pr_img.rows * src_img.rows);
}
cv::rectangle(src_img, cv::Point(order_rect[0].x,order_rect[0].y), cv::Point(order_rect[2].x,order_rect[2].y), cv::Scalar(0, 0, 255), 2, 8);
//std::cout << "After LT = " << order_rect[0] << ", After RD = " << order_rect[2] << std::endl;
}
cv::imwrite("_" + f, src_img);
std::cout << "write image done." << std::endl;
//cv::waitKey(0);
delete prob;
delete data;
}
return 0;
}

503
dbnet/logging.h Normal file
View File

@ -0,0 +1,503 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

94
dbnet/utils.h Normal file
View File

@ -0,0 +1,94 @@
#ifndef __TRT_UTILS_H_
#define __TRT_UTILS_H_
#include <iostream>
#include <vector>
#include <algorithm>
#include <cudnn.h>
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr) \
{ \
cudaError_t error_code = callstr; \
if (error_code != cudaSuccess) { \
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
assert(0); \
} \
}
#endif
namespace Tn
{
class Profiler : public nvinfer1::IProfiler
{
public:
void printLayerTimes(int itrationsTimes)
{
float totalTime = 0;
for (size_t i = 0; i < mProfile.size(); i++)
{
printf("%-40.40s %4.3fms\n", mProfile[i].first.c_str(), mProfile[i].second / itrationsTimes);
totalTime += mProfile[i].second;
}
printf("Time over all layers: %4.3f\n", totalTime / itrationsTimes);
}
private:
typedef std::pair<std::string, float> Record;
std::vector<Record> mProfile;
virtual void reportLayerTime(const char* layerName, float ms)
{
auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; });
if (record == mProfile.end())
mProfile.push_back(std::make_pair(layerName, ms));
else
record->second += ms;
}
};
//Logger for TensorRT info/warning/errors
class Logger : public nvinfer1::ILogger
{
public:
Logger(): Logger(Severity::kWARNING) {}
Logger(Severity severity): reportableSeverity(severity) {}
void log(Severity severity, const char* msg) override
{
// suppress messages with severity enum value greater than the reportable
if (severity > reportableSeverity) return;
switch (severity)
{
case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break;
case Severity::kERROR: std::cerr << "ERROR: "; break;
case Severity::kWARNING: std::cerr << "WARNING: "; break;
case Severity::kINFO: std::cerr << "INFO: "; break;
default: std::cerr << "UNKNOWN: "; break;
}
std::cerr << msg << std::endl;
}
Severity reportableSeverity{Severity::kWARNING};
};
template<typename T>
void write(char*& buffer, const T& val)
{
*reinterpret_cast<T*>(buffer) = val;
buffer += sizeof(T);
}
template<typename T>
void read(const char*& buffer, T& val)
{
val = *reinterpret_cast<const T*>(buffer);
buffer += sizeof(T);
}
}
#endif

36
densenet/CMakeLists.txt Normal file
View File

@ -0,0 +1,36 @@
cmake_minimum_required(VERSION 3.10)
# set the project name
project(densenet)
add_definitions(-std=c++11)
# get main project dir to include common files
get_filename_component(MAIN_DIR ../ ABSOLUTE)
# When enabled the static version of the
# CUDA runtime library will be used in CUDA_LIBRARIES
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
# specify the C++ standard
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED True)
set(CMAKE_BUILD_TYPE Debug)
# include
# include and link cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# include and link tensorrt
include_directories(/usr/include/x86_64-linux-gnu)
link_directories(/usr/lib/x86_64-linux-gnu)
# add the executable
add_executable(densenet ${PROJECT_SOURCE_DIR}/densenet121.cpp)
target_link_libraries(densenet nvinfer)
target_link_libraries(densenet cudart)
add_definitions(-O2 -pthread)

39
densenet/README.md Normal file
View File

@ -0,0 +1,39 @@
# Densenet121
The Pytorch implementation is [makaveli10/densenet](https://github.com/makaveli10/torchtrtz/tree/main/densenet). Model from torchvision.
The tensorrt implemenation is taken from [makaveli10/cpptensorrtz](https://github.com/makaveli10/cpptensorrtz/).
## How to Run
1. generate densenet121.wts from pytorch
```
git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/makaveli10/torchtrtz.git
// go to torchtrtz/densenet
// Enter these two commands to create densenet121.wts
python models.py
python gen_trtwts.py
```
2. build densenet and run
```
// put densenet121.wts into tensorrtx/densenet
// go to tensorrtx/densenet
mkdir build
cd build
cmake ..
make
sudo ./densenet -s // serialize model to file i.e. 'densenet.engine'
sudo ./densenet -d // deserialize model and run inference
```
3. Verify output from [torch impl](https://github.com/makaveli10/torchtrtz/blob/main/densenet/README.md)
TensorRT output[:5]:
```
[-0.587389, -0.329202, -1.83404, -1.89935, -0.928404]
```

404
densenet/densenet121.cpp Normal file
View File

@ -0,0 +1,404 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <cmath>
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 224;
static const int INPUT_W = 224;
static const int OUTPUT_SIZE = 1000;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
using namespace nvinfer1;
static Logger gLogger;
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
float *gamma = (float*)weightMap[lname + ".weight"].values;
float *beta = (float*)weightMap[lname + ".bias"].values;
float *mean = (float*)weightMap[lname + ".running_mean"].values;
float *var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
std::cout << "len " << len << std::endl;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IConvolutionLayer* addDenseLayer(INetworkDefinition* network, ITensor* input, std::map<std::string, Weights>& weightMap, std::string lname, float eps)
{
// add Batchnorm
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *input, lname + ".norm1", eps);
// add relu
IActivationLayer* relu1 = network -> addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// add conv
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network -> addConvolutionNd(*relu1->getOutput(0), 128, DimsHW{1, 1}, weightMap[lname + ".conv1.weight"], emptywts);
assert(conv1);
conv1 -> setStrideNd(DimsHW{1, 1});
// add Batchnorm
IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv1 -> getOutput(0), lname + ".norm2", eps);
// add relu
IActivationLayer* relu2 = network -> addActivation(*bn2->getOutput(0), ActivationType::kRELU);
assert(relu2);
// add conv
IConvolutionLayer* conv2 = network -> addConvolutionNd(*relu2->getOutput(0), 32, DimsHW{3, 3}, weightMap[lname + ".conv2.weight"], emptywts);
assert(conv2);
conv2 -> setStrideNd(DimsHW{1, 1});
conv2 -> setPaddingNd(DimsHW{1, 1});
return conv2;
}
IPoolingLayer* addTransition(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap, int outch, std::string lname, float eps)
{
// add batch norm
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap,input, lname + ".norm", eps);
// add relu activation
IActivationLayer* relu1 = network -> addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// add convolution layer
// empty weights for no bias
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network -> addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{1, 1}, weightMap[lname + ".conv.weight"], emptywts);
assert(conv1);
conv1 -> setStrideNd(DimsHW{1, 1});
// add pooling
IPoolingLayer* pool1 = network->addPoolingNd(*conv1->getOutput(0), PoolingType::kAVERAGE, DimsHW{2, 2});
assert(pool1);
pool1 -> setStrideNd(DimsHW{2, 2});
pool1 -> setPaddingNd(DimsHW{0,0});
return pool1;
}
IConcatenationLayer* addDenseBlock(INetworkDefinition* network, ITensor* input, std::map<std::string, Weights>& weightMap, int numDenseLayers, std::string lname, float eps)
{
IConvolutionLayer* c{nullptr};
IConcatenationLayer* concat{nullptr};
ITensor* inputTensors[numDenseLayers+1];
inputTensors[0] = input;
c = addDenseLayer(network, input, weightMap, lname + ".denselayer" + std::to_string(1), eps);
int i;
for(i=1; i<numDenseLayers; i++)
{
// inch += 32;
inputTensors[i] = c -> getOutput(0);
concat = network -> addConcatenation(inputTensors, i+1);
assert(concat);
c = addDenseLayer(network, concat->getOutput(0), weightMap, lname + ".denselayer" + std::to_string(i+1), eps);
}
inputTensors[numDenseLayers] = c -> getOutput(0);
concat = network -> addConcatenation(inputTensors, numDenseLayers+1);
assert(concat);
return concat;
}
/**
* Uses the TensorRT API to create the network engine.
**/
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)
{
// Initialize NetworkDefinition
INetworkDefinition* network = builder -> createNetworkV2(0U);
auto data = network -> addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../densenet121.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
auto conv0 = network -> addConvolutionNd(*data, 64, DimsHW{7, 7}, weightMap["features.conv0.weight"], emptywts);
assert(conv0);
conv0 -> setStrideNd(DimsHW{2, 2});
conv0 -> setPaddingNd(DimsHW{3, 3});
auto norm0 = addBatchNorm2d(network, weightMap, *conv0 -> getOutput(0), "features.norm0", 1e-5);
auto relu0 = network -> addActivation(*norm0 -> getOutput(0), ActivationType::kRELU);
assert(relu0);
auto pool0 = network -> addPoolingNd(*relu0 -> getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
assert(pool0);
pool0 -> setStrideNd(DimsHW{2, 2});
pool0 -> setPaddingNd(DimsHW{1, 1});
auto dense1 = addDenseBlock(network, pool0 -> getOutput(0), weightMap, 6, "features.denseblock1", 1e-5);
auto transition1 = addTransition(network, *dense1 -> getOutput(0), weightMap, 128, "features.transition1", 1e-5);
auto dense2 = addDenseBlock(network, transition1 -> getOutput(0), weightMap, 12, "features.denseblock2", 1e-5);
auto transition2 = addTransition(network, *dense2 -> getOutput(0), weightMap, 256, "features.transition2", 1e-5);
auto dense3 = addDenseBlock(network, transition2 -> getOutput(0), weightMap, 24, "features.denseblock3", 1e-5);
auto transition3 = addTransition(network, *dense3 -> getOutput(0), weightMap, 512, "features.transition3", 1e-5);
auto dense4 = addDenseBlock(network, transition3 -> getOutput(0), weightMap, 16, "features.denseblock4", 1e-5);
auto bn5 = addBatchNorm2d(network, weightMap, *dense4 -> getOutput(0), "features.norm5", 1e-5);
auto relu5 = network -> addActivation(*bn5 -> getOutput(0), ActivationType::kRELU);
// adaptive average pool => pytorch (F.adaptive_avg_pool2d(input, (1, 1)))
auto pool5 = network -> addPoolingNd(*relu5 -> getOutput(0), PoolingType::kAVERAGE, DimsHW{7,7});
auto fc1 = network -> addFullyConnected(*pool5 -> getOutput(0), 1000, weightMap["classifier.weight"], weightMap["classifier.bias"]);
assert(fc1);
// set ouput blob name
fc1 -> getOutput(0) -> setName(OUTPUT_BLOB_NAME);
std::cout << "set name out" << std::endl;
// mark the output
network -> markOutput(*fc1 -> getOutput(0));
// set batchsize and workspace size
builder -> setMaxBatchSize(maxBatchSize);
config -> setMaxWorkspaceSize(1 << 28); // 256 MiB
// build engine
ICudaEngine* engine = builder -> buildEngineWithConfig(*network, *config);
std::cout << "build out" << std::endl;
// destroy
network -> destroy();
// fere host mem
for(auto& mem: weightMap)
{
free((void*)(mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
{
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
config->destroy();
}
/**
* Performs inference on the given input and
* writes the output from device to host memory.
**/
void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv)
{
if (argc != 2) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./densenet -s // serialize model to plan file" << std::endl;
std::cerr << "./densenet -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// create a model using the API directly and serialize it to a stream
char *trtModelStream{nullptr};
size_t size{0};
if (std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(1, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("densenet.engine", std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 1;
} else if (std::string(argv[1]) == "-d") {
std::ifstream file("densenet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
return -1;
}
// Subtract mean from image
static float data[3 * INPUT_H * INPUT_W];
for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
data[i] = 1.0;
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Run inference
static float prob[OUTPUT_SIZE];
for (int i = 0; i < 100; i++) {
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, 1);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
// Print histogram of the output distribution
std::cout << "\nOutput:\n\n";
for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
{
std::cout << prob[i] << ", ";
if (i % 10 == 0) std::cout << i / 10 << std::endl;
}
std::cout << std::endl;
return 0;
}

293
densenet/densenet121.py Normal file
View File

@ -0,0 +1,293 @@
import os
import sys
import struct
import argparse
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
BATCH_SIZE = 1
INPUT_H = 224
INPUT_W = 224
OUTPUT_SIZE = 1000
INPUT_BLOB_NAME = "data"
OUTPUT_BLOB_NAME = "prob"
EPS = 1e-5
WEIGHT_PATH = "./densenet121.wts"
ENGINE_PATH = "./densenet121.engine"
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
def load_weights(file):
print(f"Loading weights: {file}")
assert os.path.exists(file), 'Unable to load weight file.'
weight_map = {}
with open(file, "r") as f:
lines = [line.strip() for line in f]
count = int(lines[0])
assert count == len(lines) - 1
for i in range(1, count + 1):
splits = lines[i].split(" ")
name = splits[0]
cur_count = int(splits[1])
assert cur_count + 2 == len(splits)
values = []
for j in range(2, len(splits)):
# hex string to bytes to float
values.append(struct.unpack(">f", bytes.fromhex(splits[j])))
weight_map[name] = np.array(values, dtype=np.float32)
return weight_map
def add_batch_norm_2d(network, weight_map, input, layer_name):
gamma = weight_map[layer_name + ".weight"]
beta = weight_map[layer_name + ".bias"]
mean = weight_map[layer_name + ".running_mean"]
var = weight_map[layer_name + ".running_var"]
var = np.sqrt(var + EPS)
scale = gamma / var
shift = -mean / var * gamma + beta
return network.add_scale(input=input,
mode=trt.ScaleMode.CHANNEL,
shift=shift,
scale=scale)
def add_dense_layer(network, input, weight_map, lname):
bn1 = add_batch_norm_2d(network, weight_map, input, lname + ".norm1")
relu1 = network.add_activation(bn1.get_output(0), type=trt.ActivationType.RELU)
assert relu1
conv1 = network.add_convolution(input=relu1.get_output(0),
num_output_maps=128,
kernel_shape=(1, 1),
kernel=weight_map[lname + ".conv1.weight"],
bias=trt.Weights())
assert conv1
conv1.stride = (1, 1)
bn2 = add_batch_norm_2d(network, weight_map, conv1.get_output(0), lname + ".norm2")
relu2 = network.add_activation(bn2.get_output(0), type=trt.ActivationType.RELU)
assert relu2
conv2 = network.add_convolution(input=relu2.get_output(0),
num_output_maps=32,
kernel_shape=(3, 3),
kernel=weight_map[lname + ".conv2.weight"],
bias=trt.Weights())
assert conv2
conv2.stride = (1, 1)
conv2.padding = (1, 1)
return conv2
def add_transition(network, input, weight_map, outch, lname):
bn1 = add_batch_norm_2d(network, weight_map, input, lname + ".norm")
relu1 = network.add_activation(bn1.get_output(0), type=trt.ActivationType.RELU)
assert relu1
conv1 = network.add_convolution(input=relu1.get_output(0),
num_output_maps=outch,
kernel_shape=(1, 1),
kernel=weight_map[lname + ".conv.weight"],
bias=trt.Weights())
assert conv1
conv1.stride = (1, 1)
pool1 = network.add_pooling(input=conv1.get_output(0),
type=trt.PoolingType.AVERAGE,
window_size=trt.DimsHW(2, 2))
assert pool1
pool1.stride_nd = (2, 2)
pool1.padding_nd = (0, 0)
return pool1
def add_dense_block(network, input, weight_map, num_dense_layers, lname):
input_tensors = [None for _ in range(num_dense_layers+1)]
input_tensors[0] = input
c = add_dense_layer(network, input, weight_map, lname + ".denselayer" + str(1))
for i in range(1, num_dense_layers):
input_tensors[i] = c.get_output(0)
concat = network.add_concatenation(input_tensors[:i+1])
assert concat
c = add_dense_layer(network, concat.get_output(0), weight_map, lname + ".denselayer" + str(i+1))
input_tensors[num_dense_layers] = c.get_output(0)
concat = network.add_concatenation(input_tensors)
assert concat
return concat
def create_engine(max_batch_size, builder, config, dt):
weight_map = load_weights(WEIGHT_PATH)
network = builder.create_network()
data = network.add_input(INPUT_BLOB_NAME, dt, (3, INPUT_H, INPUT_W))
assert data
conv0 = network.add_convolution(input=data,
num_output_maps=64,
kernel_shape=(7, 7),
kernel=weight_map["features.conv0.weight"],
bias=trt.Weights())
assert conv0
conv0.stride = (2, 2)
conv0.padding = (3, 3)
bn0 = add_batch_norm_2d(network, weight_map, conv0.get_output(0), "features.norm0")
relu0 = network.add_activation(bn0.get_output(0), type=trt.ActivationType.RELU)
assert relu0
pool0 = network.add_pooling(input=relu0.get_output(0),
type=trt.PoolingType.MAX,
window_size=trt.DimsHW(3, 3))
assert pool0
pool0.stride_nd = (2, 2)
pool0.padding_nd = (1, 1)
dense1 = add_dense_block(network, pool0.get_output(0), weight_map, 6, "features.denseblock1")
transition1 = add_transition(network, dense1.get_output(0), weight_map, 128, "features.transition1")
dense2 = add_dense_block(network, transition1.get_output(0), weight_map, 12, "features.denseblock2")
transition2 = add_transition(network, dense2.get_output(0), weight_map, 256, "features.transition2")
dense3 = add_dense_block(network, transition2.get_output(0), weight_map, 24, "features.denseblock3")
transition3 = add_transition(network, dense3.get_output(0), weight_map, 512, "features.transition3")
dense4 = add_dense_block(network, transition3.get_output(0), weight_map, 16, "features.denseblock4")
bn5 = add_batch_norm_2d(network, weight_map, dense4.get_output(0), "features.norm5")
relu5 = network.add_activation(bn5.get_output(0), type=trt.ActivationType.RELU)
pool5 = network.add_pooling(relu5.get_output(0), type=trt.PoolingType.AVERAGE, window_size=trt.DimsHW(7, 7))
fc1 = network.add_fully_connected(input=pool5.get_output(0),
num_outputs=OUTPUT_SIZE,
kernel=weight_map["classifier.weight"],
bias=weight_map["classifier.bias"])
assert fc1
fc1.get_output(0).name = OUTPUT_BLOB_NAME
network.mark_output(fc1.get_output(0))
# Build Engine
builder.max_batch_size = max_batch_size
builder.max_workspace_size = 1 << 20
engine = builder.build_engine(network, config)
del network
del weight_map
return engine
def API_to_model(max_batch_size):
builder = trt.Builder(TRT_LOGGER)
config = builder.create_builder_config()
engine = create_engine(max_batch_size, builder, config, trt.float32)
assert engine
with open(ENGINE_PATH, "wb") as f:
f.write(engine.serialize())
del engine
del builder
del config
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-s", action='store_true')
parser.add_argument("-d", action='store_true')
args = parser.parse_args()
if not (args.s ^ args.d):
print(
"arguments not right!\n"
"python densenet121.py -s # serialize model to plan file\n"
"python densenet121.py -d # deserialize plan file and run inference"
)
sys.exit()
if args.s:
API_to_model(BATCH_SIZE)
else:
runtime = trt.Runtime(TRT_LOGGER)
assert runtime
with open(ENGINE_PATH, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
assert engine
context = engine.create_execution_context()
assert context
data = np.ones((BATCH_SIZE * 3 * INPUT_H * INPUT_W), dtype=np.float32)
inputs, outputs, bindings, stream = allocate_buffers(engine)
inputs[0].host = data
trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
print(f'Output: \n{trt_outputs[0][:10]}\n{trt_outputs[0][-10:]}')

507
densenet/logging.h Normal file
View File

@ -0,0 +1,507 @@
/*
* Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
, mPrefix(other.mPrefix)
, mShouldLog(other.mShouldLog)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
{
ss << " ";
}
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
//! ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

33
detr/CMakeLists.txt Normal file
View File

@ -0,0 +1,33 @@
cmake_minimum_required(VERSION 2.6)
project(detr)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/data/app/TensorRT-8.4.3.1/include)
link_directories(/data/app/TensorRT-8.4.3.1/lib)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(detr ${PROJECT_SOURCE_DIR}/detr.cpp)
target_link_libraries(detr nvinfer)
target_link_libraries(detr cudart)
target_link_libraries(detr ${OpenCV_LIBS})
add_definitions(-O2 -pthread)

89
detr/README.md Normal file
View File

@ -0,0 +1,89 @@
# DETR
The Pytorch implementation is [facebookresearch/detr](https://github.com/facebookresearch/detr).
For details see [End-to-End Object Detection with Transformers](https://ai.facebook.com/research/publications/end-to-end-object-detection-with-transformers).
## Test Environment
- GTX2080Ti / Ubuntu16.04 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2
- GTX2080Ti / win10 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2 / VS2017
## How to Run
1. generate .wts from pytorch with .pth
```
// git clone https://github.com/facebookresearch/detr.git
// go to facebookresearch/detr
// download https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth
// download https://raw.githubusercontent.com/freedenS/TestImage/main/demo.jpg
// copy tensorrtx/detr/gen_wts.py and demo.jpg into facebookresearch/detr
python gen_wts.py
// a file 'detr.wts' will be generated.
```
2. build tensorrtx/detr and run
```
// put detr.wts into tensorrtx/detr
// go to tensorrtx/detr
// update parameters in detr.cpp if your model is trained on custom dataset.The parameters are corresponding to config in detr.
mkdir build
cd build
cmake ..
make
sudo ./detr -s [.wts] // serialize model to plan file
sudo ./detr -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed
// For example
sudo ./detr -s ../detr.wts detr.engine
sudo ./detr -d detr.engine ../samples
```
3. check the images generated, as follows. _demo.jpg and so on.
## Backbone
#### R50
```
1.download pretrained model
https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth
2.export wts
set first parameter in Backbone in gen_wts.py(line 23) to resnet50
set path of pretrained model(line 87 in gen_wts.py)
3.set resnet_type in BuildResNet(line 546 in detr.cpp) to R50
```
#### R101
```
1.download pretrained model
https://dl.fbaipublicfiles.com/detr/detr-r101-2c7b67e5.pth
2.export wts
set first parameter in Backbone in gen_wts.py(line 23) to resnet101
set path of pretrained model(line 87 in gen_wts.py)
3.set resnet_type in BuildResNet(line 546 in detr.cpp) to R101
```
## NOTE
- tensorrt use fixed input size, if the size of your data is different from the engine, you need to adjust your data and the result.
- image preprocessing with c++ is a little different with python(opencv vs PIL)
## Quantization
1. quantizationType:fp32,fp16,int8. see BuildDETRModel(detr.cpp line 613) for detail.
2. the usage of int8 is same with [tensorrtx/yolov5](../yolov5/README.md).
## Latency
average cost of doInference(in detr.cpp) from second time with batch=1 under the ubuntu environment above
| | fp32 | fp16 | int8 |
| ---- | ------- | ------- | ------ |
| R50 | 19.57ms | 9.424ms | 8.38ms |
| R101 | 30.82ms | 12.4ms | 9.59ms |

326
detr/backbone.hpp Normal file
View File

@ -0,0 +1,326 @@
#pragma once
#include <map>
#include "common.hpp"
enum RESNETTYPE {
R18 = 0,
R34,
R50,
R101,
R152
};
const std::map<RESNETTYPE, std::vector<int>> num_blocks_per_stage = {
{R18, {2, 2, 2, 2}},
{R34, {3, 4, 6, 3}},
{R50, {3, 4, 6, 3}},
{R101, {3, 4, 23, 3}},
{R152, {3, 8, 36, 3}}
};
IScaleLayer* addBatchNorm2d(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
ITensor& input,
const std::string& lname,
float eps = 1e-5
) {
float *gamma = (float*)(weightMap[lname + ".weight"].values);
float *beta = (float*)(weightMap[lname + ".bias"].values);
float *mean = (float*)(weightMap[lname + ".running_mean"].values);
float *var = (float*)(weightMap[lname + ".running_var"].values);
int len = weightMap[lname + ".running_var"].count;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{ DataType::kFLOAT, scval, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
ILayer* BasicStem(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& input,
int out_channels,
int group_num = 1
) {
// conv1
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
IConvolutionLayer* conv1 = network->addConvolutionNd(
input,
out_channels,
DimsHW{ 7, 7 },
weightMap[lname + ".conv1.weight"],
emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{ 2, 2 });
conv1->setPaddingNd(DimsHW{ 3, 3 });
conv1->setNbGroups(group_num);
auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1");
assert(bn1);
auto r1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(r1);
auto max_pool2d = network->addPoolingNd(*r1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
max_pool2d->setStrideNd(DimsHW{ 2, 2 });
max_pool2d->setPaddingNd(DimsHW{ 1, 1 });
auto mp_dim = max_pool2d->getOutput(0)->getDimensions();
return max_pool2d;
}
ITensor* BasicBlock(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& input,
int in_channels,
int out_channels,
int stride = 1
) {
// conv1
IConvolutionLayer* conv1 = network->addConvolutionNd(
input,
out_channels,
DimsHW{ 3, 3 },
weightMap[lname + ".conv1.weight"],
weightMap[lname + ".conv1.bias"]);
assert(conv1);
conv1->setStrideNd(DimsHW{ stride, stride });
conv1->setPaddingNd(DimsHW{ 1, 1 });
auto r1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(r1);
// conv2
IConvolutionLayer* conv2 = network->addConvolutionNd(
*r1->getOutput(0),
out_channels, DimsHW{ 3, 3 },
weightMap[lname + ".conv2.weight"],
weightMap[lname + ".conv2.bias"]);
assert(conv2);
conv2->setStrideNd(DimsHW{ 1, 1 });
conv2->setPaddingNd(DimsHW{ 1, 1 });
// shortcut
ITensor* shortcut_value = nullptr;
if (in_channels != out_channels) {
auto shortcut = network->addConvolutionNd(
input,
out_channels,
DimsHW{ 1, 1 },
weightMap[lname + ".shortcut.weight"],
weightMap[lname + ".shortcut.bias"]);
assert(shortcut);
shortcut->setStrideNd(DimsHW{ stride, stride });
shortcut_value = shortcut->getOutput(0);
} else {
shortcut_value = &input;
}
// add
auto ew = network->addElementWise(*conv2->getOutput(0), *shortcut_value, ElementWiseOperation::kSUM);
assert(ew);
auto r3 = network->addActivation(*ew->getOutput(0), ActivationType::kRELU);
assert(r3);
return r3->getOutput(0);
}
ITensor* BottleneckBlock(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& input,
int in_channels,
int bottleneck_channels,
int out_channels,
int stride = 1,
int dilation = 1,
int group_num = 1
) {
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
// conv1
IConvolutionLayer* conv1 = network->addConvolutionNd(
input,
bottleneck_channels,
DimsHW{ 1, 1 },
weightMap[lname + ".conv1.weight"],
emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{ 1, 1 });
conv1->setNbGroups(group_num);
auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1");
assert(bn1);
auto r1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(r1);
// conv2
IConvolutionLayer* conv2 = network->addConvolutionNd(
*r1->getOutput(0),
bottleneck_channels,
DimsHW{ 3, 3 },
weightMap[lname + ".conv2.weight"],
emptywts);
assert(conv2);
conv2->setStrideNd(DimsHW{ stride, stride });
conv2->setPaddingNd(DimsHW{ 1 * dilation, 1 * dilation });
conv2->setDilationNd(DimsHW{ dilation, dilation });
conv2->setNbGroups(group_num);
auto bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2");
assert(bn2);
auto r2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
assert(r2);
// conv3
IConvolutionLayer* conv3 = network->addConvolutionNd(
*r2->getOutput(0),
out_channels,
DimsHW{ 1, 1 },
weightMap[lname + ".conv3.weight"],
emptywts);
assert(conv3);
conv3->setStrideNd(DimsHW{ 1, 1 });
conv3->setNbGroups(group_num);
auto bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3");
assert(bn3);
// shortcut
ITensor* shortcut_value = nullptr;
if (in_channels != out_channels) {
auto shortcut = network->addConvolutionNd(
input,
out_channels,
DimsHW{ 1, 1 },
weightMap[lname + ".downsample.0.weight"],
emptywts);
assert(shortcut);
shortcut->setStrideNd(DimsHW{stride, stride});
shortcut->setNbGroups(group_num);
auto shortcut_bn = addBatchNorm2d(network, weightMap, *shortcut->getOutput(0), lname + ".downsample.1");
assert(shortcut_bn);
shortcut_value = shortcut_bn->getOutput(0);
} else {
shortcut_value = &input;
}
// add
auto ew = network->addElementWise(*bn3->getOutput(0), *shortcut_value, ElementWiseOperation::kSUM);
assert(ew);
auto r3 = network->addActivation(*ew->getOutput(0), ActivationType::kRELU);
assert(r3);
return r3->getOutput(0);
}
ITensor* MakeStage(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& input,
int stage,
RESNETTYPE resnet_type,
int in_channels,
int bottleneck_channels,
int out_channels,
int first_stride = 1,
int dilation = 1
) {
ITensor* out = &input;
for (int i = 0; i < stage; i++) {
std::string layerName = lname + "." + std::to_string(i);
int stride = i == 0 ? first_stride : 1;
if (resnet_type == R18 || resnet_type == R34)
out = BasicBlock(network, weightMap, layerName, *out, in_channels, out_channels, stride);
else
out = BottleneckBlock(
network,
weightMap,
layerName,
*out,
in_channels,
bottleneck_channels,
out_channels,
stride,
dilation);
in_channels = out_channels;
}
return out;
}
ITensor* BuildResNet(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
ITensor& input,
RESNETTYPE resnet_type,
int stem_out_channels,
int bottleneck_channels,
int res2_out_channels,
int res5_dilation = 1
) {
assert(res5_dilation == 1 || res5_dilation == 2); // "res5_dilation must be 1 or 2"
if (resnet_type == R18 || resnet_type == R34) {
assert(res2_out_channels == 64); // "res2_out_channels must be 64 for R18/R34")
assert(res5_dilation == 1); // "res5_dilation must be 1 for R18/R34")
}
int out_channels = res2_out_channels;
ITensor* out = nullptr;
// stem
auto stem = BasicStem(network, weightMap, "backbone.0.body", input, stem_out_channels);
out = stem->getOutput(0);
// res
for (int i = 0; i < 4; i++) {
int dilation = (i == 3) ? res5_dilation : 1;
int first_stride = (i == 0 || (i == 3 && dilation == 2)) ? 1 : 2;
out = MakeStage(
network,
weightMap,
"backbone.0.body.layer" + std::to_string(i + 1),
*out,
num_blocks_per_stage.at(resnet_type)[i],
resnet_type,
stem_out_channels,
bottleneck_channels,
out_channels,
first_stride,
dilation);
stem_out_channels = out_channels;
bottleneck_channels *= 2;
out_channels *= 2;
}
return out;
}

117
detr/calibrator.hpp Normal file
View File

@ -0,0 +1,117 @@
#pragma once
#include "NvInfer.h"
#include <string>
#include <vector>
#include <iostream>
#include <iterator>
#include <fstream>
#include <algorithm>
#include "common.hpp"
#include "macros.h"
//! \class Int8EntropyCalibrator2
//!
//! \brief Implements Entropy calibrator 2.
//! CalibrationAlgoType is kENTROPY_CALIBRATION_2.
//!
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
public:
Int8EntropyCalibrator2(int batchsize, int input_w, int input_h,
const char* img_dir, const char* calib_table_name,
const char* input_blob_name, bool read_cache = true);
virtual ~Int8EntropyCalibrator2();
int getBatchSize() const TRT_NOEXCEPT override;
bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override;
const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override;
void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override;
private:
int batchsize_;
int input_w_;
int input_h_;
int img_idx_;
std::string img_dir_;
std::vector<std::string> img_files_;
size_t input_count_;
std::string calib_table_name_;
const char* input_blob_name_;
bool read_cache_;
void* device_input_;
std::vector<char> calib_cache_;
};
Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize,
int input_w, int input_h, const char* img_dir,
const char* calib_table_name, const char* input_blob_name,
bool read_cache)
: batchsize_(batchsize)
, input_w_(input_w)
, input_h_(input_h)
, img_idx_(0)
, img_dir_(img_dir)
, calib_table_name_(calib_table_name)
, input_blob_name_(input_blob_name)
, read_cache_(read_cache) {
input_count_ = 3 * input_w * input_h * batchsize;
CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float)));
read_files_in_dir(img_dir, img_files_);
}
Int8EntropyCalibrator2::~Int8EntropyCalibrator2() {
CUDA_CHECK(cudaFree(device_input_));
}
int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT {
return batchsize_;
}
bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT {
if (img_idx_ + batchsize_ > static_cast<int>(img_files_.size())) {
return false;
}
std::vector<float> input_imgs_(input_count_, 0);
for (int i = img_idx_; i < img_idx_ + batchsize_; i++) {
std::cout << img_files_[i] << " " << i << std::endl;
cv::Mat temp = cv::imread(img_dir_ + img_files_[i]);
if (temp.empty()) {
std::cerr << "Fatal error: image cannot open!" << std::endl;
return false;
}
preprocessImg(temp, input_w_, input_h_);
for (int c = 0; c < 3; c++) {
for (int h = 0; h < input_h_; h++) {
for (int w = 0; w < input_w_; w++) {
input_imgs_[(i-img_idx_)*input_w_*input_h_*3 +
c * input_h_ * input_w_ + h * input_w_ + w] = temp.at<cv::Vec3f>(h, w)[c];
}
}
}
}
img_idx_ += batchsize_;
CUDA_CHECK(cudaMemcpy(device_input_, input_imgs_.data(), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
assert(!strcmp(names[0], input_blob_name_));
bindings[0] = device_input_;
return true;
}
const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT {
std::cout << "reading calib cache: " << calib_table_name_ << std::endl;
calib_cache_.clear();
std::ifstream input(calib_table_name_, std::ios::binary);
input >> std::noskipws;
if (read_cache_ && input.good()) {
std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(calib_cache_));
}
length = calib_cache_.size();
return length ? calib_cache_.data() : nullptr;
}
void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT {
std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl;
std::ofstream output(calib_table_name_, std::ios::binary);
output.write(reinterpret_cast<const char*>(cache), length);
}

100
detr/common.hpp Normal file
View File

@ -0,0 +1,100 @@
#pragma once
#include <dirent.h>
#include <cuda_runtime_api.h>
#include <fstream>
#include <sstream>
#include <iostream>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include "./logging.h"
#include <NvInfer.h>
#include <opencv2/opencv.hpp>
static Logger gLogger;
using namespace nvinfer1;
void loadWeights(const std::string file, std::unordered_map<std::string, Weights>& weightMap) {
std::cout << "Loading weights: " << file << std::endl;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--) {
Weights wt{ DataType::kFLOAT, nullptr, 0 };
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
}
int CalculateSize(Dims a) {
int res = 1;
for (int i = 0; i < a.nbDims; i++) {
res *= a.d[i];
}
return res;
}
static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
// std::string cur_file_name(p_dir_name);
// cur_file_name += "/";
// cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}
void preprocessImg(cv::Mat& img, int newh, int neww) {
// convert to rgb
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
cv::resize(img, img, cv::Size(neww, newh));
img.convertTo(img, CV_32FC3);
img /= 255;
img -= cv::Scalar(0.485, 0.456, 0.406);
img /= cv::Scalar(0.229, 0.224, 0.225);
}
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr)\
{\
cudaError_t error_code = callstr;\
if (error_code != cudaSuccess) {\
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
assert(0);\
}\
}
#endif // CUDA_CHECK

818
detr/detr.cpp Normal file
View File

@ -0,0 +1,818 @@
#pragma once
#include <iostream>
#include <unordered_map>
#include "./logging.h"
#include "backbone.hpp"
#include "calibrator.hpp"
#define DEVICE 0
#define BATCH_SIZE 1
// 1 / math.sqrt(head_dim) https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/nn/functional/activation.h#623
static const float SCALING = 0.17677669529663687;
static const int INPUT_H = 800;
static const int INPUT_W = 1066;
static const int NUM_CLASS = 92; // include background
static const float SCALING_ONE = 1.0;
static const float SHIFT_ZERO = 0.0;
static const float POWER_TWO = 2.0;
static const float EPS = 0.00001;
static const int D_MODEL = 256;
static const int NHEAD = 8;
static const int DIM_FEEDFORWARD = 2048;
static const int NUM_ENCODE_LAYERS = 6;
static const int NUM_DECODE_LAYERS = 6;
static const int NUM_QUERIES = 100;
static const float SCORE_THRESH = 0.5;
const char* INPUT_NODE_NAME = "images";
const std::vector<std::string> OUTPUT_NAMES = { "scores", "boxes"};
ITensor* PositionEmbeddingSine(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
ITensor& input,
int num_pos_feats = 64,
int temperature = 10000
) {
// refer to https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py#12
// TODO: improve this implementation
auto mask_dim = input.getDimensions();
int h = mask_dim.d[1], w = mask_dim.d[2];
std::vector<std::vector<float>> y_embed(h);
for (int i = 0; i < h; i++)
y_embed[i] = std::vector<float>(w, i + 1);
std::vector<float> sub_embed(w, 0);
for (int i = 0; i < w; i++)
sub_embed[i] = i + 1;
std::vector<std::vector<float>> x_embed(h, sub_embed);
// normalize
float eps = 1e-6, scale = 2.0 * 3.1415926;
for (int i = 0; i < h; i++) {
for (int j = 0; j < w; j++) {
y_embed[i][j] = y_embed[i][j] / (h + eps) * scale;
x_embed[i][j] = x_embed[i][j] / (w + eps) * scale;
}
}
// dim_t
std::vector<float> dim_t(num_pos_feats, 0);
for (int i = 0; i < num_pos_feats; i++) {
dim_t[i] = pow(temperature, (2 * (i / 2) / static_cast<float>(num_pos_feats)));
}
// pos_x, pos_y
std::vector<std::vector<std::vector<float>>> pos_x(h,
std::vector<std::vector<float>>(w,
std::vector<float>(num_pos_feats, 0)));
std::vector<std::vector<std::vector<float>>> pos_y(h,
std::vector<std::vector<float>>(w,
std::vector<float>(num_pos_feats, 0)));
for (int i = 0; i < h; i++) {
for (int j = 0; j < w; j++) {
for (int k = 0; k < num_pos_feats; k++) {
float value_x = x_embed[i][j] / dim_t[k];
float value_y = y_embed[i][j] / dim_t[k];
if (k & 1) {
pos_x[i][j][k] = std::cos(value_x);
pos_y[i][j][k] = std::cos(value_y);
} else {
pos_x[i][j][k] = std::sin(value_x);
pos_y[i][j][k] = std::sin(value_y);
}
}
}
}
// pos
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * h * w * num_pos_feats * 2));
float *pNext = pval;
for (int i = 0; i < h; i++) {
for (int j = 0; j < w; j++) {
for (int k = 0; k < num_pos_feats; k++) {
*pNext = pos_y[i][j][k];
++pNext;
}
for (int k = 0; k < num_pos_feats; k++) {
*pNext = pos_x[i][j][k];
++pNext;
}
}
}
Weights pos_embed_weight{ DataType::kFLOAT, pval, h * w * num_pos_feats * 2 };
weightMap["pos"] = pos_embed_weight;
auto pos_embed = network->addConstant(Dims4{ h * w, num_pos_feats * 2, 1, 1 }, pos_embed_weight);
assert(pos_embed);
return pos_embed->getOutput(0);
}
ITensor* MultiHeadAttention(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& query,
ITensor& key,
ITensor& value,
int embed_dim = 256,
int num_heads = 8
) {
int tgt_len = query.getDimensions().d[0];
int head_dim = embed_dim / num_heads;
// q
auto linear_q = network->addFullyConnected(
query,
embed_dim,
weightMap[lname + ".in_proj_weight_q"],
weightMap[lname + ".in_proj_bias_q"]);
assert(linear_q);
// k
auto linear_k = network->addFullyConnected(
key,
embed_dim,
weightMap[lname + ".in_proj_weight_k"],
weightMap[lname + ".in_proj_bias_k"]);
assert(linear_k);
// v
auto linear_v = network->addFullyConnected(
value,
embed_dim,
weightMap[lname + ".in_proj_weight_v"],
weightMap[lname + ".in_proj_bias_v"]);
assert(linear_v);
auto scaling_t = network->addConstant(Dims4{ 1, 1, 1, 1 }, Weights{ DataType::kFLOAT, &SCALING, 1 });
assert(scaling_t);
auto q_scaling = network->addElementWise(
*linear_q->getOutput(0),
*scaling_t->getOutput(0),
ElementWiseOperation::kPROD);
assert(q_scaling);
auto q_shuffle = network->addShuffle(*q_scaling->getOutput(0));
assert(q_shuffle);
q_shuffle->setName((lname + ".q_shuffle").c_str());
q_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim });
q_shuffle->setSecondTranspose(Permutation{1, 0, 2});
auto k_shuffle = network->addShuffle(*linear_k->getOutput(0));
assert(k_shuffle);
k_shuffle->setName((lname + ".k_shuffle").c_str());
k_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim });
k_shuffle->setSecondTranspose(Permutation{ 1, 0, 2 });
auto v_shuffle = network->addShuffle(*linear_v->getOutput(0));
assert(v_shuffle);
v_shuffle->setName((lname + ".v_shuffle").c_str());
v_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim });
v_shuffle->setSecondTranspose(Permutation{ 1, 0, 2 });
#if NV_TENSORRT_MAJOR >= 8
auto q_product_k = network->addMatrixMultiply(*q_shuffle->getOutput(0), nvinfer1::MatrixOperation::kNONE, *k_shuffle->getOutput(0), nvinfer1::MatrixOperation::kTRANSPOSE);
#else
auto q_product_k = network->addMatrixMultiply(*q_shuffle->getOutput(0), false, *k_shuffle->getOutput(0), true);
#endif
assert(q_product_k);
// src_key_padding_mask are all false, so do nothing here
// see https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/nn/functional/activation.h#826-#839
auto softmax = network->addSoftMax(*q_product_k->getOutput(0));
assert(softmax);
softmax->setAxes(4);
#if NV_TENSORRT_MAJOR >= 8
auto attn_product_v = network->addMatrixMultiply(*softmax->getOutput(0), nvinfer1::MatrixOperation::kNONE, *v_shuffle->getOutput(0), nvinfer1::MatrixOperation::kNONE);
#else
auto attn_product_v = network->addMatrixMultiply(*softmax->getOutput(0), false, *v_shuffle->getOutput(0), false);
#endif
assert(attn_product_v);
auto attn_shuffle = network->addShuffle(*attn_product_v->getOutput(0));
assert(attn_shuffle);
attn_shuffle->setName((lname + ".attn_shuffle").c_str());
attn_shuffle->setFirstTranspose(Permutation{ 1, 0, 2 });
attn_shuffle->setReshapeDimensions(Dims4{ tgt_len, -1, 1, 1 });
auto linear_attn = network->addFullyConnected(
*attn_shuffle->getOutput(0),
embed_dim,
weightMap[lname + ".out_proj.weight"],
weightMap[lname + ".out_proj.bias"]);
assert(linear_attn);
return linear_attn->getOutput(0);
}
ITensor* LayerNorm(
INetworkDefinition *network,
ITensor& input,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
int d_model = 256
) {
// TODO: maybe a better implementation https://github.com/NVIDIA/TensorRT/blob/master/plugin/common/common.cuh#212
auto mean = network->addReduce(input, ReduceOperation::kAVG, 2, true);
assert(mean);
auto sub_mean = network->addElementWise(input, *mean->getOutput(0), ElementWiseOperation::kSUB);
assert(sub_mean);
// implement pow2 with scale
Weights scale{ DataType::kFLOAT, &SCALING_ONE, 1 };
Weights shift{ DataType::kFLOAT, &SHIFT_ZERO, 1 };
Weights power{ DataType::kFLOAT, &POWER_TWO, 1 };
auto pow2 = network->addScaleNd(*sub_mean->getOutput(0), ScaleMode::kUNIFORM, shift, scale, power, 0);
assert(pow2);
auto pow_mean = network->addReduce(*pow2->getOutput(0), ReduceOperation::kAVG, 2, true);
assert(pow_mean);
auto eps = network->addConstant(Dims4{ 1, 1, 1, 1 }, Weights{ DataType::kFLOAT, &EPS, 1 });
assert(eps);
auto add_eps = network->addElementWise(*pow_mean->getOutput(0), *eps->getOutput(0), ElementWiseOperation::kSUM);
assert(add_eps);
auto sqrt = network->addUnary(*add_eps->getOutput(0), UnaryOperation::kSQRT);
assert(sqrt);
auto div = network->addElementWise(*sub_mean->getOutput(0), *sqrt->getOutput(0), ElementWiseOperation::kDIV);
assert(div);
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * d_model));
for (int i = 0; i < d_model; i++) {
pval[i] = 1.0;
}
Weights norm1_power{ DataType::kFLOAT, pval, d_model };
weightMap[lname + ".power"] = norm1_power;
auto affine = network->addScaleNd(
*div->getOutput(0),
ScaleMode::kCHANNEL,
weightMap[lname + ".bias"],
weightMap[lname + ".weight"],
norm1_power,
1);
assert(affine);
return affine->getOutput(0);
}
ITensor* TransformerEncoderLayer(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& src,
ITensor& pos,
int d_model = 256,
int nhead = 8,
int dim_feedforward = 2048
) {
auto pos_embed = network->addElementWise(src, pos, ElementWiseOperation::kSUM);
assert(pos_embed);
ITensor* src2 = MultiHeadAttention(
network,
weightMap,
lname + ".self_attn",
*pos_embed->getOutput(0),
*pos_embed->getOutput(0),
src,
d_model,
nhead);
auto shortcut1 = network->addElementWise(src, *src2, ElementWiseOperation::kSUM);
assert(shortcut1);
ITensor* norm1 = LayerNorm(network, *shortcut1->getOutput(0), weightMap, lname + ".norm1");
auto linear1 = network->addFullyConnected(
*norm1,
dim_feedforward,
weightMap[lname + ".linear1.weight"],
weightMap[lname + ".linear1.bias"]);
assert(linear1);
auto relu = network->addActivation(*linear1->getOutput(0), ActivationType::kRELU);
assert(relu);
auto linear2 = network->addFullyConnected(
*relu->getOutput(0),
d_model,
weightMap[lname + ".linear2.weight"],
weightMap[lname + ".linear2.bias"]);
assert(linear2);
auto shortcut2 = network->addElementWise(*norm1, *linear2->getOutput(0), ElementWiseOperation::kSUM);
assert(shortcut2);
ITensor* norm2 = LayerNorm(network, *shortcut2->getOutput(0), weightMap, lname + ".norm2");
return norm2;
}
ITensor* TransformerEncoder(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& src,
ITensor& pos,
int num_layers = 6
) {
ITensor* out = &src;
for (int i = 0; i < num_layers; i++) {
std::string layer_name = lname + ".layers." + std::to_string(i);
out = TransformerEncoderLayer(network, weightMap, layer_name, *out, pos);
}
return out;
}
ITensor* TransformerDecoderLayer(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& tgt,
ITensor& memory,
ITensor& pos,
ITensor& query_pos,
int d_model = 256,
int nhead = 8,
int dim_feedforward = 2048
) {
auto pos_embed = network->addElementWise(tgt, query_pos, ElementWiseOperation::kSUM);
assert(pos_embed);
ITensor* tgt2 = MultiHeadAttention(
network,
weightMap,
lname + ".self_attn",
*pos_embed->getOutput(0),
*pos_embed->getOutput(0),
tgt);
auto shortcut1 = network->addElementWise(tgt, *tgt2, ElementWiseOperation::kSUM);
assert(shortcut1);
ITensor* norm1 = LayerNorm(network, *shortcut1->getOutput(0), weightMap, lname + ".norm1");
auto query_embed = network->addElementWise(*norm1, query_pos, ElementWiseOperation::kSUM);
assert(query_embed);
auto key_embed = network->addElementWise(memory, pos, ElementWiseOperation::kSUM);
assert(key_embed);
ITensor* mha2 = MultiHeadAttention(
network,
weightMap,
lname + ".multihead_attn",
*query_embed->getOutput(0),
*key_embed->getOutput(0),
memory);
auto shortcut2 = network->addElementWise(*norm1, *mha2, ElementWiseOperation::kSUM);
assert(shortcut2);
ITensor* norm2 = LayerNorm(network, *shortcut2->getOutput(0), weightMap, lname + ".norm2");
auto linear1 = network->addFullyConnected(
*norm2,
dim_feedforward,
weightMap[lname + ".linear1.weight"],
weightMap[lname + ".linear1.bias"]);
assert(linear1);
auto relu = network->addActivation(*linear1->getOutput(0), ActivationType::kRELU);
assert(relu);
auto linear2 = network->addFullyConnected(
*relu->getOutput(0),
d_model,
weightMap[lname + ".linear2.weight"],
weightMap[lname + ".linear2.bias"]);
assert(linear2);
auto shortcut3 = network->addElementWise(*norm2, *linear2->getOutput(0), ElementWiseOperation::kSUM);
assert(shortcut3);
ITensor* norm3 = LayerNorm(network, *shortcut3->getOutput(0), weightMap, lname + ".norm3");
return norm3;
}
ITensor* TransformerDecoder(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& tgt,
ITensor& memory,
ITensor& pos,
ITensor& query_pos,
int num_layers = 6,
int d_model = 256,
int nhead = 8,
int dim_feedforward = 2048
) {
ITensor* out = &tgt;
for (int i = 0; i < num_layers; i++) {
std::string layer_name = lname + ".layers." + std::to_string(i);
out = TransformerDecoderLayer(
network,
weightMap,
layer_name,
*out,
memory,
pos,
query_pos,
d_model,
nhead,
dim_feedforward);
}
ITensor* norm = LayerNorm(network, *out, weightMap, lname + ".norm", d_model);
return norm;
}
ITensor* Transformer(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& src,
ITensor& pos_embed,
int num_queries = 100,
int num_encoder_layers = 6,
int num_decoder_layers = 6,
int d_model = 256,
int nhead = 8,
int dim_feedforward = 2048
) {
auto memory = TransformerEncoder(network, weightMap, lname + ".encoder", src, pos_embed, num_encoder_layers);
// construct tgt
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * num_queries * d_model));
for (int i = 0; i < num_queries * d_model; i++) {
pval[i] = 0.0;
}
Weights tgt_weight{ DataType::kFLOAT, pval, num_queries * d_model };
weightMap[lname + ".tgt_weight"] = tgt_weight;
auto tgt = network->addConstant(Dims4{ num_queries, d_model, 1, 1 }, tgt_weight);
assert(tgt);
// construct query_pos
auto query_pos = network->addConstant(Dims4{ num_queries, d_model, 1, 1 }, weightMap["query_embed.weight"]);
assert(query_pos);
auto out = TransformerDecoder(
network,
weightMap,
lname + ".decoder",
*tgt->getOutput(0),
*memory, pos_embed,
*query_pos->getOutput(0),
num_decoder_layers,
d_model,
nhead,
dim_feedforward);
return out;
}
ITensor* MLP(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
const std::string& lname,
ITensor& src,
int num_layers = 3,
int hidden_dim = 256,
int output_dim = 4
) {
ITensor* out = &src;
for (int i = 0; i < num_layers; i++) {
std::string layer_name = lname + "." + std::to_string(i);
if (i != num_layers - 1) {
auto fc = network->addFullyConnected(
*out,
hidden_dim,
weightMap[layer_name + ".weight"],
weightMap[layer_name + ".bias"]);
assert(fc);
auto relu = network->addActivation(*fc->getOutput(0), ActivationType::kRELU);
assert(relu);
out = relu->getOutput(0);
} else {
auto fc = network->addFullyConnected(
*out,
output_dim,
weightMap[layer_name + ".weight"],
weightMap[layer_name + ".bias"]);
assert(fc);
out = fc->getOutput(0);
}
}
return out;
}
std::vector<ITensor*> Predict(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& weightMap,
ITensor& src
) {
auto class_embed = network->addFullyConnected(
src,
NUM_CLASS,
weightMap["class_embed.weight"],
weightMap["class_embed.bias"]);
assert(class_embed);
auto class_softmax = network->addSoftMax(*class_embed->getOutput(0));
assert(class_softmax);
class_softmax->setAxes(2);
ITensor* bbox = MLP(network, weightMap, "bbox_embed.layers", src);
auto bbox_sig = network->addActivation(*bbox, ActivationType::kSIGMOID);
assert(bbox_sig);
std::vector<ITensor*> output = { class_softmax->getOutput(0), bbox_sig->getOutput(0) };
return output;
}
ICudaEngine* createEngine_r50detr(
unsigned int maxBatchSize,
const std::string& wtsfile,
IBuilder* builder,
IBuilderConfig* config,
DataType dt,
const std::string& modelType = "fp16"
) {
/*
description: after fuse bn
*/
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_NODE_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
// preprocess
std::unordered_map<std::string, Weights> weightMap;
loadWeights(wtsfile, weightMap);
// backbone
auto features = BuildResNet(network, weightMap, *data, R50, 64, 64, 256);
ITensor* pos_embed = PositionEmbeddingSine(network, weightMap, *features, 128);
auto input_proj = network->addConvolutionNd(
*features,
D_MODEL,
DimsHW{ 1, 1 },
weightMap["input_proj.weight"],
weightMap["input_proj.bias"]);
assert(input_proj);
input_proj->setStrideNd(DimsHW{ 1, 1 });
auto flatten = network->addShuffle(*input_proj->getOutput(0));
assert(flatten);
flatten->setReshapeDimensions(Dims4{ input_proj->getOutput(0)->getDimensions().d[0], -1, 1, 1 });
flatten->setSecondTranspose(Permutation{ 1, 0, 2, 3 });
auto out1 = Transformer(
network,
weightMap,
"transformer",
*flatten->getOutput(0),
*pos_embed,
NUM_QUERIES,
NUM_ENCODE_LAYERS,
NUM_DECODE_LAYERS,
D_MODEL,
NHEAD,
DIM_FEEDFORWARD);
std::vector<ITensor*> results = Predict(network, weightMap, *out1);
// build output
for (int i = 0; i < results.size(); i++) {
network->markOutput(*results[i]);
results[i]->setName(OUTPUT_NAMES[i].c_str());
}
// build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(1ULL << 30);
if (modelType == "fp32") {
} else if (modelType == "fp16") {
config->setFlag(BuilderFlag::kFP16);
} else if (modelType == "int8") {
std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
assert(builder->platformHasFastInt8());
config->setFlag(BuilderFlag::kINT8);
Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(BATCH_SIZE, INPUT_W, INPUT_H, "./coco_calib/",
"int8calib.table", INPUT_NODE_NAME);
config->setInt8Calibrator(calibrator);
} else {
throw("does not support model type");
}
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// destroy network
network->destroy();
// Release host memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
}
void BuildDETRModel(unsigned int maxBatchSize, IHostMemory** modelStream,
const std::string& wtsfile, std::string modelType = "fp32") {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine_r50detr(maxBatchSize,
wtsfile, builder, config, DataType::kFLOAT, modelType);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, cudaStream_t& stream, std::vector<void*>& buffers,
std::vector<float>& input, std::vector<float*>& output) {
CUDA_CHECK(cudaMemcpyAsync(buffers[0], input.data(), input.size() * sizeof(float),
cudaMemcpyHostToDevice, stream));
context.enqueue(BATCH_SIZE, buffers.data(), stream, nullptr);
CUDA_CHECK(cudaMemcpyAsync(output[0], buffers[1], BATCH_SIZE * NUM_QUERIES * NUM_CLASS * sizeof(float),
cudaMemcpyDeviceToHost, stream));
CUDA_CHECK(cudaMemcpyAsync(output[1], buffers[2], BATCH_SIZE * NUM_QUERIES * 4 * sizeof(float),
cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
}
bool parse_args(int argc, char** argv, std::string& wtsFile, std::string& engineFile, std::string& imgDir) {
if (argc < 4) return false;
if (std::string(argv[1]) == "-s") {
wtsFile = std::string(argv[2]);
engineFile = std::string(argv[3]);
} else if (std::string(argv[1]) == "-d") {
engineFile = std::string(argv[2]);
imgDir = std::string(argv[3]);
} else {
return false;
}
return true;
}
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
std::string wtsFile = "";
std::string engineFile = "";
std::string imgDir;
if (!parse_args(argc, argv, wtsFile, engineFile, imgDir)) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./detr -s [.wts] [.engine] // serialize model to plan file" << std::endl;
std::cerr << "./detr -d [.engine] ../samples // deserialize plan file and run inference" << std::endl;
return -1;
}
if (!wtsFile.empty()) {
IHostMemory* modelStream{ nullptr };
BuildDETRModel(BATCH_SIZE, &modelStream, wtsFile, "fp32");
assert(modelStream != nullptr);
std::ofstream p(engineFile, std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
}
// deserialize the .engine and run inference
std::ifstream file(engineFile, std::ios::binary);
if (!file.good()) {
std::cerr << "read " << engineFile << " error!" << std::endl;
return -1;
}
std::string trtModelStream;
size_t modelSize{ 0 };
file.seekg(0, file.end);
modelSize = file.tellg();
file.seekg(0, file.beg);
trtModelStream.resize(modelSize);
assert(!trtModelStream.empty());
file.read(const_cast<char*>(trtModelStream.c_str()), modelSize);
file.close();
// build engine
std::cout << "build engine" << std::endl;
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream.c_str(), modelSize);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
runtime->destroy();
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
// prepare input file
std::vector<std::string> fileList;
if (read_files_in_dir(imgDir.c_str(), fileList) < 0) {
std::cerr << "read_files_in_dir failed." << std::endl;
return -1;
}
// calculate input size
int input_size = CalculateSize(context->getBindingDimensions(0));
// prepare input data
std::vector<float> data(BATCH_SIZE * input_size, 0);
void *data_d, *scores_d, *boxes_d;
CUDA_CHECK(cudaMalloc(&data_d, BATCH_SIZE * input_size * sizeof(float)));
CUDA_CHECK(cudaMalloc(&scores_d, BATCH_SIZE * NUM_QUERIES * NUM_CLASS * sizeof(float)));
CUDA_CHECK(cudaMalloc(&boxes_d, BATCH_SIZE * NUM_QUERIES * 4 * sizeof(float)));
std::vector<float> scores_h(BATCH_SIZE * NUM_QUERIES * NUM_CLASS);
std::vector<float> boxes_h(BATCH_SIZE * NUM_QUERIES * 4);
std::vector<void*> buffers = { data_d, scores_d, boxes_d };
std::vector<float*> outputs = {scores_h.data(), boxes_h.data()};
int fcount = 0;
int fileLen = fileList.size();
for (int f = 0; f < fileLen; f++) {
fcount++;
if (fcount < BATCH_SIZE && f + 1 != fileLen) continue;
for (int b = 0; b < fcount; b++) {
cv::Mat img = cv::imread(imgDir + "/" + fileList[f - fcount + 1 + b]);
if (img.empty()) continue;
preprocessImg(img, INPUT_H, INPUT_W);
assert(img.cols * img.rows * 3 == input_size);
for (int c = 0; c < 3; c++) {
for (int h = 0; h < img.rows; h++) {
for (int w = 0; w < img.cols; w++) {
data[b * input_size +
c * img.rows * img.cols + h * img.cols + w] = img.at<cv::Vec3f>(h, w)[c];
}
}
}
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, stream, buffers, data, outputs);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
for (int b = 0; b < fcount; b++) {
cv::Mat img = cv::imread(imgDir + "/" + fileList[f - fcount + 1 + b]);
for (int i = 0; i < scores_h.size(); i += NUM_CLASS) {
int label = -1;
float score = -1;
for (int j = i; j < i + NUM_CLASS; j++) {
if (score < scores_h[j]) {
label = j;
score = scores_h[j];
}
}
if (score > SCORE_THRESH && (label % NUM_CLASS != NUM_CLASS - 1)) {
int ind = label / NUM_CLASS;
label = label % NUM_CLASS;
float cx = boxes_h[ind * 4];
float cy = boxes_h[ind * 4 + 1];
float w = boxes_h[ind * 4 + 2];
float h = boxes_h[ind * 4 + 3];
float x1 = (cx - w / 2.0) * img.cols;
float y1 = (cy - h / 2.0) * img.rows;
float x2 = (cx + w / 2.0) * img.cols;
float y2 = (cy + h / 2.0) * img.rows;
cv::Rect r(x1, y1, x2 - x1, y2 - y1);
cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
cv::putText(img, std::to_string(label), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2,
cv::Scalar(0xFF, 0xFF, 0xFF), 2);
}
}
cv::imwrite("_" + fileList[f - fcount + 1 + b], img);
}
fcount = 0;
}
cudaStreamDestroy(stream);
CUDA_CHECK(cudaFree(data_d));
CUDA_CHECK(cudaFree(scores_d));
CUDA_CHECK(cudaFree(boxes_d));
context->destroy();
engine->destroy();
return 0;
}

123
detr/gen_wts.py Normal file
View File

@ -0,0 +1,123 @@
import cv2
import torch
from models.transformer import Transformer
from models.position_encoding import PositionEmbeddingSine
from models.backbone import Backbone, Joiner
from models.detr import DETR
import torchvision.transforms as T
from PIL import Image
import struct
def box_cxcywh_to_xyxy(x):
x_c, y_c, w, h = x.unbind(-1)
b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
(x_c + 0.5 * w), (y_c + 0.5 * h)]
return torch.stack(b, dim=-1)
def build_backbone():
N_steps = 256 // 2
position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
train_backbone = True
return_interm_layers = False
backbone = Backbone('resnet50', train_backbone, return_interm_layers, False)
model = Joiner(backbone, position_embedding)
model.num_channels = backbone.num_channels
return model
def gen_wts(model, filename):
f = open(filename + '.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys()) + 72))
for k, v in model.state_dict().items():
if 'in_proj' in k:
dim = int(v.size(0) / 3)
q_weight = v[:dim].reshape(-1).cpu().numpy()
k_weight = v[dim:2*dim].reshape(-1).cpu().numpy()
v_weight = v[2*dim:].reshape(-1).cpu().numpy()
f.write('{} {} '.format(k + '_q', len(q_weight)))
for vv in q_weight:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')
f.write('{} {} '.format(k + '_k', len(k_weight)))
for vv in k_weight:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')
f.write('{} {} '.format(k + '_v', len(v_weight)))
for vv in v_weight:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')
else:
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')
f.close()
def main():
num_classes = 91
device = torch.device('cuda')
backbone = build_backbone()
transformer = Transformer(
d_model=256,
dropout=0.1,
nhead=8,
dim_feedforward=2048,
num_encoder_layers=6,
num_decoder_layers=6,
normalize_before=False,
return_intermediate_dec=True,
)
model = DETR(
backbone,
transformer,
num_classes=num_classes,
num_queries=100,
aux_loss=True,
)
checkpoint = torch.load('./detr-r50-e632da11.pth')
model.load_state_dict(checkpoint['model'])
model.to(device)
model.eval()
gen_wts(model, "detr")
# test
# with torch.no_grad():
# transform = T.Compose([T.Resize(800), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# im = Image.open('./image/demo.jpg')
# img = transform(im).unsqueeze(0)
# img = img.to(device)
# res = model(img)
# logits = res['pred_logits']
# pred_boxes = res['pred_boxes']
# out_prob = logits.softmax(-1)[0, :, :-1]
# keep = out_prob.max(-1).values > 0.5
# label = out_prob[keep].argmax(dim=1)
# out_bbox = pred_boxes[0, keep]
# out_bbox = out_bbox.to(torch.device('cpu'))
# out_bbox = box_cxcywh_to_xyxy(out_bbox)
# out_bbox = out_bbox * torch.tensor([640, 480, 640, 480])
# image = cv2.imread('./image/demo.jpg')
# for ob in out_bbox:
# x0 = int(ob[0].item())
# y0 = int(ob[1].item())
# x1 = int(ob[2].item())
# y1 = int(ob[3].item())
# cv2.rectangle(image, (x0, y0), (x1, y1), (0,0,255), 1)
# cv2.imwrite('res.jpg', image)
if __name__ == '__main__':
main()

450
detr/logging.h Normal file
View File

@ -0,0 +1,450 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
#include "macros.h"
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf {
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog) {}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput) {}
~LogStreamConsumerBuffer() {
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr()) {
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync() {
putOutput();
return 0;
}
void putOutput() {
if (mShouldLog) {
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog) {
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase {
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog) {}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity) {}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity) {}
void setReportableSeverity(Severity reportableSeverity) {
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger {
public:
explicit Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity) {}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult {
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger() {
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) TRT_NOEXCEPT override {
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity) {
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom {
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline) {}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline) {
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) {
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom) {
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result) {
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass) {
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const {
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result) {
switch (result) {
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result) {
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv) {
std::stringstream ss;
for (int i = 0; i < argc; i++) {
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace {
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

29
detr/macros.h Normal file
View File

@ -0,0 +1,29 @@
#ifndef __MACROS_H
#define __MACROS_H
#include "NvInfer.h"
#ifdef API_EXPORTS
#if defined(_MSC_VER)
#define API __declspec(dllexport)
#else
#define API __attribute__((visibility("default")))
#endif
#else
#if defined(_MSC_VER)
#define API __declspec(dllimport)a
#else
#define API
#endif
#endif // API_EXPORTS
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif
#endif // __MACROS_H

6
docker/.env Normal file
View File

@ -0,0 +1,6 @@
COMPOSE_PROJECT_NAME=tensorrtx
HOME=$HOME
EUID=$(id -u)
## (optional) a local mount point path
DATA_DIR=""

74
docker/README.md Normal file
View File

@ -0,0 +1,74 @@
# Tutorials
## Introduction
This folder contains the docker and docker-compose file to build the development environment without pain.
## Prerequisites
* OS: Linux or WSL2
* docker
* nvidia-container-toolkit
* (Optional but **recommended**) docker-compose
## Usage
1. (With docker-compose) configure the `.env` file, change `DATA_DIR` to your mount point, such as your code or data folder, etc, comment the `volumes` in docker compose file if not necessariy needed
2. Build image:
```bash
docker compose -f docker-compose.yml build
```
3. Run a container at background:
```bash
docker compose -f docker-compose.yml up -d
```
4. Attach to this container with your IDE and have fun!
## HowTos
### How to build and run with docker?
``` bash
docker build -f docker/x86_64.dockerfile -v .
docker run -it --gpus all --privileged --net=host --ipc=host -v /bin/bash
```
### How to build image with other TensorRT version?
Change the `TAG` on top of the `.dockerfile`. Note: all images are officially owned by NVIDIA NGC, which requires a registration before pulling. For this repo, the mainly used `TAG` would be:
| Container Image | Container OS | Driver | CUDA | TensorRT | Torch | Recommended |
| :----: | :----: | :----: | :----: | :----: | :----: | :----: |
| 20.12-py3 | Ubuntu 20.04 | 455 | 11.2 | 7.2.2 | 1.8.0 | ❌ |
| 24.01-py3 | Ubuntu 22.04 | 545 | 12.3 | 8.6.1 | 2.2.0 | ✅ |
| 24.04-py3 | Ubuntu 22.04 | 545 | 12.4 | 8.6.3 | 2.3.0 | ✅ |
| 24.09-py3 | Ubuntu 22.04 | 560 | 12.6 | 10.4.0 | 2.5.0 | ✅ |
For more detail of the support matrix, please check [HERE](https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html)
### How to customize opencv?
If prebuilt package from apt cannot meet your requirements, please refer to the demo code in `.dockerfile` to build opencv from source.
### How to solve image build fail issues?
For *443 timeout* or any similar network issues, a proxy may required. To make your host proxy work for building env of docker, please change the `build` node inside docker-compose file like this:
```YAML
build:
dockerfile: x86_64.dockerfile
args:
HTTP_PROXY: ${PROXY}
HTTPS_PROXY: ${PROXY}
ALL_PROXY: ${PROXY}
http_proxy: ${PROXY}
https_proxy: ${PROXY}
all_proxy: ${PROXY}
```
then add `PROXY="http://xxx:xxx"` in `.env` file
## Note
The older version support, like TensorRT version **< 8**, may be deprecated in the future.

46
docker/tensorrtx-docker-compose.yml Normal file
View File

@ -0,0 +1,46 @@
services:
tensorrt:
image: tensortx:1.0.0
container_name: tensortx
environment:
- NVIDIA_VISIBLE_DEVICES=all
build:
dockerfile: x86_64.dockerfile
cap_add:
- CAP_SYS_ADMIN
security_opt:
- seccomp:unconfined
privileged: true
stdin_open: true
tty: true
shm_size: '8gb'
ulimits:
memlock:
soft: -1
hard: -1
devices:
- /dev:/dev:rw
volumes:
#### user ####
- ${HOME}:/workspace/localhome:rw
#### custom ####
- mount:/mnt:rw
deploy:
restart_policy:
condition: on-failure
max_attempts: 1
delay: 5s
resources:
reservations:
devices:
- driver: nvidia
capabilities: [gpu]
count: all
volumes:
mount:
driver: local
driver_opts:
type: none
o: bind
device: ${DATA_DIR}

39
docker/x86_64.dockerfile Normal file
View File

@ -0,0 +1,39 @@
ARG TAG=24.01-py3
FROM nvcr.io/nvidia/tensorrt:${TAG} AS tensorrtx
ENV DEBIAN_FRONTEND noninteractive
# basic tools
RUN apt update && apt-get install -y --fix-missing --no-install-recommends \
sudo wget curl git ca-certificates ninja-build tzdata pkg-config \
gdb libglib2.0-dev libmount-dev \
&& rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir yapf isort cmake-format pre-commit
## override older cmake
RUN find /usr/local/share -type d -name "cmake-*" -exec rm -rf {} + \
&& curl -fsSL "https://github.com/Kitware/CMake/releases/download/v3.29.0/cmake-3.29.0-linux-x86_64.sh" \
-o cmake.sh && bash cmake.sh --skip-license --exclude-subdir --prefix=/usr/local && rm cmake.sh
RUN apt update && apt-get install -y \
libopencv-dev \
&& rm -rf /var/lib/apt/lists/*
## a template to build opencv and opencv_contrib from source
# RUN git clone -b 4.x https://github.com/opencv/opencv_contrib.git \
# && git clone -b 4.x https://github.com/opencv/opencv.git opencv \
# && cmake -S opencv -B opencv/build -G Ninja \
# -DBUILD_LIST=core,calib3d,imgproc,imgcodecs,highgui \
# -DOPENCV_EXTRA_MODULES_PATH="/workspace/opencv_contrib/modules" \
# -DCMAKE_BUILD_TYPE=RELEASE \
# -DCMAKE_INSTALL_PREFIX=/usr/local \
# -DENABLE_FAST_MATH=ON \
# -DOPENCV_GENERATE_PKGCONFIG=ON \
# -DBUILD_opencv_python2=OFF \
# -DBUILD_opencv_python3=OFF \
# -DBUILD_JAVA=OFF \
# -DBUILD_DOCS=OFF \
# -DBUILD_PERF_TESTS=OFF \
# -DBUILD_TESTS=OFF \
# && ninja -C opencv/build install

37
efficient_ad/CMakeLists.txt Normal file
View File

@ -0,0 +1,37 @@
cmake_minimum_required(VERSION 3.12)
project(EfficientAD-M)
add_definitions(-w)
add_definitions(-D API_EXPORTS)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE "Debug")
set(CMAKE_CUDA_ARCHITECTURES 61 75 86 89)
set(THREADS_PREFER_PTHREAD_FLAG ON)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /od")
### nvcc
set(CMAKE_CUDA_COMPILER "D:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8/bin/nvcc.exe")
enable_language(CUDA)
### cuda
include_directories("D:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8/include")
link_directories("D:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8/lib/x64")
### tensorrt
set(TRT_DIR "D:/Program Files/NVIDIA GPU Computing Toolkit/TensorRT-8.5.3.1/")
include_directories(${TRT_DIR}/include)
link_directories(${TRT_DIR}/lib)
### opencv
set(OpenCV_DIR "E:/OpenCV/OpenCV_4.6.0/opencv/build")
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
### dirent
include_directories("E:/SDK/dirent-1.24/include")
include_directories(${PROJECT_SOURCE_DIR}/src/)
file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu)
add_executable(efficientAD_det "./efficientAD_det.cpp" ${SRCS})
target_link_libraries(efficientAD_det nvinfer
cudart
nvinfer_plugin
${OpenCV_LIBS}
)

50
efficient_ad/README.md Normal file
View File

@ -0,0 +1,50 @@
# EfficientAd
EfficientAd: Accurate Visual Anomaly Detection at Millisecond-Level Latencies.
The Pytorch implementation is [openvinotoolkit/anomalib](https://github.com/openvinotoolkit/anomalib).
<p align="center">
<img src="https://github.com/wang-xinyu/tensorrtx/assets/15235574/061c90a7-fe59-48e0-a8d0-6bddc4296cf1">
</p>
# Test Environment
GTX3080 / Windows10 22H2 / cuda11.8 / cudnn8.9.7 / TensorRT8.5.3 / OpenCV4.6
# How to Run
1. training to generate weight files (`efficientAD_[category].pt`)
```
// Please refer to Anomalib's tutorial for details:
// https://github.com/openvinotoolkit/anomalib?tab=readme-ov-file#-training
```
2. generate `.wts` from pytorch with `.pt`
```
cd ./datas/models/
// copy your `.pt` file to the current directory.
python gen_wts.py
// a file `efficientAD_[category].wts` will be generated.
```
3. build and run
```
mkdir build
cd build
cmake ..
make
sudo ./EfficientAD-M -s [.wts] // serialize model to plan file
sudo ./EfficientAD-M -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed
```
# Latency
average cost of doInference(in `efficientad_detect.cpp`) from second time with batch=1 under the windows environment above
| | FP32 |
| :-----------: | :--: |
| EfficientAD-M | 12ms |

20
efficient_ad/datas/models/gen_wts.py Normal file
View File

@ -0,0 +1,20 @@
import torch
import struct
import sys
# Initialize
pt_file = sys.argv[1]
device = torch.device('cuda')
# Load model
model = torch.load(pt_file, map_location=torch.device('cpu'))['model'].float() # load to FP32
model.to(device).eval()
with open(pt_file.split('.')[0] + '.wts', 'w') as f:
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')

256
efficient_ad/efficientAD_det.cpp Normal file
View File

@ -0,0 +1,256 @@
#include <cuda_runtime.h>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <opencv2/opencv.hpp>
#include "config.h"
#include "cuda_utils.h"
#include "logging.h"
#include "model.h"
#include "postprocess.h"
#include "utils.h"
using namespace nvinfer1;
static Logger gLogger;
// const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1;
const static int kInputSize = 3 * 256 * 256;
const static int kOutputSize = 1 * 256 * 256;
bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw,
std::string& img_dir) {
if (argc != 4)
return false;
if (std::string(argv[1]) == "-s") {
wts = std::string(argv[2]);
engine = std::string(argv[3]);
} else if (std::string(argv[1]) == "-d") {
engine = std::string(argv[2]);
img_dir = std::string(argv[3]);
} else {
return false;
}
return true;
}
void prepare_infer_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer,
float** cpu_output_buffer) {
// assert(engine->getNbIOTensors() == 2);
assert(engine->getNbBindings() == 2);
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine->getBindingIndex(kInputTensorName);
const int outputIndex = engine->getBindingIndex(kOutputTensorName);
// nvinfer1::Dims outputDims = engine->getBindingDimensions(outputIndex);
assert(inputIndex == 0);
assert(outputIndex == 1);
// Create GPU in/output buffers on device
CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float)));
CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * 1 * kOutputSize * sizeof(float))); // 3 or 1 ??
// Create CPU output buffers on host
*cpu_output_buffer = new float[kBatchSize * kOutputSize];
}
void preprocessImg(cv::Mat& img, int newh, int neww) {
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
cv::resize(img, img, cv::Size(neww, newh));
img.convertTo(img, CV_32FC3);
// ImageNet normalize
img /= 255.0f;
img -= cv::Scalar(0.485, 0.456, 0.406);
img /= cv::Scalar(0.229, 0.224, 0.225);
}
void infer(IExecutionContext& context, cudaStream_t& stream, std::vector<void*>& gpu_buffers,
std::vector<float>& cpu_input_data, std::vector<float>& cpu_output_data, int batchsize) {
// copy input data from host (CPU) to device (GPU)
CUDA_CHECK(cudaMemcpyAsync(gpu_buffers[0], cpu_input_data.data(), cpu_input_data.size() * sizeof(float),
cudaMemcpyHostToDevice, stream));
// execute inference using context provided by engine
context.enqueue(batchsize, gpu_buffers.data(), stream, nullptr);
// copy output back from device (GPU) to host (CPU)
CUDA_CHECK(cudaMemcpyAsync(cpu_output_data.data(), gpu_buffers[1], batchsize * kOutputSize * sizeof(float),
cudaMemcpyDeviceToHost, stream));
// synchronize the stream to prevent issues (block CUDA and wait for CUDA operations to be completed)
cudaStreamSynchronize(stream);
}
void serialize_engine(unsigned int max_batchsize, float& gd, float& gw, std::string& wts_name,
std::string& engine_name) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = nullptr;
engine = build_efficientAD_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
assert(engine != nullptr);
// Serialize the engine
IHostMemory* serialized_engine = engine->serialize();
assert(serialized_engine != nullptr);
// Save engine to file
std::ofstream p(engine_name, std::ios::binary);
if (!p) {
std::cerr << "Could not open plan output file" << std::endl;
assert(false);
}
p.write(reinterpret_cast<const char*>(serialized_engine->data()), serialized_engine->size());
// Close everything down
engine->destroy();
config->destroy();
serialized_engine->destroy();
builder->destroy();
}
void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine,
IExecutionContext** context) {
std::ifstream file(engine_name, std::ios::binary);
if (!file.good()) {
std::cerr << "read " << engine_name << " error!" << std::endl;
assert(false);
}
size_t size = 0;
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
char* serialized_engine = new char[size];
assert(serialized_engine);
file.read(serialized_engine, size);
file.close();
*runtime = createInferRuntime(gLogger);
assert(*runtime);
*engine = (*runtime)->deserializeCudaEngine(serialized_engine, size);
assert(*engine != nullptr);
*context = (*engine)->createExecutionContext();
assert(*context);
delete[] serialized_engine;
}
int main(int argc, char** argv) {
cudaSetDevice(kGpuId);
std::string wts_name = "";
std::string engine_name = "";
float gd = 1.0f, gw = 1.0f;
std::string img_dir;
if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./efficientad_det -s [.wts] [.engine] // serialize model to plan file" << std::endl;
std::cerr
<< "./efficientad_det -d [.engine] [../../datas/images/...] // deserialize plan file and run inference"
<< std::endl;
return -1;
}
// Create a model using the API directly and serialize it to a file
if (!wts_name.empty()) {
serialize_engine(kBatchSize, gd, gw, wts_name, engine_name);
return 0;
}
// Deserialize the engine from file
IRuntime* runtime = nullptr;
ICudaEngine* engine = nullptr;
IExecutionContext* context = nullptr;
deserialize_engine(engine_name, &runtime, &engine, &context);
// create CUDA stream for simultaneous CUDA operations
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
// prepare cpu and gpu buffers
void *gpu_input_buffer, *gpu_output_buffer;
CUDA_CHECK(cudaMalloc(&gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float)));
CUDA_CHECK(cudaMalloc(&gpu_output_buffer, kBatchSize * 1 * kOutputSize * sizeof(float))); // 3 or 1 ??
std::vector<void*> gpu_buffers = {gpu_input_buffer, gpu_output_buffer};
std::vector<float> cpu_input_data(kBatchSize * kInputSize, 0);
std::vector<float> cpu_output_data(kBatchSize * kOutputSize, 0);
// read images from directory
std::vector<std::string> file_names;
if (read_files_in_dir(img_dir.c_str(), file_names) < 0) {
std::cerr << "read_files_in_dir failed." << std::endl;
return -1;
}
std::vector<cv::Mat> originImg_batch;
for (size_t i = 0; i < file_names.size(); i += kBatchSize) {
// get a batch of images
std::vector<cv::Mat> img_batch;
std::vector<std::string> img_name_batch;
for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) {
cv::Mat img = cv::imread(img_dir + "/" + file_names[j]);
originImg_batch.push_back(img.clone());
preprocessImg(img, kInputW, kInputH);
assert(img.cols * img.rows * 3 == 3 * 256 * 256);
for (int c = 0; c < 3; c++) {
for (int h = 0; h < img.rows; h++) {
for (int w = 0; w < img.cols; w++) {
cpu_input_data[c * img.rows * img.cols + h * img.cols + w] = img.at<cv::Vec3f>(h, w)[c];
}
}
}
img_batch.push_back(img);
img_name_batch.push_back(file_names[j]);
}
// Run inference
auto start = std::chrono::system_clock::now();
// infer(*context, stream, (void**)gpu_buffers, cpu_input_data, cpu_output_buffer, kBatchSize);
infer(*context, stream, gpu_buffers, cpu_input_data, cpu_output_data,
kBatchSize); // change to save into vec `cpu_output_data`
auto end = std::chrono::system_clock::now();
std::cout << "inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count()
<< "ms" << std::endl;
// postProcess
cv::Mat img_1(256, 256, CV_8UC1);
for (int row = 0; row < 256; row++) {
for (int col = 0; col < 256; col++) {
float value = cpu_output_data[row * 256 + col];
if (value < 0) // clip(0,1)
value = 0;
else if (value > 1)
value = 1;
img_1.at<uchar>(row, col) = static_cast<uchar>(value * 255);
}
}
cv::Mat HeatMap, colorMap;
// genHeatMap(img_batch[0], img_1, HeatMap);
cv::applyColorMap(img_1, colorMap, cv::COLORMAP_JET);
cv::resize(originImg_batch[i], originImg_batch[i], cv::Size(256, 256));
cv::cvtColor(originImg_batch[i], originImg_batch[i], cv::COLOR_RGB2BGR);
cv::addWeighted(originImg_batch[i], 0.5, colorMap, 0.5, 0, HeatMap);
// Save images
for (size_t j = 0; j < img_batch.size(); j++) {
cv::imwrite("_output" + img_name_batch[j], img_1);
cv::imwrite("_heatmap" + img_name_batch[j], HeatMap);
}
}
// Release stream and buffers
cudaStreamDestroy(stream);
CUDA_CHECK(cudaFree(gpu_buffers[0]));
CUDA_CHECK(cudaFree(gpu_buffers[1]));
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
return 0;
}

31
efficient_ad/src/config.h Normal file
View File

@ -0,0 +1,31 @@
#pragma once
/* --------------------------------------------------------
* These configs are related to tensorrt model, if these are changed,
* please re-compile and re-serialize the tensorrt model.
* --------------------------------------------------------*/
// For INT8, you need prepare the calibration dataset, please refer to
#define USE_FP32 // set USE_INT8 or USE_FP16 or USE_FP32
// These are used to define input/output tensor names,
// you can set them to whatever you want.
const static char* kInputTensorName = "data";
const static char* kOutputTensorName = "prob";
constexpr static int kBatchSize = 1;
// input width and height must by divisible by 32
constexpr static int kInputH = 256;
constexpr static int kInputW = 256;
/* --------------------------------------------------------
* These configs are NOT related to tensorrt model, if these are changed,
* please re-compile, but no need to re-serialize the tensorrt model.
* --------------------------------------------------------*/
// default GPU_id
const static int kGpuId = 0;
// If your image size is larger than 4096 * 3112, please increase this value
const static int kMaxInputImageSize = 4096 * 3112;

17
efficient_ad/src/cuda_utils.h Normal file
View File

@ -0,0 +1,17 @@
#ifndef TRTX_CUDA_UTILS_H_
#define TRTX_CUDA_UTILS_H_
#include <cuda_runtime_api.h>
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr) \
{ \
cudaError_t error_code = callstr; \
if (error_code != cudaSuccess) { \
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
assert(0); \
} \
}
#endif // CUDA_CHECK
#endif // TRTX_CUDA_UTILS_H_

456
efficient_ad/src/logging.h Normal file
View File

@ -0,0 +1,456 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
#include "NvInferRuntimeCommon.h"
#include "macros.h"
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf {
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) : mOutput(other.mOutput) {}
~LogStreamConsumerBuffer() {
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr()) {
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync() {
putOutput();
return 0;
}
void putOutput() {
if (mShouldLog) {
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase {
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog) {}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(severity <= reportableSeverity),
mSeverity(severity) {}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(other.mShouldLog),
mSeverity(other.mSeverity) {}
void setReportableSeverity(Severity reportableSeverity) {
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger {
public:
Logger(Severity severity = Severity::kWARNING) : mReportableSeverity(severity) {}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult {
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger() { return *this; }
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) TRT_NOEXCEPT override {
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity) { mReportableSeverity = severity; }
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom {
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started), mName(name), mCmdline(cmdline) {}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline) {
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) {
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom) {
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result) {
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass) {
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const { return mReportableSeverity; }
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result) {
switch (result) {
case TestResult::kRUNNING:
return "RUNNING";
case TestResult::kPASSED:
return "PASSED";
case TestResult::kFAILED:
return "FAILED";
case TestResult::kWAIVED:
return "WAIVED";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result) {
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv) {
std::stringstream ss;
for (int i = 0; i < argc; i++) {
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace {
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

29
efficient_ad/src/macros.h Normal file
View File

@ -0,0 +1,29 @@
#ifndef __MACROS_H
#define __MACROS_H
#include <NvInfer.h>
#ifdef API_EXPORTS
#if defined(_MSC_VER)
#define API __declspec(dllexport)
#else
#define API __attribute__((visibility("default")))
#endif
#else
#if defined(_MSC_VER)
#define API __declspec(dllimport)
#else
#define API
#endif
#endif // API_EXPORTS
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif
#endif // __MACROS_H

436
efficient_ad/src/model.cpp Normal file
View File

@ -0,0 +1,436 @@
#include "model.h"
#include <cassert>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iostream>
#include <map>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>
#include "config.h"
using namespace nvinfer1;
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
static std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--) {
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
void printNetworkLayers(INetworkDefinition* network) {
int numLayers = network->getNbLayers();
// std::cout << "currently num of layers: " << numLayers << std::endl;
auto dataTypeToString = [](DataType type) {
switch (type) {
case DataType::kFLOAT:
return "kFLOAT";
case DataType::kHALF:
return "kHALF";
case DataType::kINT8:
return "kINT8";
case DataType::kINT32:
return "kINT32";
case DataType::kBOOL:
return "kBOOL";
default:
return "Unknown";
}
};
for (int i = 0; i < numLayers; ++i) {
ILayer* layer = network->getLayer(i);
std::cout << "--- Layer" << i << " = " << layer->getName() << std::endl;
std::cout << "input & output tensor type: " << dataTypeToString(layer->getInput(0)->getType()) << "\t"
<< dataTypeToString(layer->getOutput(0)->getType()) << std::endl;
// input
int inTensorNum = layer->getNbInputs();
for (int j = 0; j < inTensorNum; ++j) {
// std::cout << layer->getInput(j)->getDimensions().nbDims;
Dims dims_in = layer->getInput(j)->getDimensions();
std::cout << "input shape[" << j << "]: (";
for (int k = 0; k < dims_in.nbDims; ++k) {
std::cout << dims_in.d[k];
if (k < dims_in.nbDims - 1) {
std::cout << ", ";
}
}
std::cout << ")\t";
}
std::cout << std::endl;
// output
int outTensorNum = layer->getNbOutputs();
for (int j = 0; j < outTensorNum; ++j) {
// std::cout << layer->getOutput(j)->getName();
Dims dims_out = layer->getOutput(j)->getDimensions();
std::cout << "output shape: (";
for (int k = 0; k < dims_out.nbDims; ++k) {
std::cout << dims_out.d[k];
if (k < dims_out.nbDims - 1) {
std::cout << ", ";
}
}
std::cout << ")";
}
std::cout << "\n" << std::endl;
}
}
static IScaleLayer* NormalizeInput(INetworkDefinition* network, ITensor& input) {
float meanValues[3] = {-0.485f, -0.456f, -0.406f};
float stdValues[3] = {1.0f / 0.229f, 1.0f / 0.224f, 1.0f / 0.225f};
Weights meanWeights{DataType::kFLOAT, meanValues, 3};
Weights stdWeights{DataType::kFLOAT, stdValues, 3};
IScaleLayer* NormaLayer = network->addScale(input, ScaleMode::kCHANNEL, meanWeights, stdWeights, Weights{});
assert(NormaLayer != nullptr);
return NormaLayer;
}
static IScaleLayer* NormalizeTeacherMap(INetworkDefinition* network, std::map<std::string, Weights>& weightMap,
ITensor& input) {
float* mean = (float*)weightMap["mean_std.mean"].values;
float* std = (float*)weightMap["mean_std.std"].values;
int len = weightMap["mean_std.mean"].count;
// 1.scale
float* scaleVal = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scaleVal[i] = 1.0 / std[i];
}
Weights scale{DataType::kFLOAT, scaleVal, len};
// 2.shift
float* shiftVal = nullptr;
shiftVal = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shiftVal[i] = -mean[i];
}
Weights shift{DataType::kFLOAT, shiftVal, len};
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, Weights{}, Weights{});
assert(scale_1);
IScaleLayer* scale_2 = network->addScale(*scale_1->getOutput(0), ScaleMode::kCHANNEL, Weights{}, scale, Weights{});
assert(scale_2);
return scale_2;
}
static ILayer* NormalizeFinalMap(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
std::string name) {
float* qa = (float*)weightMap["quantiles.qa_" + name].values;
float* qb = (float*)weightMap["quantiles.qb_" + name].values;
int len = weightMap["quantiles.qa_" + name].count;
Weights qbWeight_2{DataType::kFLOAT, qb, len};
// fmap_st - qa_st
float* shiftVal_1 = nullptr;
shiftVal_1 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shiftVal_1[i] = -qa[i];
}
Weights qa_shiftWeight_1{DataType::kFLOAT, shiftVal_1, len};
IScaleLayer* mapNorm_subLayer_1 =
network->addScale(input, ScaleMode::kUNIFORM, qa_shiftWeight_1, Weights{}, Weights{});
assert(mapNorm_subLayer_1);
// qb_st - qa_st
float* shiftVal_2 = nullptr;
shiftVal_2 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shiftVal_2[i] = qb[i] - qa[i];
}
// (fmap_st - qa_st) / (qb_st - qa_st)
float* scaleVal_1 = nullptr;
scaleVal_1 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scaleVal_1[i] = 1.0f / shiftVal_2[i];
}
Weights scaleWeight_1{DataType::kFLOAT, scaleVal_1, len};
IScaleLayer* mapNorm_divLayer_1 = network->addScale(*mapNorm_subLayer_1->getOutput(0), ScaleMode::kUNIFORM,
Weights{}, scaleWeight_1, Weights{});
assert(mapNorm_divLayer_1);
// ((fmap_st - qa_st) / (qb_st - qa_st)) * 0.1
float* scaleVal_2 = nullptr;
scaleVal_2 = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scaleVal_2[i] = 0.1f;
}
Weights scaleWeight_2{DataType::kFLOAT, scaleVal_2, 1};
IScaleLayer* mapNorm_Layer = network->addScale(*mapNorm_divLayer_1->getOutput(0), ScaleMode::kUNIFORM, Weights{},
scaleWeight_2, Weights{});
assert(mapNorm_Layer);
return mapNorm_Layer;
}
static ILayer* convRelu(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int outch, int ksize, int s, int p, int g, std::string lname, bool withRelu) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(
input, outch, DimsHW{ksize, ksize}, weightMap[lname + ".weight"],
weightMap[lname + ".bias"]); // if without bias weights, the results won't match with torch version
assert(conv1);
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
conv1->setNbGroups(g);
conv1->setName((lname).c_str());
if (!withRelu)
return conv1;
auto relu = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(relu);
return relu;
}
static IResizeLayer* interpolate(INetworkDefinition* network, ITensor& input, Dims upsampleScale,
ResizeMode resizeMode) {
IResizeLayer* interpolateLayer = network->addResize(input);
assert(interpolateLayer);
interpolateLayer->setOutputDimensions(upsampleScale);
interpolateLayer->setResizeMode(resizeMode);
return interpolateLayer;
}
static ILayer* interpConvRelu(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int outch, int ksize, int s, int p, int g, std::string lname, int dim) {
IResizeLayer* interpolateLayer = network->addResize(input);
assert(interpolateLayer != nullptr);
interpolateLayer->setOutputDimensions(Dims3{input.getDimensions().d[0], dim, dim});
interpolateLayer->setResizeMode(ResizeMode::kLINEAR);
IConvolutionLayer* conv1 = network->addConvolutionNd(*interpolateLayer->getOutput(0), outch, DimsHW{ksize, ksize},
weightMap[lname + ".weight"], weightMap[lname + ".bias"]);
assert(conv1);
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
conv1->setNbGroups(g);
conv1->setName((lname + ".conv").c_str());
auto relu = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
assert(relu);
return relu;
}
static IPoolingLayer* avgPool2d(INetworkDefinition* network, ITensor& input, int kernelSize, int stride, int padding) {
IPoolingLayer* poolLayer = network->addPooling(input, PoolingType::kAVERAGE, DimsHW{kernelSize, kernelSize});
assert(poolLayer);
poolLayer->setStride(DimsHW{stride, stride});
poolLayer->setPadding(DimsHW{padding, padding});
return poolLayer;
}
static void slice(INetworkDefinition* network, ITensor& input, std::vector<ITensor*>& layer_vec) {
Dims inputDims = input.getDimensions();
ISliceLayer* slice1 = network->addSlice(input, Dims3{0, 0, 0},
Dims3{inputDims.d[0] / 2, inputDims.d[1], inputDims.d[2]}, Dims3{1, 1, 1});
assert(slice1);
ISliceLayer* slice2 = network->addSlice(input, Dims3{inputDims.d[0] / 2, 0, 0},
Dims3{inputDims.d[0] / 2, inputDims.d[1], inputDims.d[2]}, Dims3{1, 1, 1});
assert(slice2);
layer_vec.push_back(slice1->getOutput(0));
layer_vec.push_back(slice2->getOutput(0));
}
static IElementWiseLayer* mergeMap(INetworkDefinition* network, ITensor& input1, ITensor& input2) {
float* scaleVal = nullptr;
scaleVal = reinterpret_cast<float*>(malloc(sizeof(float) * 1));
for (int i = 0; i < 1; i++) {
scaleVal[i] = 0.5f;
}
Weights scaleWeight{DataType::kFLOAT, scaleVal, 1};
IScaleLayer* mergeMapLayer1 = network->addScale(input1, ScaleMode::kUNIFORM, Weights{}, scaleWeight, Weights{});
assert(mergeMapLayer1);
IScaleLayer* mergeMapLayer2 = network->addScale(input2, ScaleMode::kUNIFORM, Weights{}, scaleWeight, Weights{});
assert(mergeMapLayer2);
IElementWiseLayer* mergedMapLayer = network->addElementWise(
*mergeMapLayer1->getOutput(0), *mergeMapLayer2->getOutput(0), ElementWiseOperation::kSUM);
assert(mergedMapLayer);
return mergedMapLayer;
}
ICudaEngine* build_efficientAD_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt,
float& gd, float& gw, std::string& wts_name) {
/* create network object */
INetworkDefinition* network = builder->createNetworkV2(0U);
/* create input tensor {3, kInputH, kInputW} */
ITensor* InputData = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW});
assert(InputData);
/* create weight map */
std::map<std::string, Weights> weightMap = loadWeights(wts_name);
/* AE */
// auto BN1 = NormalizeInput(network, *InputData);
// encoder
auto enconv1 = convRelu(network, weightMap, *InputData, 32, 4, 2, 1, 1, "ae.encoder.enconv1", true);
auto enconv2 = convRelu(network, weightMap, *enconv1->getOutput(0), 32, 4, 2, 1, 1, "ae.encoder.enconv2", true);
auto enconv3 = convRelu(network, weightMap, *enconv2->getOutput(0), 64, 4, 2, 1, 1, "ae.encoder.enconv3", true);
auto enconv4 = convRelu(network, weightMap, *enconv3->getOutput(0), 64, 4, 2, 1, 1, "ae.encoder.enconv4", true);
auto enconv5 = convRelu(network, weightMap, *enconv4->getOutput(0), 64, 4, 2, 1, 1, "ae.encoder.enconv5", true);
auto enconv6 = convRelu(network, weightMap, *enconv5->getOutput(0), 64, 8, 1, 0, 1, "ae.encoder.enconv6", false);
// decoder
auto deconv1 = interpConvRelu(network, weightMap, *enconv6->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv1", 3);
auto deconv2 = interpConvRelu(network, weightMap, *deconv1->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv2", 8);
auto deconv3 = interpConvRelu(network, weightMap, *deconv2->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv3", 15);
auto deconv4 = interpConvRelu(network, weightMap, *deconv3->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv4", 32);
auto deconv5 = interpConvRelu(network, weightMap, *deconv4->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv5", 63);
auto deconv6 =
interpConvRelu(network, weightMap, *deconv5->getOutput(0), 64, 4, 1, 2, 1, "ae.decoder.deconv6", 127);
auto deconv7 = interpConvRelu(network, weightMap, *deconv6->getOutput(0), 64, 3, 1, 1, 1, "ae.decoder.deconv7", 56);
auto deconv8 = convRelu(network, weightMap, *deconv7->getOutput(0), 384, 3, 1, 1, 1, "ae.decoder.deconv8", false);
/* PDN_medium_teacher */
// no BN added after the convolutional layer
auto teacher1 = convRelu(network, weightMap, *InputData, 256, 4, 1, 0, 1, "teacher.conv1", true);
auto avgPool1 = avgPool2d(network, *teacher1->getOutput(0), 2, 2, 0);
auto teacher2 = convRelu(network, weightMap, *avgPool1->getOutput(0), 512, 4, 1, 0, 1, "teacher.conv2", true);
auto avgPool2 = avgPool2d(network, *teacher2->getOutput(0), 2, 2, 0);
auto teacher3 = convRelu(network, weightMap, *avgPool2->getOutput(0), 512, 1, 1, 0, 1, "teacher.conv3", true);
auto teacher4 = convRelu(network, weightMap, *teacher3->getOutput(0), 512, 3, 1, 0, 1, "teacher.conv4", true);
auto teacher5 = convRelu(network, weightMap, *teacher4->getOutput(0), 384, 4, 1, 0, 1, "teacher.conv5", true);
auto teacher6 = convRelu(network, weightMap, *teacher5->getOutput(0), 384, 1, 1, 0, 1, "teacher.conv6", false);
/* PDN_medium_student */
auto student1 = convRelu(network, weightMap, *InputData, 256, 4, 1, 0, 1, "student.conv1", true);
auto avgPool3 = avgPool2d(network, *student1->getOutput(0), 2, 2, 0);
auto student2 = convRelu(network, weightMap, *avgPool3->getOutput(0), 512, 4, 1, 0, 1, "student.conv2", true);
auto avgPool4 = avgPool2d(network, *student2->getOutput(0), 2, 2, 0);
auto student3 = convRelu(network, weightMap, *avgPool4->getOutput(0), 512, 1, 1, 0, 1, "student.conv3", true);
auto student4 = convRelu(network, weightMap, *student3->getOutput(0), 512, 3, 1, 0, 1, "student.conv4", true);
auto student5 = convRelu(network, weightMap, *student4->getOutput(0), 768, 4, 1, 0, 1, "student.conv5", true);
auto student6 = convRelu(network, weightMap, *student5->getOutput(0), 768, 1, 1, 0, 1, "student.conv6", false);
/* postCalculate */
auto normal_teacher_output = NormalizeTeacherMap(network, weightMap, *teacher6->getOutput(0));
std::vector<ITensor*> layer_vec{};
slice(network, *student6->getOutput(0), layer_vec);
ITensor* y_st = layer_vec[0];
ITensor* y_stae = layer_vec[1];
// distance_st
IElementWiseLayer* sub_st =
network->addElementWise(*normal_teacher_output->getOutput(0), *y_st, ElementWiseOperation::kSUB);
assert(sub_st);
IElementWiseLayer* distance_st =
network->addElementWise(*sub_st->getOutput(0), *sub_st->getOutput(0), ElementWiseOperation::kPROD);
assert(distance_st);
// distance_stae
IElementWiseLayer* sub_stae = network->addElementWise(*deconv8->getOutput(0), *y_stae, ElementWiseOperation::kSUB);
assert(sub_stae);
IElementWiseLayer* distance_stae =
network->addElementWise(*sub_stae->getOutput(0), *sub_stae->getOutput(0), ElementWiseOperation::kPROD);
assert(distance_stae);
IReduceLayer* map_st = network->addReduce(*distance_st->getOutput(0), ReduceOperation::kAVG, 1, true);
assert(map_st);
IReduceLayer* map_stae = network->addReduce(*distance_stae->getOutput(0), ReduceOperation::kAVG, 1, true);
assert(map_stae);
IPaddingLayer* padMap_st = network->addPadding(*map_st->getOutput(0), DimsHW{4, 4}, DimsHW{4, 4});
assert(padMap_st);
IPaddingLayer* padMap_stae = network->addPadding(*map_stae->getOutput(0), DimsHW{4, 4}, DimsHW{4, 4});
assert(padMap_stae);
IResizeLayer* interpMap_st =
interpolate(network, *padMap_st->getOutput(0),
Dims3{padMap_st->getOutput(0)->getDimensions().d[0], 256, 256}, ResizeMode::kLINEAR);
assert(interpMap_st);
IResizeLayer* interpMap_stae =
interpolate(network, *padMap_stae->getOutput(0),
Dims3{padMap_stae->getOutput(0)->getDimensions().d[0], 256, 256}, ResizeMode::kLINEAR);
assert(interpMap_stae);
ILayer* normalizedMap_st = NormalizeFinalMap(network, weightMap, *interpMap_st->getOutput(0), "st");
assert(normalizedMap_st);
ILayer* normalizedMap_stae = NormalizeFinalMap(network, weightMap, *interpMap_stae->getOutput(0), "ae");
assert(normalizedMap_stae);
IElementWiseLayer* mergedMapLayer =
mergeMap(network, *normalizedMap_st->getOutput(0), *normalizedMap_st->getOutput(0));
printNetworkLayers(network);
/* ouput */
mergedMapLayer->getOutput(0)->setName(kOutputTensorName);
network->markOutput(*mergedMapLayer->getOutput(0));
/* Engine config */
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#if defined(USE_FP16)
config->setFlag(BuilderFlag::kFP16);
#elif defined(USE_INT8)
std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
assert(builder->platformHasFastInt8());
config->setFlag(BuilderFlag::kINT8);
Int8EntropyCalibrator2* calibrator =
new Int8EntropyCalibrator2(1, kInputW, kInputH, "./coco_calib/", "int8calib.table", kInputTensorName);
config->setInt8Calibrator(calibrator);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
}

9
efficient_ad/src/model.h Normal file
View File

@ -0,0 +1,9 @@
#pragma once
#include <NvInfer.h>
#include <string>
nvinfer1::ICudaEngine* build_efficientAD_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, float& gd,
float& gw, std::string& wts_name);

9
efficient_ad/src/postprocess.h Normal file
View File

@ -0,0 +1,9 @@
#pragma once
#include <opencv2/opencv.hpp>
void genHeatMap(cv::Mat originImg, cv::Mat& anomalyGrayMap, cv::Mat& HeatMap) {
cv::Mat colorMap;
cv::applyColorMap(colorMap, anomalyGrayMap, cv::COLORMAP_JET);
cv::addWeighted(originImg, 0.5, colorMap, 0.5, 0, HeatMap);
}

30
efficient_ad/src/utils.h Normal file
View File

@ -0,0 +1,30 @@
#pragma once
#include <dirent.h>
#include <cstring>
#include <fstream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& file_names) {
DIR* p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {
return -1;
}
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 && strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;
std::string cur_file_name(p_file->d_name);
file_names.push_back(cur_file_name);
}
}
closedir(p_dir);
return 0;
}

27
efficientnet/CMakeLists.txt Normal file
View File

@ -0,0 +1,27 @@
cmake_minimum_required(VERSION 2.6)
project(efficientnet)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
add_executable(efficientnet ${PROJECT_SOURCE_DIR}/efficientnet.cpp)
target_link_libraries(efficientnet nvinfer)
target_link_libraries(efficientnet cudart)
add_definitions(-O2 -pthread)

45
efficientnet/README.md Normal file
View File

@ -0,0 +1,45 @@
# EfficientNet
A TensorRT implementation of EfficientNet.
For the Pytorch implementation, you can refer to [EfficientNet-PyTorch](https://github.com/lukemelas/EfficientNet-PyTorch)
## How to run
1. install `efficientnet_pytorch`
```
pip install efficientnet_pytorch
```
2. gennerate `.wts` file
```
python gen_wts.py
```
3. build
```
mkdir build
cd build
cmake ..
make
```
4. serialize model to engine
```
./efficientnet -s [.wts] [.engine] [b0 b1 b2 b3 ... b7] // serialize model to engine file
```
such as
```
./efficientnet -s ../efficientnet-b3.wts efficientnet-b3.engine b3
```
5. deserialize and do infer
```
./efficientnet -d [.engine] [b0 b1 b2 b3 ... b7] // deserialize engine file and run inference
```
such as
```
./efficientnet -d efficientnet-b3.engine b3
```
6. see if the output is same as pytorch side
For more models, please refer to [tensorrtx](https://github.com/wang-xinyu/tensorrtx)

280
efficientnet/efficientnet.cpp Normal file
View File

@ -0,0 +1,280 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include "utils.hpp"
#define USE_FP32 //USE_FP16
#define INPUT_NAME "data"
#define OUTPUT_NAME "prob"
#define MAX_BATCH_SIZE 8
using namespace nvinfer1;
static Logger gLogger;
static std::vector<BlockArgs>
block_args_list = {
BlockArgs{1, 3, 1, 1, 32, 16, 0.25, true},
BlockArgs{2, 3, 2, 6, 16, 24, 0.25, true},
BlockArgs{2, 5, 2, 6, 24, 40, 0.25, true},
BlockArgs{3, 3, 2, 6, 40, 80, 0.25, true},
BlockArgs{3, 5, 1, 6, 80, 112, 0.25, true},
BlockArgs{4, 5, 2, 6, 112, 192, 0.25, true},
BlockArgs{1, 3, 1, 6, 192, 320, 0.25, true}};
static std::map<std::string, GlobalParams>
global_params_map = {
// input_h,input_w,num_classes,batch_norm_epsilon,
// width_coefficient,depth_coefficient,depth_divisor, min_depth
{"b0", GlobalParams{224, 224, 1000, 0.001, 1.0, 1.0, 8, -1}},
{"b1", GlobalParams{240, 240, 1000, 0.001, 1.0, 1.1, 8, -1}},
{"b2", GlobalParams{260, 260, 1000, 0.001, 1.1, 1.2, 8, -1}},
{"b3", GlobalParams{300, 300, 1000, 0.001, 1.2, 1.4, 8, -1}},
{"b4", GlobalParams{380, 380, 1000, 0.001, 1.4, 1.8, 8, -1}},
{"b5", GlobalParams{456, 456, 1000, 0.001, 1.6, 2.2, 8, -1}},
{"b6", GlobalParams{528, 528, 1000, 0.001, 1.8, 2.6, 8, -1}},
{"b7", GlobalParams{600, 600, 1000, 0.001, 2.0, 3.1, 8, -1}},
{"b8", GlobalParams{672, 672, 1000, 0.001, 2.2, 3.6, 8, -1}},
{"l2", GlobalParams{800, 800, 1000, 0.001, 4.3, 5.3, 8, -1}},
};
ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt, std::string path_wts, std::vector<BlockArgs> block_args_list, GlobalParams global_params)
{
float bn_eps = global_params.batch_norm_epsilon;
DimsHW image_size = DimsHW{global_params.input_h, global_params.input_w};
std::map<std::string, Weights> weightMap = loadWeights(path_wts);
Weights emptywts{DataType::kFLOAT, nullptr, 0};
INetworkDefinition *network = builder->createNetworkV2(0U);
ITensor *data = network->addInput(INPUT_NAME, dt, Dims3{3, global_params.input_h, global_params.input_w});
assert(data);
int out_channels = roundFilters(32, global_params);
auto conv_stem = addSamePaddingConv2d(network, weightMap, *data, out_channels, 3, 2, 1, 1, image_size, "_conv_stem");
auto bn0 = addBatchNorm2d(network, weightMap, *conv_stem->getOutput(0), "_bn0", bn_eps);
auto swish0 = addSwish(network, *bn0->getOutput(0));
ITensor *x = swish0->getOutput(0);
image_size = calculateOutputImageSize(image_size, 2);
int block_id = 0;
for (int i = 0; i < block_args_list.size(); i++)
{
BlockArgs block_args = block_args_list[i];
block_args.input_filters = roundFilters(block_args.input_filters, global_params);
block_args.output_filters = roundFilters(block_args.output_filters, global_params);
block_args.num_repeat = roundRepeats(block_args.num_repeat, global_params);
x = MBConvBlock(network, weightMap, *x, "_blocks." + std::to_string(block_id), block_args, global_params, image_size);
assert(x);
block_id++;
image_size = calculateOutputImageSize(image_size, block_args.stride);
if (block_args.num_repeat > 1)
{
block_args.input_filters = block_args.output_filters;
block_args.stride = 1;
}
for (int r = 0; r < block_args.num_repeat - 1; r++)
{
x = MBConvBlock(network, weightMap, *x, "_blocks." + std::to_string(block_id), block_args, global_params, image_size);
block_id++;
}
}
out_channels = roundFilters(1280, global_params);
auto conv_head = addSamePaddingConv2d(network, weightMap, *x, out_channels, 1, 1, 1, 1, image_size, "_conv_head", false);
auto bn1 = addBatchNorm2d(network, weightMap, *conv_head->getOutput(0), "_bn1", bn_eps);
auto swish1 = addSwish(network, *bn1->getOutput(0));
auto avg_pool = network->addPoolingNd(*swish1->getOutput(0), PoolingType::kAVERAGE, image_size);
IFullyConnectedLayer *final = network->addFullyConnected(*avg_pool->getOutput(0), global_params.num_classes, weightMap["_fc.weight"], weightMap["_fc.bias"]);
assert(final);
final->getOutput(0)->setName(OUTPUT_NAME);
network->markOutput(*final->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(1 << 20);
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "build engine ..." << std::endl;
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
assert(engine != nullptr);
std::cout << "build finished" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto &mem : weightMap)
{
free((void *)(mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream, std::string wtsPath, std::vector<BlockArgs> block_args_list, GlobalParams global_params)
{
// Create builder
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine *engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT, wtsPath, block_args_list, global_params);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
config->destroy();
}
void doInference(IExecutionContext &context, float *input, float *output, int batchSize, GlobalParams global_params)
{
const ICudaEngine &engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void *buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * global_params.input_h * global_params.input_w * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * global_params.num_classes * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * global_params.input_h * global_params.input_w * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * global_params.num_classes * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &backbone)
{
if (std::string(argv[1]) == "-s" && argc == 5)
{
wts = std::string(argv[2]);
engine = std::string(argv[3]);
backbone = std::string(argv[4]);
}
else if (std::string(argv[1]) == "-d" && argc == 4)
{
engine = std::string(argv[2]);
backbone = std::string(argv[3]);
}
else
{
return false;
}
return true;
}
int main(int argc, char **argv)
{
std::string wtsPath = "";
std::string engine_name = "";
std::string backbone = "";
if (!parse_args(argc, argv, wtsPath, engine_name, backbone))
{
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./efficientnet -s [.wts] [.engine] [b0 b1 b2 b3 ... b7] // serialize model to engine file" << std::endl;
std::cerr << "./efficientnet -d [.engine] [b0 b1 b2 b3 ... b7] // deserialize engine file and run inference" << std::endl;
return -1;
}
GlobalParams global_params = global_params_map[backbone];
// create a model using the API directly and serialize it to a stream
if (!wtsPath.empty())
{
IHostMemory *modelStream{nullptr};
APIToModel(MAX_BATCH_SIZE, &modelStream, wtsPath, block_args_list, global_params);
assert(modelStream != nullptr);
std::ofstream p(engine_name, std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 1;
}
char *trtModelStream{nullptr};
size_t size{0};
std::ifstream file(engine_name, std::ios::binary);
if (file.good())
{
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
else
{
std::cerr << "could not open plan file" << std::endl;
return -1;
}
// dummy input
float *data = new float[3 * global_params.input_h * global_params.input_w];
for (int i = 0; i < 3 * global_params.input_h * global_params.input_w; i++)
data[i] = 0.1;
IRuntime *runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext *context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Run inference
float *prob = new float[global_params.num_classes];
for (int i = 0; i < 100; i++)
{
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, 1, global_params);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
for (unsigned int i = 0; i < 20; i++)
{
std::cout << prob[i] << ", ";
}
std::cout << std::endl;
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
delete data;
delete prob;
return 0;
}

16
efficientnet/gen_wts.py Normal file
View File

@ -0,0 +1,16 @@
import torch
import struct
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b3')
model.eval()
f = open('efficientnet-b3.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')
f.close()

503
efficientnet/logging.h Normal file
View File

@ -0,0 +1,503 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf
{
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream)
, mPrefix(prefix)
, mShouldLog(shouldLog)
{
}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)
: mOutput(other.mOutput)
{
}
~LogStreamConsumerBuffer()
{
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr())
{
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync()
{
putOutput();
return 0;
}
void putOutput()
{
if (mShouldLog)
{
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog)
{
mShouldLog = shouldLog;
}
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase
{
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog)
{
}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream
{
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(severity <= reportableSeverity)
, mSeverity(severity)
{
}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)
, std::ostream(&mBuffer) // links the stream buffer with the stream
, mShouldLog(other.mShouldLog)
, mSeverity(other.mSeverity)
{
}
void setReportableSeverity(Severity reportableSeverity)
{
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger
{
public:
Logger(Severity severity = Severity::kWARNING)
: mReportableSeverity(severity)
{
}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult
{
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger()
{
return *this;
}
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity)
{
mReportableSeverity = severity;
}
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom
{
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started)
, mName(name)
, mCmdline(cmdline)
{
}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline)
{
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)
{
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom)
{
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result)
{
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom)
{
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass)
{
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const
{
return mReportableSeverity;
}
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity)
{
switch (severity)
{
case Severity::kINTERNAL_ERROR: return "[F] ";
case Severity::kERROR: return "[E] ";
case Severity::kWARNING: return "[W] ";
case Severity::kINFO: return "[I] ";
case Severity::kVERBOSE: return "[V] ";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result)
{
switch (result)
{
case TestResult::kRUNNING: return "RUNNING";
case TestResult::kPASSED: return "PASSED";
case TestResult::kFAILED: return "FAILED";
case TestResult::kWAIVED: return "WAIVED";
default: assert(0); return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity)
{
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result)
{
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv)
{
std::stringstream ss;
for (int i = 0; i < argc; i++)
{
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace
{
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger)
{
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

251
efficientnet/utils.hpp Normal file
View File

@ -0,0 +1,251 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <math.h>
#include <string>
#include <algorithm>
using namespace nvinfer1;
#define CHECK(status) \
do \
{ \
auto ret = (status); \
if (ret != 0) \
{ \
std::cerr << "Cuda failure: " << ret << std::endl; \
abort(); \
} \
} while (0)
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
struct BlockArgs
{
int num_repeat;
int kernel_size;
int stride;
float expand_ratio;
int input_filters;
int output_filters;
float se_ratio;
bool id_skip;
};
struct GlobalParams
{
int input_h;
int input_w;
int num_classes;
float batch_norm_epsilon;
float width_coefficient;
float depth_coefficient;
int depth_divisor;
int min_depth;
};
int roundFilters(int filters, GlobalParams global_params)
{
float multiplier = global_params.width_coefficient;
int divisor = global_params.depth_divisor;
int min_depth = global_params.min_depth;
filters = int(filters * multiplier);
if (min_depth < 0)
{
min_depth = divisor;
}
// follow the formula transferred from official TensorFlow implementation
int new_filters = std::max(min_depth, int(int(filters + divisor / 2) / divisor) * divisor);
if (new_filters < 0.9 * filters) // prevent rounding by more than 10%
new_filters += divisor;
return int(new_filters);
}
DimsHW calculateOutputImageSize(DimsHW image_size, int stride)
{
int image_h = int(ceil(float(image_size.h()) / float(stride)));
int image_w = int(ceil(float(image_size.w()) / float(stride)));
return DimsHW{image_h, image_w};
}
int roundRepeats(int repeats, GlobalParams global_params)
{
float multiplier = global_params.depth_coefficient;
// follow the formula transferred from official TensorFlow implementation
int new_repeats = int(ceil(multiplier * repeats));
return new_repeats;
}
IScaleLayer *addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, std::string lname, float eps)
{
float *gamma = (float *)weightMap[lname + ".weight"].values;
float *beta = (float *)weightMap[lname + ".bias"].values;
float *mean = (float *)weightMap[lname + ".running_mean"].values;
float *var = (float *)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
float *scval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++)
{
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++)
{
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++)
{
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer *scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IConvolutionLayer *addSamePaddingConv2d(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int outch, int kernel_size, int stride, int dilation, int groups, DimsHW image_size, std::string lname, bool bias = true)
{
int ih = image_size.h();
int iw = image_size.w();
int kh = kernel_size;
int kw = kernel_size;
int sh = stride;
int sw = stride;
int oh = ceil(float(ih) / float(sh));
int ow = ceil(float(iw) / float(sw));
int pad_h = std::max((oh - 1) * stride + (kh - 1) * dilation + 1 - ih, 0);
int pad_w = std::max((ow - 1) * stride + (kw - 1) * dilation + 1 - iw, 0);
int pad_left = 0;
int pad_right = 0;
int pad_top = 0;
int pad_bottom = 0;
if (pad_h > 0 || pad_w > 0)
{
pad_left = int(pad_w / 2);
pad_right = pad_w - int(pad_w / 2);
pad_top = int(pad_h / 2);
pad_bottom = pad_h - int(pad_h / 2);
}
Weights bias_wt{DataType::kFLOAT, nullptr, 0};
if (bias)
{
bias_wt = weightMap[lname + ".bias"];
}
IConvolutionLayer *conv = network->addConvolutionNd(input, outch, DimsHW{kh, kw}, weightMap[lname + ".weight"], bias_wt);
conv->setPrePadding(DimsHW{pad_top, pad_left});
conv->setPostPadding(DimsHW{pad_bottom, pad_right});
conv->setStrideNd(DimsHW{stride, stride});
conv->setDilationNd(DimsHW{dilation, dilation});
conv->setNbGroups(groups);
return conv;
}
ILayer *addSwish(INetworkDefinition *network, ITensor &input)
{
//swish
auto *sigmoid = network->addActivation(input, ActivationType::kSIGMOID);
auto *ew = network->addElementWise(input, *sigmoid->getOutput(0), ElementWiseOperation::kPROD);
return ew;
}
ITensor *MBConvBlock(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, std::string lname, BlockArgs block_args, GlobalParams global_params, DimsHW image_size)
{
bool has_se = block_args.se_ratio > 0 && block_args.se_ratio <= 1;
bool id_skip = block_args.id_skip;
float bn_eps = global_params.batch_norm_epsilon;
int input_filters = block_args.input_filters;
int output_filters = block_args.output_filters;
Weights emptywts{DataType::kFLOAT, nullptr, 0};
ITensor *x = &input;
int inp = block_args.input_filters;
int oup = int(block_args.input_filters * block_args.expand_ratio);
// expand_ratio != 1
if (fabs(block_args.expand_ratio - 1) > 1e-5)
{
auto expand_conv = addSamePaddingConv2d(network, weightMap, input, oup, 1, 1, 1, 1, image_size, lname + "._expand_conv");
auto bn0 = addBatchNorm2d(network, weightMap, *expand_conv->getOutput(0), lname + "._bn0", bn_eps);
auto swish0 = addSwish(network, *bn0->getOutput(0));
x = swish0->getOutput(0);
}
int k = block_args.kernel_size;
int s = block_args.stride;
auto depthwise_conv = addSamePaddingConv2d(network, weightMap, *x, oup, k, s, 1, oup, image_size, lname + "._depthwise_conv", false);
auto bn1 = addBatchNorm2d(network, weightMap, *depthwise_conv->getOutput(0), lname + "._bn1", bn_eps);
//swish
auto swish1 = addSwish(network, *bn1->getOutput(0));
x = swish1->getOutput(0);
image_size = calculateOutputImageSize(image_size, s);
if (has_se)
{
auto avg_pool = network->addPoolingNd(*x, PoolingType::kAVERAGE, image_size);
int num_squeezed_channels = std::max(1, int(input_filters * block_args.se_ratio));
auto se_reduce = addSamePaddingConv2d(network, weightMap, *avg_pool->getOutput(0), num_squeezed_channels, 1, 1, 1, 1, DimsHW{1, 1}, lname + "._se_reduce");
auto swish2 = addSwish(network, *se_reduce->getOutput(0));
auto se_expand = addSamePaddingConv2d(network, weightMap, *swish2->getOutput(0), oup, 1, 1, 1, 1, DimsHW{1, 1}, lname + "._se_expand");
auto *sigmoid = network->addActivation(*se_expand->getOutput(0), ActivationType::kSIGMOID);
auto *ew = network->addElementWise(*x, *sigmoid->getOutput(0), ElementWiseOperation::kPROD);
x = ew->getOutput(0);
}
int final_oup = block_args.output_filters;
auto project_conv = addSamePaddingConv2d(network, weightMap, *x, final_oup, 1, 1, 1, 1, image_size, lname + "._project_conv");
auto bn2 = addBatchNorm2d(network, weightMap, *project_conv->getOutput(0), lname + "._bn2", bn_eps);
x = bn2->getOutput(0);
if (id_skip && block_args.stride == 1 && input_filters == output_filters)
{
auto *ew = network->addElementWise(input, *x, ElementWiseOperation::kSUM);
x = ew->getOutput(0);
}
return x;
}

82
ghostnet/README.md Normal file
View File

@ -0,0 +1,82 @@
# GhostNet
GhostNetv1 architecture is from the paper "GhostNet: More Features from Cheap Operations" [(https://arxiv.org/abs/1911.11907)](https://arxiv.org/abs/1911.11907).
GhostNetv2 architecture is from the paper "GhostNetV2: Enhance Cheap Operation with Long-Range Attention" [(https://arxiv.org/abs/2211.12905)](https://arxiv.org/abs/2211.12905).
For the PyTorch implementations, you can refer to [huawei-noah/ghostnet](https://github.com/huawei-noah/ghostnet).
Both versions use the following techniques in their TensorRT implementations:
- **BatchNorm** layer is implemented by TensorRT's **Scale** layer.
- **Ghost Modules** are used to generate more features from cheap operations, as described in the paper.
- Replacing `IPoolingLayer` with `IReduceLayer` in TensorRT for Global Average Pooling. The `IReduceLayer` allows you to perform reduction operations (such as sum, average, max) over specified dimensions without being constrained by the kernel size limitations of pooling layers.
## Project Structure
```plaintext
ghostnet
│
├── ghostnetv1
│ ├── CMakeLists.txt
│ ├── gen_wts.py
│ ├── ghostnetv1.cpp
│ └── logging.h
│
├── ghostnetv2
│ ├── CMakeLists.txt
│ ├── gen_wts.py
│ ├── ghostnetv2.cpp
│ └── logging.h
│
└── README.md
```
## Steps to use GhostNet in TensorRT
### 1. Generate `.wts` files for both GhostNetv1 and GhostNetv2
```bash
# For ghostnetv1
python ghostnetv1/gen_wts.py
# For ghostnetv2
python ghostnetv2/gen_wts.py
```
### 2. Build the project
```bash
cd tensorrtx/ghostnet
mkdir build
cd build
cmake ..
make
```
### 3. Serialize the models to engine files
Use the following commands to serialize the PyTorch models into TensorRT engine files (`ghostnetv1.engine` and `ghostnetv2.engine`):
```bash
# For ghostnetv1
sudo ./ghostnetv1 -s
# For ghostnetv2
sudo ./ghostnetv2 -s
```
### 4. Run inference using the engine files
Once the engine files are generated, you can run inference with the following commands:
```bash
# For ghostnetv1
sudo ./ghostnetv1 -d
# For ghostnetv2
sudo ./ghostnetv2 -d
```
### 5. Verify output
Compare the output with the PyTorch implementation from [huawei-noah/ghostnet](https://github.com/huawei-noah/ghostnet) to ensure that the TensorRT results are consistent with the PyTorch model.

24
ghostnet/ghostnetv1/CMakeLists.txt Normal file
View File

@ -0,0 +1,24 @@
cmake_minimum_required(VERSION 2.6)
project(ghostnetv1)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
add_executable(ghostnetv1 ${PROJECT_SOURCE_DIR}/ghostnetv1.cpp)
target_link_libraries(ghostnetv1 nvinfer)
target_link_libraries(ghostnetv1 cudart)
add_definitions(-O2 -pthread)

292
ghostnet/ghostnetv1/gen_wts.py Normal file
View File

@ -0,0 +1,292 @@
"""
Creates a GhostNet Model as defined in:
GhostNet: More Features from Cheap Operations By Kai Han, Yunhe Wang, Qi Tian, Jianyuan Guo, Chunjing Xu, Chang Xu.
https://arxiv.org/abs/1911.11907
Modified from https://github.com/d-li14/mobilenetv3.pytorch and https://github.com/rwightman/pytorch-image-models
"""
import torch
import torch.nn as nn
import torch.onnx
import struct
import torch
import torch.nn.functional as F
import math
def _make_divisible(v, divisor, min_value=None):
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
def hard_sigmoid(x, inplace: bool = False):
if inplace:
return x.add_(3.).clamp_(0., 6.).div_(6.)
else:
return F.relu6(x + 3.) / 6.
class SqueezeExcite(nn.Module):
def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None,
act_layer=nn.ReLU, gate_fn=hard_sigmoid, divisor=4, **_):
super(SqueezeExcite, self).__init__()
self.gate_fn = gate_fn
reduced_chs = _make_divisible((reduced_base_chs or in_chs) * se_ratio, divisor)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)
self.act1 = act_layer(inplace=True)
self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)
def forward(self, x):
x_se = self.avg_pool(x)
x_se = self.conv_reduce(x_se)
x_se = self.act1(x_se)
x_se = self.conv_expand(x_se)
x = x * self.gate_fn(x_se)
return x
class ConvBnAct(nn.Module):
def __init__(self, in_chs, out_chs, kernel_size,
stride=1, act_layer=nn.ReLU):
super(ConvBnAct, self).__init__()
self.conv = nn.Conv2d(in_chs, out_chs, kernel_size, stride, kernel_size//2, bias=False)
self.bn1 = nn.BatchNorm2d(out_chs)
self.act1 = act_layer(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn1(x)
x = self.act1(x)
return x
class GhostModule(nn.Module):
def __init__(self, inp, oup, kernel_size=1, ratio=2, dw_size=3, stride=1, relu=True):
super(GhostModule, self).__init__()
self.oup = oup
init_channels = math.ceil(oup / ratio)
new_channels = init_channels*(ratio-1)
self.primary_conv = nn.Sequential(
nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size//2, bias=False),
nn.BatchNorm2d(init_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
self.cheap_operation = nn.Sequential(
nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size//2, groups=init_channels, bias=False),
nn.BatchNorm2d(new_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
def forward(self, x):
x1 = self.primary_conv(x)
x2 = self.cheap_operation(x1)
out = torch.cat([x1, x2], dim=1)
return out[:, :self.oup, :, :]
class GhostBottleneck(nn.Module):
""" Ghost bottleneck w/ optional SE"""
def __init__(self, in_chs, mid_chs, out_chs, dw_kernel_size=3,
stride=1, act_layer=nn.ReLU, se_ratio=0.):
super(GhostBottleneck, self).__init__()
has_se = se_ratio is not None and se_ratio > 0.
self.stride = stride
# Point-wise expansion
self.ghost1 = GhostModule(in_chs, mid_chs, relu=True)
# Depth-wise convolution
if self.stride > 1:
self.conv_dw = nn.Conv2d(mid_chs, mid_chs, dw_kernel_size, stride=stride,
padding=(dw_kernel_size-1)//2, groups=mid_chs, bias=False)
self.bn_dw = nn.BatchNorm2d(mid_chs)
# Squeeze-and-excitation
if has_se:
self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio)
else:
self.se = None
# Point-wise linear projection
self.ghost2 = GhostModule(mid_chs, out_chs, relu=False)
# shortcut
if (in_chs == out_chs and self.stride == 1):
self.shortcut = nn.Sequential()
else:
self.shortcut = nn.Sequential(
nn.Conv2d(in_chs, in_chs, dw_kernel_size, stride=stride,
padding=(dw_kernel_size-1)//2, groups=in_chs, bias=False),
nn.BatchNorm2d(in_chs),
nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(out_chs),
)
def forward(self, x):
residual = x
# 1st ghost bottleneck
x = self.ghost1(x)
# Depth-wise convolution
if self.stride > 1:
x = self.conv_dw(x)
x = self.bn_dw(x)
# Squeeze-and-excitation
if self.se is not None:
x = self.se(x)
# 2nd ghost bottleneck
x = self.ghost2(x)
x += self.shortcut(residual)
return x
class GhostNet(nn.Module):
def __init__(self, cfgs, num_classes=1000, width=1.0, dropout=0.2):
super(GhostNet, self).__init__()
# setting of inverted residual blocks
self.cfgs = cfgs
self.dropout = dropout
# building first layer
output_channel = _make_divisible(16 * width, 4)
self.conv_stem = nn.Conv2d(3, output_channel, 3, 2, 1, bias=False)
self.bn1 = nn.BatchNorm2d(output_channel)
self.act1 = nn.ReLU(inplace=True)
input_channel = output_channel
# building inverted residual blocks
stages = []
block = GhostBottleneck
for cfg in self.cfgs:
layers = []
for k, exp_size, c, se_ratio, s in cfg:
output_channel = _make_divisible(c * width, 4)
hidden_channel = _make_divisible(exp_size * width, 4)
layers.append(block(input_channel, hidden_channel, output_channel, k, s,
se_ratio=se_ratio))
input_channel = output_channel
stages.append(nn.Sequential(*layers))
output_channel = _make_divisible(exp_size * width, 4)
stages.append(nn.Sequential(ConvBnAct(input_channel, output_channel, 1)))
input_channel = output_channel
self.blocks = nn.Sequential(*stages)
# building last several layers
output_channel = 1280
self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
self.conv_head = nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True)
self.act2 = nn.ReLU(inplace=True)
self.classifier = nn.Linear(output_channel, num_classes)
def forward(self, x):
x = self.conv_stem(x)
x = self.bn1(x)
x = self.act1(x)
x = self.blocks(x)
x = self.global_pool(x)
x = self.conv_head(x)
x = self.act2(x)
x = x.view(x.size(0), -1)
if self.dropout > 0.:
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.classifier(x)
return x
def ghostnet(**kwargs):
"""
Constructs a GhostNet model
"""
cfgs = [
# k, t, c, SE, s
# stage1
[[3, 16, 16, 0, 1]],
# stage2
[[3, 48, 24, 0, 2]],
[[3, 72, 24, 0, 1]],
# stage3
[[5, 72, 40, 0.25, 2]],
[[5, 120, 40, 0.25, 1]],
# stage4
[[3, 240, 80, 0, 2]],
[[3, 200, 80, 0, 1],
[3, 184, 80, 0, 1],
[3, 184, 80, 0, 1],
[3, 480, 112, 0.25, 1],
[3, 672, 112, 0.25, 1]],
# stage5
[[5, 672, 160, 0.25, 2]],
[[5, 960, 160, 0, 1],
[5, 960, 160, 0.25, 1],
[5, 960, 160, 0, 1],
[5, 960, 160, 0.25, 1]]
]
return GhostNet(cfgs, **kwargs)
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
# Function to export weights in the specified format
def export_weight(model):
f = open("ghostnetv1.weights", 'w')
f.write("{}\n".format(len(model.state_dict().keys())))
# Convert weights to hexadecimal format
for k, v in model.state_dict().items():
print('exporting ... {}: {}'.format(k, v.shape))
# Reshape the weights to 1D
vr = v.reshape(-1).cpu().numpy()
f.write("{} {}".format(k, len(vr)))
for vv in vr:
f.write(" ")
f.write(struct.pack(">f", float(vv)).hex())
f.write("\n")
f.close()
# Function to evaluate the model (optional)
def eval_model(input, model):
output = model(input)
print("------from inference------")
print(input)
print(output)
if __name__ == "__main__":
setup_seed(1)
model = ghostnet(num_classes=1000, width=1.0, dropout=0.2)
model.eval()
input = torch.full((32, 3, 320, 256), 10.0)
export_weight(model)
eval_model(input, model)

516
ghostnet/ghostnetv1/ghostnetv1.cpp Normal file
View File

@ -0,0 +1,516 @@
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
using namespace std;
#define CHECK(status) \
do { \
auto ret = (status); \
if (ret != 0) { \
std::cerr << "Cuda failure: " << ret << std::endl; \
abort(); \
} \
} while (0)
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 256;
static const int INPUT_W = 320;
static const int OUTPUT_SIZE = 1000;
static const int batchSize = 32;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
using namespace nvinfer1;
static Logger gLogger;
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
if (!input.is_open()) {
std::cerr << "Unable to load weight file." << std::endl;
exit(EXIT_FAILURE);
}
// Read number of weight blobs
int32_t count;
input >> count;
if (count <= 0) {
std::cerr << "Invalid weight map file." << std::endl;
exit(EXIT_FAILURE);
}
while (count--) {
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(uint32_t) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
int _make_divisible(int v, int divisor, int min_value = -1) {
if (min_value == -1) {
min_value = divisor;
}
int new_v = std::max(min_value, (v + divisor / 2) / divisor * divisor);
if (new_v < static_cast<int>(0.9 * v)) {
new_v += divisor;
}
return new_v;
}
ILayer* hardSigmoid(INetworkDefinition* network, ITensor& input) {
IActivationLayer* scale_layer = network->addActivation(input, ActivationType::kHARD_SIGMOID);
return scale_layer;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
std::string lname, float eps) {
float* gamma = (float*)weightMap[lname + ".weight"].values;
float* beta = (float*)weightMap[lname + ".bias"].values;
float* mean = (float*)weightMap[lname + ".running_mean"].values;
float* var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
float* scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float* shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float* pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IActivationLayer* convBnReluStem(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int outch, std::string lname) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 =
network->addConvolutionNd(input, outch, DimsHW{3, 3}, weightMap[lname + ".weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{2, 2}); // Stride = 2
conv1->setPaddingNd(DimsHW{1, 1}); // Padding = 1
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "bn1", 1e-5);
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
return relu1;
}
ILayer* convBnAct(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int out_channels, std::string lname, ActivationType actType = ActivationType::kRELU) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv =
network->addConvolutionNd(input, out_channels, DimsHW{1, 1}, weightMap[lname + ".conv.weight"], emptywts);
assert(conv);
conv->setStrideNd(DimsHW{1, 1});
IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv->getOutput(0), lname + ".bn1", 1e-5);
IActivationLayer* act = network->addActivation(*bn->getOutput(0), actType);
assert(act);
return act;
}
ILayer* squeezeExcite(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap,
int in_chs, float se_ratio = 0.25, std::string lname = "", float eps = 1e-5) {
IReduceLayer* avg_pool = network->addReduce(input, ReduceOperation::kAVG, 1 << 2 | 1 << 3, true);
assert(avg_pool);
// Reduce channels with 1x1 convolution
int reduced_chs = _make_divisible(static_cast<int>(in_chs * se_ratio), 4);
IConvolutionLayer* conv_reduce =
network->addConvolutionNd(*avg_pool->getOutput(0), reduced_chs, DimsHW{1, 1},
weightMap[lname + ".conv_reduce.weight"], weightMap[lname + ".conv_reduce.bias"]);
assert(conv_reduce);
IActivationLayer* relu1 = network->addActivation(*conv_reduce->getOutput(0), ActivationType::kRELU);
assert(relu1);
// Expand channels back with another 1x1 convolution
IConvolutionLayer* conv_expand =
network->addConvolutionNd(*relu1->getOutput(0), in_chs, DimsHW{1, 1},
weightMap[lname + ".conv_expand.weight"], weightMap[lname + ".conv_expand.bias"]);
assert(conv_expand);
cout << "SE conv_expand -> " << printTensorShape(conv_expand->getOutput(0)) << endl;
// Apply hardSigmoid function
ILayer* hard_sigmoid = hardSigmoid(network, *conv_expand->getOutput(0));
cout << "hard_sigmoid conv_expand -> " << printTensorShape(hard_sigmoid->getOutput(0)) << endl;
// Elementwise multiplication of input and gated SE output
IElementWiseLayer* scale = network->addElementWise(input, *hard_sigmoid->getOutput(0), ElementWiseOperation::kPROD);
assert(scale);
return scale;
}
ILayer* ghostModule(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap, int inp,
int oup, int kernel_size = 1, int ratio = 2, int dw_size = 3, int stride = 1, bool relu = true,
std::string lname = "") {
int init_channels = std::ceil(oup / ratio);
int new_channels = init_channels * (ratio - 1);
// Primary convolution
IConvolutionLayer* primary_conv = network->addConvolutionNd(input, init_channels, DimsHW{kernel_size, kernel_size},
weightMap[lname + ".primary_conv.0.weight"], Weights{});
primary_conv->setStrideNd(DimsHW{stride, stride});
primary_conv->setPaddingNd(DimsHW{kernel_size / 2, kernel_size / 2});
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *primary_conv->getOutput(0), lname + ".primary_conv.1", 1e-5);
// Cheap operation (Depthwise Convolution)
IConvolutionLayer* cheap_conv =
network->addConvolutionNd(*bn1->getOutput(0), new_channels, DimsHW{dw_size, dw_size},
weightMap[lname + ".cheap_operation.0.weight"], Weights{});
cheap_conv->setStrideNd(DimsHW{1, 1});
cheap_conv->setPaddingNd(DimsHW{dw_size / 2, dw_size / 2});
cheap_conv->setNbGroups(init_channels);
IScaleLayer* bn2 =
addBatchNorm2d(network, weightMap, *cheap_conv->getOutput(0), lname + ".cheap_operation.1", 1e-5);
// Define relu1 and relu2
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
// Initialize inputs array based on the `relu` flag
std::vector<ITensor*> inputs_vec;
if (relu) {
inputs_vec = {relu1->getOutput(0), relu2->getOutput(0)};
} else {
inputs_vec = {bn1->getOutput(0), bn2->getOutput(0)};
}
ITensor* inputs[] = {inputs_vec[0], inputs_vec[1]};
IConcatenationLayer* concat = network->addConcatenation(inputs, 2);
std::cout << printTensorShape(concat->getOutput(0)) << std::endl;
// Slice the output to keep only the first `oup` channels
Dims start{4, {0, 0, 0, 0}}; // Starting from batch=0, channel=0, height=0, width=0
Dims size{4,
{concat->getOutput(0)->getDimensions().d[0], oup, concat->getOutput(0)->getDimensions().d[2],
concat->getOutput(0)
->getDimensions()
.d[3]}}; // Keep all batches, first `oup` channels, all heights and widths
Dims stride_{4, {1, 1, 1, 1}}; // Stride is 1 for all dimensions
ISliceLayer* slice = network->addSlice(*concat->getOutput(0), start, size, stride_);
cout << "slice" << printTensorShape(slice->getOutput(0)) << endl;
return slice;
}
ILayer* ghostBottleneck(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap,
int in_chs, int mid_chs, int out_chs, int dw_kernel_size = 3, int stride = 1,
float se_ratio = 0.0f, std::string lname = "") {
ILayer* ghost1 = ghostModule(network, input, weightMap, in_chs, mid_chs, 1, 2, 3, 1, true, lname + ".ghost1");
ILayer* depthwise_conv = ghost1;
if (stride > 1) {
IConvolutionLayer* conv_dw =
network->addConvolutionNd(*ghost1->getOutput(0), mid_chs, DimsHW{dw_kernel_size, dw_kernel_size},
weightMap[lname + ".conv_dw.weight"], Weights{});
conv_dw->setStrideNd(DimsHW{stride, stride});
conv_dw->setPaddingNd(DimsHW{(dw_kernel_size - 1) / 2, (dw_kernel_size - 1) / 2});
conv_dw->setNbGroups(mid_chs); // Depth-wise convolution
IScaleLayer* bn_dw = addBatchNorm2d(network, weightMap, *conv_dw->getOutput(0), lname + ".bn_dw", 1e-5);
depthwise_conv = bn_dw;
}
ILayer* se_layer = depthwise_conv;
if (se_ratio > 0.0f) {
se_layer = squeezeExcite(network, *depthwise_conv->getOutput(0), weightMap, mid_chs, se_ratio, lname + ".se");
}
ILayer* ghost2 = ghostModule(network, *se_layer->getOutput(0), weightMap, mid_chs, out_chs, 1, 2, 3, 1, false,
lname + ".ghost2");
ILayer* shortcut_layer = nullptr;
if (in_chs == out_chs && stride == 1) {
shortcut_layer = network->addIdentity(input);
} else {
IConvolutionLayer* conv_shortcut_dw =
network->addConvolutionNd(input, in_chs, DimsHW{dw_kernel_size, dw_kernel_size},
weightMap[lname + ".shortcut.0.weight"], Weights{});
conv_shortcut_dw->setStrideNd(DimsHW{stride, stride});
conv_shortcut_dw->setPaddingNd(DimsHW{(dw_kernel_size - 1) / 2, (dw_kernel_size - 1) / 2});
conv_shortcut_dw->setNbGroups(in_chs); // Depth-wise convolution
IScaleLayer* bn_shortcut_dw =
addBatchNorm2d(network, weightMap, *conv_shortcut_dw->getOutput(0), lname + ".shortcut.1", 1e-5);
IConvolutionLayer* conv_shortcut_pw =
network->addConvolutionNd(*bn_shortcut_dw->getOutput(0), out_chs, DimsHW{1, 1},
weightMap[lname + ".shortcut.2.weight"], Weights{});
IScaleLayer* bn_shortcut_pw =
addBatchNorm2d(network, weightMap, *conv_shortcut_pw->getOutput(0), lname + ".shortcut.3", 1e-5);
shortcut_layer = bn_shortcut_pw;
}
IElementWiseLayer* ew_sum =
network->addElementWise(*ghost2->getOutput(0), *shortcut_layer->getOutput(0), ElementWiseOperation::kSUM);
return ew_sum;
}
ICudaEngine* createEngine(IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network =
builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
// Create input tensor of shape {batchSize, 3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims4{batchSize, 3, INPUT_H, INPUT_W});
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../ghostnetv1.weights");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// Conv Stem
IActivationLayer* conv_stem = convBnReluStem(network, weightMap, *data, 16, "conv_stem");
ILayer* current_layer = conv_stem;
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 16, 16, 16, 3, 1, 0, "blocks.0.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 16, 48, 24, 3, 2, 0, "blocks.1.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 24, 72, 24, 3, 1, 0, "blocks.2.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 24, 72, 40, 5, 2, 0.25, "blocks.3.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 40, 120, 40, 5, 1, 0.25, "blocks.4.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 40, 240, 80, 3, 2, 0, "blocks.5.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 200, 80, 3, 1, 0, "blocks.6.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 184, 80, 3, 1, 0, "blocks.6.1");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 184, 80, 3, 1, 0, "blocks.6.2");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 480, 112, 3, 1, 0.25, "blocks.6.3");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 112, 672, 112, 3, 1, 0.25, "blocks.6.4");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 112, 672, 160, 5, 2, 0.25, "blocks.7.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0, "blocks.8.0");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.25, "blocks.8.1");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0, "blocks.8.2");
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.25, "blocks.8.3");
// Apply ConvBnAct
current_layer = convBnAct(network, weightMap, *current_layer->getOutput(0), 960, "blocks.9.0");
// Global Average Pooling
IReduceLayer* global_pool =
network->addReduce(*current_layer->getOutput(0), ReduceOperation::kAVG, 1 << 2 | 1 << 3, true);
assert(global_pool);
// Conv Head
IConvolutionLayer* conv_head = network->addConvolutionNd(
*global_pool->getOutput(0), 1280, DimsHW{1, 1}, weightMap["conv_head.weight"], weightMap["conv_head.bias"]);
IActivationLayer* act2 = network->addActivation(*conv_head->getOutput(0), ActivationType::kRELU);
// Fully Connected Layer (Classifier)
IFullyConnectedLayer* classifier = network->addFullyConnected(
*act2->getOutput(0), 1000, weightMap["classifier.weight"], weightMap["classifier.bias"]);
classifier->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*classifier->getOutput(0));
// Build engine
config->setMaxWorkspaceSize(1 << 24);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
}
void APIToModel(IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
config->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
const ICudaEngine& engine = context.getEngine();
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Pointers to input and output device buffers to pass to engine.
void* buffers[2];
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float),
cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv) {
if (argc != 2) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./ghostnetv1 -s // serialize model to plan file" << std::endl;
std::cerr << "./ghostnetv1 -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// create a model using the API directly and serialize it to a stream
char* trtModelStream{nullptr};
size_t size{0};
if (std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(&modelStream);
assert(modelStream != nullptr);
std::ofstream p("ghostnetv1.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (std::string(argv[1]) == "-d") {
std::ifstream file("ghostnetv1.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
return -1;
}
float* data = new float[batchSize * 3 * INPUT_H * INPUT_W];
for (int i = 0; i < batchSize * 3 * INPUT_H * INPUT_W; i++)
data[i] = 10.0;
float* prob = new float[batchSize * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
doInference(*context, data, prob, batchSize);
std::cout << "\nOutput:\n\n";
for (int i = 0; i < batchSize; i++) {
std::cout << "Batch " << i << ":\n";
for (unsigned int j = 0; j < OUTPUT_SIZE; j++) {
std::cout << prob[i * OUTPUT_SIZE + j] << ", ";
if (j % 10 == 0)
std::cout << j / 10 << std::endl;
}
std::cout << "\n";
}
context->destroy();
engine->destroy();
runtime->destroy();
delete[] data;
delete[] prob;
return 0;
}

455
ghostnet/ghostnetv1/logging.h Normal file
View File

@ -0,0 +1,455 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
#include "NvInferRuntimeCommon.h"
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf {
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) : mOutput(other.mOutput) {}
~LogStreamConsumerBuffer() {
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr()) {
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync() {
putOutput();
return 0;
}
void putOutput() {
if (mShouldLog) {
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase {
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog) {}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(severity <= reportableSeverity),
mSeverity(severity) {}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(other.mShouldLog),
mSeverity(other.mSeverity) {}
void setReportableSeverity(Severity reportableSeverity) {
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger {
public:
Logger(Severity severity = Severity::kWARNING) : mReportableSeverity(severity) {}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult {
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger() { return *this; }
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) noexcept override {
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity) { mReportableSeverity = severity; }
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom {
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started), mName(name), mCmdline(cmdline) {}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline) {
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) {
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom) {
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result) {
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass) {
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const { return mReportableSeverity; }
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result) {
switch (result) {
case TestResult::kRUNNING:
return "RUNNING";
case TestResult::kPASSED:
return "PASSED";
case TestResult::kFAILED:
return "FAILED";
case TestResult::kWAIVED:
return "WAIVED";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result) {
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv) {
std::stringstream ss;
for (int i = 0; i < argc; i++) {
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace {
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

24
ghostnet/ghostnetv2/CMakeLists.txt Normal file
View File

@ -0,0 +1,24 @@
cmake_minimum_required(VERSION 2.6)
project(ghostnetv2)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
add_executable(ghostnetv2 ${PROJECT_SOURCE_DIR}/ghostnetv2.cpp)
target_link_libraries(ghostnetv2 nvinfer)
target_link_libraries(ghostnetv2 cudart)
add_definitions(-O2 -pthread)

312
ghostnet/ghostnetv2/gen_wts.py Normal file
View File

@ -0,0 +1,312 @@
import torch
import torch.nn as nn
import torch.onnx
import struct
import torch
import torch.nn.functional as F
import math
from timm.models.registry import register_model
def _make_divisible(v, divisor, min_value=None):
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
def hard_sigmoid(x, inplace: bool = False):
if inplace:
return x.add_(3.).clamp_(0., 6.).div_(6.)
else:
return F.relu6(x + 3.) / 6.
class SqueezeExcite(nn.Module):
def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None,
act_layer=nn.ReLU, gate_fn=hard_sigmoid, divisor=4, **_):
super(SqueezeExcite, self).__init__()
self.gate_fn = gate_fn
reduced_chs = _make_divisible((reduced_base_chs or in_chs) * se_ratio, divisor)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)
self.act1 = act_layer(inplace=True)
self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)
def forward(self, x):
x_se = self.avg_pool(x)
x_se = self.conv_reduce(x_se)
x_se = self.act1(x_se)
x_se = self.conv_expand(x_se)
x = x * self.gate_fn(x_se)
return x
class ConvBnAct(nn.Module):
def __init__(self, in_chs, out_chs, kernel_size,
stride=1, act_layer=nn.ReLU):
super(ConvBnAct, self).__init__()
self.conv = nn.Conv2d(in_chs, out_chs, kernel_size, stride, kernel_size//2, bias=False)
self.bn1 = nn.BatchNorm2d(out_chs)
self.act1 = act_layer(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn1(x)
x = self.act1(x)
return x
class GhostModuleV2(nn.Module):
def __init__(self, inp, oup, kernel_size=1, ratio=2, dw_size=3, stride=1, relu=True, mode=None, args=None):
super(GhostModuleV2, self).__init__()
self.mode = mode
self.gate_fn = nn.Sigmoid()
if self.mode in ['original']:
self.oup = oup
init_channels = math.ceil(oup / ratio)
new_channels = init_channels*(ratio-1)
self.primary_conv = nn.Sequential(
nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size//2, bias=False),
nn.BatchNorm2d(init_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
self.cheap_operation = nn.Sequential(
nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size//2, groups=init_channels, bias=False),
nn.BatchNorm2d(new_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
elif self.mode in ['attn']:
self.oup = oup
init_channels = math.ceil(oup / ratio)
new_channels = init_channels*(ratio-1)
self.primary_conv = nn.Sequential(
nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size//2, bias=False),
nn.BatchNorm2d(init_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
self.cheap_operation = nn.Sequential(
nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size//2, groups=init_channels, bias=False),
nn.BatchNorm2d(new_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
self.short_conv = nn.Sequential(
nn.Conv2d(inp, oup, kernel_size, stride, kernel_size//2, bias=False),
nn.BatchNorm2d(oup),
nn.Conv2d(oup, oup, kernel_size=(1, 5), stride=1, padding=(0, 2), groups=oup, bias=False),
nn.BatchNorm2d(oup),
nn.Conv2d(oup, oup, kernel_size=(5, 1), stride=1, padding=(2, 0), groups=oup, bias=False),
nn.BatchNorm2d(oup),
)
def forward(self, x):
if self.mode in ['original']:
x1 = self.primary_conv(x)
x2 = self.cheap_operation(x1)
out = torch.cat([x1, x2], dim=1)
return out[:, :self.oup, :, :]
elif self.mode in ['attn']:
res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2))
x1 = self.primary_conv(x)
x2 = self.cheap_operation(x1)
out = torch.cat([x1, x2], dim=1)
return out[:, :self.oup, :, :]*F.interpolate(self.gate_fn(res),
size=(out.shape[-2], out.shape[-1]), mode='nearest')
class GhostBottleneckV2(nn.Module):
def __init__(self, in_chs, mid_chs, out_chs, dw_kernel_size=3,
stride=1, act_layer=nn.ReLU, se_ratio=0., layer_id=None, args=None):
super(GhostBottleneckV2, self).__init__()
has_se = se_ratio is not None and se_ratio > 0.
self.stride = stride
# Point-wise expansion
if layer_id <= 1:
self.ghost1 = GhostModuleV2(in_chs, mid_chs, relu=True, mode='original', args=args)
else:
self.ghost1 = GhostModuleV2(in_chs, mid_chs, relu=True, mode='attn', args=args)
# Depth-wise convolution
if self.stride > 1:
self.conv_dw = nn.Conv2d(mid_chs, mid_chs, dw_kernel_size, stride=stride,
padding=(dw_kernel_size-1)//2, groups=mid_chs, bias=False)
self.bn_dw = nn.BatchNorm2d(mid_chs)
# Squeeze-and-excitation
if has_se:
self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio)
else:
self.se = None
self.ghost2 = GhostModuleV2(mid_chs, out_chs, relu=False, mode='original', args=args)
# shortcut
if (in_chs == out_chs and self.stride == 1):
self.shortcut = nn.Sequential()
else:
self.shortcut = nn.Sequential(
nn.Conv2d(in_chs, in_chs, dw_kernel_size, stride=stride,
padding=(dw_kernel_size-1)//2, groups=in_chs, bias=False),
nn.BatchNorm2d(in_chs),
nn.Conv2d(in_chs, out_chs, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(out_chs),
)
def forward(self, x):
residual = x
x = self.ghost1(x)
if self.stride > 1:
x = self.conv_dw(x)
x = self.bn_dw(x)
if self.se is not None:
x = self.se(x)
x = self.ghost2(x)
x += self.shortcut(residual)
return x
class GhostNetV2(nn.Module):
def __init__(self, cfgs, num_classes=1000, width=1.0, dropout=0.2, block=GhostBottleneckV2, args=None):
super(GhostNetV2, self).__init__()
self.cfgs = cfgs
self.dropout = dropout
# building first layer
output_channel = _make_divisible(16 * width, 4)
self.conv_stem = nn.Conv2d(3, output_channel, 3, 2, 1, bias=False)
self.bn1 = nn.BatchNorm2d(output_channel)
self.act1 = nn.ReLU(inplace=True)
input_channel = output_channel
# building inverted residual blocks
stages = []
layer_id = 0
for cfg in self.cfgs:
layers = []
for k, exp_size, c, se_ratio, s in cfg:
output_channel = _make_divisible(c * width, 4)
hidden_channel = _make_divisible(exp_size * width, 4)
layers.append(block(input_channel, hidden_channel, output_channel, k, s,
se_ratio=se_ratio, layer_id=layer_id, args=args))
input_channel = output_channel
layer_id += 1
stages.append(nn.Sequential(*layers))
output_channel = _make_divisible(exp_size * width, 4)
stages.append(nn.Sequential(ConvBnAct(input_channel, output_channel, 1)))
input_channel = output_channel
self.blocks = nn.Sequential(*stages)
# building last several layers
output_channel = 1280
self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
self.conv_head = nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=True)
self.act2 = nn.ReLU(inplace=True)
self.classifier = nn.Linear(output_channel, num_classes)
def forward(self, x):
x = self.conv_stem(x)
x = self.bn1(x)
x = self.act1(x)
x = self.blocks(x)
x = self.global_pool(x)
x = self.conv_head(x)
x = self.act2(x)
x = x.view(x.size(0), -1)
if self.dropout > 0.:
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.classifier(x)
return x
@register_model
def ghostnetv2(**kwargs):
cfgs = [
# k, t, c, SE, s
[[3, 16, 16, 0, 1]],
[[3, 48, 24, 0, 2]],
[[3, 72, 24, 0, 1]],
[[5, 72, 40, 0.25, 2]],
[[5, 120, 40, 0.25, 1]],
[[3, 240, 80, 0, 2]],
[[3, 200, 80, 0, 1],
[3, 184, 80, 0, 1],
[3, 184, 80, 0, 1],
[3, 480, 112, 0.25, 1],
[3, 672, 112, 0.25, 1]],
[[5, 672, 160, 0.25, 2]],
[[5, 960, 160, 0, 1],
[5, 960, 160, 0.25, 1],
[5, 960, 160, 0, 1],
[5, 960, 160, 0.25, 1]]
]
return GhostNetV2(cfgs, num_classes=kwargs['num_classes'],
width=kwargs['width'],
dropout=kwargs['dropout'],
args=kwargs['args'])
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
# Function to export weights in the specified format
def export_weight(model):
f = open("ghostnetv2.weights", 'w')
f.write("{}\n".format(len(model.state_dict().keys())))
# Convert weights to hexadecimal format
for k, v in model.state_dict().items():
print('exporting ... {}: {}'.format(k, v.shape))
# Reshape the weights to 1D
vr = v.reshape(-1).cpu().numpy()
f.write("{} {}".format(k, len(vr)))
for vv in vr:
f.write(" ")
f.write(struct.pack(">f", float(vv)).hex())
f.write("\n")
f.close()
# Function to evaluate the model (optional)
def eval_model(input, model):
output = model(input)
print("------from inference------")
print(input)
print(output)
if __name__ == "__main__":
setup_seed(1)
# Create an instance of GhostNetV2
model = ghostnetv2(width=1.0, num_classes=1000, dropout=0.2, args=None)
model.eval()
# Dummy input tensor (adjust the shape as per your requirement)
input = torch.full((32, 3, 320, 256), 10.0)
# Export the model weights
export_weight(model)
# Evaluate the model
eval_model(input, model)

591
ghostnet/ghostnetv2/ghostnetv2.cpp Normal file
View File

@ -0,0 +1,591 @@
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
using namespace std;
#define CHECK(status) \
do { \
auto ret = (status); \
if (ret != 0) { \
std::cerr << "Cuda failure: " << ret << std::endl; \
abort(); \
} \
} while (0)
// Define input/output parameters
static const int INPUT_H = 256;
static const int INPUT_W = 320;
static const int OUTPUT_SIZE = 1000;
static const int batchSize = 32;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
using namespace nvinfer1;
static Logger gLogger;
// Load weight file
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open the weight file
std::ifstream input(file);
if (!input.is_open()) {
std::cerr << "Unable to load weight file." << std::endl;
exit(EXIT_FAILURE);
}
// Read the number of weights
int32_t count;
input >> count;
if (count <= 0) {
std::cerr << "Invalid weight map file." << std::endl;
exit(EXIT_FAILURE);
}
while (count--) {
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read the name and size
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load weight data
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(uint32_t) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
int _make_divisible(int v, int divisor, int min_value = -1) {
// If min_value is not specified, set it to divisor
if (min_value == -1) {
min_value = divisor;
}
// Calculate new channel size to be divisible by divisor
int new_v = std::max(min_value, (v + divisor / 2) / divisor * divisor);
// Ensure rounding down does not reduce by more than 10%
if (new_v < static_cast<int>(0.9 * v)) {
new_v += divisor;
}
return new_v;
}
ILayer* hardSigmoid(INetworkDefinition* network, ITensor& input) {
// Apply Hard Sigmoid activation function
IActivationLayer* scale_layer = network->addActivation(input, ActivationType::kHARD_SIGMOID);
// Return the output after activation
return scale_layer;
}
IScaleLayer* addBatchNorm2d(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
std::string lname, float eps) {
float* gamma = (float*)weightMap[lname + ".weight"].values;
float* beta = (float*)weightMap[lname + ".bias"].values;
float* mean = (float*)weightMap[lname + ".running_mean"].values;
float* var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
float* scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float* shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float* pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IActivationLayer* convBnReluStem(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int outch, std::string lname) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// Step 1: Convolution layer
IConvolutionLayer* conv1 =
network->addConvolutionNd(input, outch, DimsHW{3, 3}, weightMap[lname + ".weight"], emptywts);
assert(conv1);
conv1->setStrideNd(DimsHW{2, 2}); // Stride of 2
conv1->setPaddingNd(DimsHW{1, 1}); // Padding of 1
// Step 2: Batch normalization layer
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "bn1", 1e-5);
// Step 3: ReLU activation
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
return relu1; // Return the result after activation
}
ILayer* convBnAct(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
int out_channels, std::string lname, ActivationType actType = ActivationType::kRELU) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// Add convolution layer
IConvolutionLayer* conv =
network->addConvolutionNd(input, out_channels, DimsHW{1, 1}, weightMap[lname + ".conv.weight"], emptywts);
assert(conv);
conv->setStrideNd(DimsHW{1, 1});
// Add batch normalization layer
IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv->getOutput(0), lname + ".bn1", 1e-5);
// Add activation layer (default is ReLU)
IActivationLayer* act = network->addActivation(*bn->getOutput(0), actType);
assert(act);
return act;
}
ILayer* squeezeExcite(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap,
int in_chs, float se_ratio = 0.25, std::string lname = "", float eps = 1e-5) {
// Step 1: Global average pooling
IReduceLayer* avg_pool = network->addReduce(input, ReduceOperation::kAVG, 1 << 2 | 1 << 3, true);
assert(avg_pool);
// Step 2: 1x1 convolution for dimension reduction
int reduced_chs = _make_divisible(static_cast<int>(in_chs * se_ratio), 4);
IConvolutionLayer* conv_reduce =
network->addConvolutionNd(*avg_pool->getOutput(0), reduced_chs, DimsHW{1, 1},
weightMap[lname + ".conv_reduce.weight"], weightMap[lname + ".conv_reduce.bias"]);
assert(conv_reduce);
// Step 3: ReLU activation
IActivationLayer* relu1 = network->addActivation(*conv_reduce->getOutput(0), ActivationType::kRELU);
assert(relu1);
// Step 4: 1x1 convolution for dimension expansion
IConvolutionLayer* conv_expand =
network->addConvolutionNd(*relu1->getOutput(0), in_chs, DimsHW{1, 1},
weightMap[lname + ".conv_expand.weight"], weightMap[lname + ".conv_expand.bias"]);
assert(conv_expand);
// Step 5: Hard Sigmoid activation
ILayer* hard_sigmoid = hardSigmoid(network, *conv_expand->getOutput(0));
// Step 6: Multiply input by the output of SE module
IElementWiseLayer* scale = network->addElementWise(input, *hard_sigmoid->getOutput(0), ElementWiseOperation::kPROD);
assert(scale);
return scale;
}
ILayer* ghostModuleV2(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap, int inp,
int oup, int kernel_size = 1, int ratio = 2, int dw_size = 3, int stride = 1, bool relu = true,
std::string lname = "", std::string mode = "original") {
int init_channels = std::ceil(oup / ratio);
int new_channels = init_channels * (ratio - 1);
// Primary convolution
IConvolutionLayer* primary_conv = network->addConvolutionNd(input, init_channels, DimsHW{kernel_size, kernel_size},
weightMap[lname + ".primary_conv.0.weight"], Weights{});
primary_conv->setStrideNd(DimsHW{stride, stride});
primary_conv->setPaddingNd(DimsHW{kernel_size / 2, kernel_size / 2});
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *primary_conv->getOutput(0), lname + ".primary_conv.1", 1e-5);
ITensor* act1_output = bn1->getOutput(0);
if (relu) {
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
act1_output = relu1->getOutput(0);
}
// Cheap operation
IConvolutionLayer* cheap_conv =
network->addConvolutionNd(*act1_output, new_channels, DimsHW{dw_size, dw_size},
weightMap[lname + ".cheap_operation.0.weight"], Weights{});
cheap_conv->setStrideNd(DimsHW{1, 1});
cheap_conv->setPaddingNd(DimsHW{dw_size / 2, dw_size / 2});
cheap_conv->setNbGroups(init_channels);
IScaleLayer* bn2 =
addBatchNorm2d(network, weightMap, *cheap_conv->getOutput(0), lname + ".cheap_operation.1", 1e-5);
ITensor* act2_output = bn2->getOutput(0);
if (relu) {
IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
act2_output = relu2->getOutput(0);
}
// Concatenate
ITensor* concat_inputs[] = {act1_output, act2_output};
IConcatenationLayer* concat = network->addConcatenation(concat_inputs, 2);
// Slice to oup channels
Dims start{4, {0, 0, 0, 0}};
Dims size = concat->getOutput(0)->getDimensions();
size.d[1] = oup;
Dims stride_{4, {1, 1, 1, 1}};
ISliceLayer* slice = network->addSlice(*concat->getOutput(0), start, size, stride_);
ITensor* out = slice->getOutput(0);
if (mode == "original") {
return slice;
} else if (mode == "attn") {
// Attention mechanism
// Average pooling
IPoolingLayer* avg_pool = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{2, 2});
avg_pool->setStrideNd(DimsHW{2, 2});
ITensor* avg_pooled = avg_pool->getOutput(0);
// Short convolution branch
IConvolutionLayer* short_conv1 =
network->addConvolutionNd(*avg_pooled, oup, DimsHW{kernel_size, kernel_size},
weightMap[lname + ".short_conv.0.weight"], Weights{});
short_conv1->setStrideNd(DimsHW{1, 1});
short_conv1->setPaddingNd(DimsHW{kernel_size / 2, kernel_size / 2});
IScaleLayer* short_bn1 =
addBatchNorm2d(network, weightMap, *short_conv1->getOutput(0), lname + ".short_conv.1", 1e-5);
// Conv with kernel size (1,5)
IConvolutionLayer* short_conv2 = network->addConvolutionNd(
*short_bn1->getOutput(0), oup, DimsHW{1, 5}, weightMap[lname + ".short_conv.2.weight"], Weights{});
short_conv2->setStrideNd(DimsHW{1, 1});
short_conv2->setPaddingNd(DimsHW{0, 2});
short_conv2->setNbGroups(oup);
IScaleLayer* short_bn2 =
addBatchNorm2d(network, weightMap, *short_conv2->getOutput(0), lname + ".short_conv.3", 1e-5);
// Conv with kernel size (5,1)
IConvolutionLayer* short_conv3 = network->addConvolutionNd(
*short_bn2->getOutput(0), oup, DimsHW{5, 1}, weightMap[lname + ".short_conv.4.weight"], Weights{});
short_conv3->setStrideNd(DimsHW{1, 1});
short_conv3->setPaddingNd(DimsHW{2, 0});
short_conv3->setNbGroups(oup);
IScaleLayer* short_bn3 =
addBatchNorm2d(network, weightMap, *short_conv3->getOutput(0), lname + ".short_conv.5", 1e-5);
ITensor* res = short_bn3->getOutput(0);
// Sigmoid activation
IActivationLayer* gate = network->addActivation(*res, ActivationType::kSIGMOID);
// Upsample to the same size as out
IResizeLayer* gate_upsampled = network->addResize(*gate->getOutput(0));
gate_upsampled->setResizeMode(ResizeMode::kNEAREST);
Dims out_dims = out->getDimensions();
gate_upsampled->setOutputDimensions(out_dims);
// Element-wise multiplication
IElementWiseLayer* scaled_out =
network->addElementWise(*out, *gate_upsampled->getOutput(0), ElementWiseOperation::kPROD);
return scaled_out;
} else {
std::cerr << "Invalid mode: " << mode << " in ghostModuleV2" << std::endl;
return nullptr;
}
}
ILayer* ghostBottleneck(INetworkDefinition* network, ITensor& input, std::map<std::string, Weights>& weightMap,
int in_chs, int mid_chs, int out_chs, int dw_kernel_size = 3, int stride = 1,
float se_ratio = 0.0f, std::string lname = "", int layer_id = 0) {
// Determine mode based on layer_id
std::string mode = (layer_id <= 1) ? "original" : "attn";
// ghost1
ILayer* ghost1 =
ghostModuleV2(network, input, weightMap, in_chs, mid_chs, 1, 2, 3, 1, true, lname + ".ghost1", mode);
ILayer* depthwise_conv = ghost1;
if (stride > 1) {
IConvolutionLayer* conv_dw =
network->addConvolutionNd(*ghost1->getOutput(0), mid_chs, DimsHW{dw_kernel_size, dw_kernel_size},
weightMap[lname + ".conv_dw.weight"], Weights{});
conv_dw->setStrideNd(DimsHW{stride, stride});
conv_dw->setPaddingNd(DimsHW{(dw_kernel_size - 1) / 2, (dw_kernel_size - 1) / 2});
conv_dw->setNbGroups(mid_chs);
IScaleLayer* bn_dw = addBatchNorm2d(network, weightMap, *conv_dw->getOutput(0), lname + ".bn_dw", 1e-5);
depthwise_conv = bn_dw;
}
ILayer* se_layer = depthwise_conv;
if (se_ratio > 0.0f) {
se_layer = squeezeExcite(network, *depthwise_conv->getOutput(0), weightMap, mid_chs, se_ratio, lname + ".se");
}
// ghost2 uses original mode
ILayer* ghost2 = ghostModuleV2(network, *se_layer->getOutput(0), weightMap, mid_chs, out_chs, 1, 2, 3, 1, false,
lname + ".ghost2", "original");
ILayer* shortcut_layer = nullptr;
if (in_chs == out_chs && stride == 1) {
shortcut_layer = network->addIdentity(input);
} else {
IConvolutionLayer* conv_shortcut_dw =
network->addConvolutionNd(input, in_chs, DimsHW{dw_kernel_size, dw_kernel_size},
weightMap[lname + ".shortcut.0.weight"], Weights{});
conv_shortcut_dw->setStrideNd(DimsHW{stride, stride});
conv_shortcut_dw->setPaddingNd(DimsHW{(dw_kernel_size - 1) / 2, (dw_kernel_size - 1) / 2});
conv_shortcut_dw->setNbGroups(in_chs);
IScaleLayer* bn_shortcut_dw =
addBatchNorm2d(network, weightMap, *conv_shortcut_dw->getOutput(0), lname + ".shortcut.1", 1e-5);
IConvolutionLayer* conv_shortcut_pw =
network->addConvolutionNd(*bn_shortcut_dw->getOutput(0), out_chs, DimsHW{1, 1},
weightMap[lname + ".shortcut.2.weight"], Weights{});
IScaleLayer* bn_shortcut_pw =
addBatchNorm2d(network, weightMap, *conv_shortcut_pw->getOutput(0), lname + ".shortcut.3", 1e-5);
shortcut_layer = bn_shortcut_pw;
}
IElementWiseLayer* ew_sum =
network->addElementWise(*ghost2->getOutput(0), *shortcut_layer->getOutput(0), ElementWiseOperation::kSUM);
return ew_sum;
}
ICudaEngine* createEngine(IBuilder* builder, IBuilderConfig* config, DataType dt) {
// Use explicit batch mode
INetworkDefinition* network =
builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
// Create input tensor
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims4{batchSize, 3, INPUT_H, INPUT_W});
assert(data);
// Load weights
std::map<std::string, Weights> weightMap = loadWeights("../ghostnetv2.weights");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// Step 1: Conv Stem
IActivationLayer* conv_stem = convBnReluStem(network, weightMap, *data, 16, "conv_stem");
ILayer* current_layer = conv_stem;
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 16, 16, 16, 3, 1, 0.0f, "blocks.0.0", 0);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 16, 48, 24, 3, 2, 0.0f, "blocks.1.0", 1);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 24, 72, 24, 3, 1, 0.0f, "blocks.2.0", 2);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 24, 72, 40, 5, 2, 0.25f, "blocks.3.0", 3);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 40, 120, 40, 5, 1, 0.25f,
"blocks.4.0", 4);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 40, 240, 80, 3, 2, 0.0f, "blocks.5.0", 5);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 200, 80, 3, 1, 0.0f, "blocks.6.0", 6);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 184, 80, 3, 1, 0.0f, "blocks.6.1", 7);
current_layer =
ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 184, 80, 3, 1, 0.0f, "blocks.6.2", 8);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 80, 480, 112, 3, 1, 0.25f,
"blocks.6.3", 9);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 112, 672, 112, 3, 1, 0.25f,
"blocks.6.4", 10);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 112, 672, 160, 5, 2, 0.25f,
"blocks.7.0", 11);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.0f,
"blocks.8.0", 12);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.25f,
"blocks.8.1", 13);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.0f,
"blocks.8.2", 14);
current_layer = ghostBottleneck(network, *current_layer->getOutput(0), weightMap, 160, 960, 160, 5, 1, 0.25f,
"blocks.8.3", 15);
// Apply ConvBnAct
current_layer = convBnAct(network, weightMap, *current_layer->getOutput(0), 960, "blocks.9.0");
// Global average pooling
IReduceLayer* global_pool =
network->addReduce(*current_layer->getOutput(0), ReduceOperation::kAVG, 1 << 2 | 1 << 3, true);
assert(global_pool);
// Conv Head
IConvolutionLayer* conv_head = network->addConvolutionNd(
*global_pool->getOutput(0), 1280, DimsHW{1, 1}, weightMap["conv_head.weight"], weightMap["conv_head.bias"]);
IActivationLayer* act2 = network->addActivation(*conv_head->getOutput(0), ActivationType::kRELU);
// Fully connected layer (classifier)
IFullyConnectedLayer* classifier = network->addFullyConnected(
*act2->getOutput(0), 1000, weightMap["classifier.weight"], weightMap["classifier.bias"]);
classifier->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*classifier->getOutput(0));
// Build the engine
config->setMaxWorkspaceSize(1 << 24);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// Destroy the network
network->destroy();
// Free memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
}
void APIToModel(IHostMemory** modelStream) {
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model and serialize
ICudaEngine* engine = createEngine(builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Release resources
engine->destroy();
config->destroy();
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
const ICudaEngine& engine = context.getEngine();
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Input and output buffers
void* buffers[2];
// Create buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// Copy input data to device, execute inference, and copy output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float),
cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv) {
if (argc != 2) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./ghostnetv2 -s // serialize model to plan file" << std::endl;
std::cerr << "./ghostnetv2 -d // deserialize plan file and run inference" << std::endl;
return -1;
}
// Create model and serialize
char* trtModelStream{nullptr};
size_t size{0};
if (std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(&modelStream);
assert(modelStream != nullptr);
std::ofstream p("ghostnetv2.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
} else if (std::string(argv[1]) == "-d") {
std::ifstream file("ghostnetv2.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
} else {
return -1;
}
// Allocate input and output data
float* data = new float[batchSize * 3 * INPUT_H * INPUT_W];
for (int i = 0; i < batchSize * 3 * INPUT_H * INPUT_W; i++)
data[i] = 10.0;
float* prob = new float[batchSize * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Execute inference
doInference(*context, data, prob, batchSize);
// Print output results
std::cout << "\nOutput:\n\n";
for (int i = 0; i < batchSize; i++) {
std::cout << "Batch " << i << ":\n";
for (unsigned int j = 0; j < OUTPUT_SIZE; j++) {
std::cout << prob[i * OUTPUT_SIZE + j] << ", ";
if (j % 10 == 0)
std::cout << j / 10 << std::endl;
}
std::cout << "\n";
}
// Release resources
context->destroy();
engine->destroy();
runtime->destroy();
delete[] data;
delete[] prob;
return 0;
}

455
ghostnet/ghostnetv2/logging.h Normal file
View File

@ -0,0 +1,455 @@
/*
* 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 TENSORRT_LOGGING_H
#define TENSORRT_LOGGING_H
#include <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
#include "NvInferRuntimeCommon.h"
using Severity = nvinfer1::ILogger::Severity;
class LogStreamConsumerBuffer : public std::stringbuf {
public:
LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) : mOutput(other.mOutput) {}
~LogStreamConsumerBuffer() {
// std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
// std::streambuf::pptr() gives a pointer to the current position of the output sequence
// if the pointer to the beginning is not equal to the pointer to the current position,
// call putOutput() to log the output to the stream
if (pbase() != pptr()) {
putOutput();
}
}
// synchronizes the stream buffer and returns 0 on success
// synchronizing the stream buffer consists of inserting the buffer contents into the stream,
// resetting the buffer and flushing the stream
virtual int sync() {
putOutput();
return 0;
}
void putOutput() {
if (mShouldLog) {
// prepend timestamp
std::time_t timestamp = std::time(nullptr);
tm* tm_local = std::localtime(&timestamp);
std::cout << "[";
std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
// std::stringbuf::str() gets the string contents of the buffer
// insert the buffer contents pre-appended by the appropriate prefix into the stream
mOutput << mPrefix << str();
// set the buffer to empty
str("");
// flush the stream
mOutput.flush();
}
}
void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
private:
std::ostream& mOutput;
std::string mPrefix;
bool mShouldLog;
};
//!
//! \class LogStreamConsumerBase
//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
//!
class LogStreamConsumerBase {
public:
LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
: mBuffer(stream, prefix, shouldLog) {}
protected:
LogStreamConsumerBuffer mBuffer;
};
//!
//! \class LogStreamConsumer
//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
//! Order of base classes is LogStreamConsumerBase and then std::ostream.
//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
//! Please do not change the order of the parent classes.
//!
class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
public:
//! \brief Creates a LogStreamConsumer which logs messages with level severity.
//! Reportable severity determines if the messages are severe enough to be logged.
LogStreamConsumer(Severity reportableSeverity, Severity severity)
: LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(severity <= reportableSeverity),
mSeverity(severity) {}
LogStreamConsumer(LogStreamConsumer&& other)
: LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog),
std::ostream(&mBuffer) // links the stream buffer with the stream
,
mShouldLog(other.mShouldLog),
mSeverity(other.mSeverity) {}
void setReportableSeverity(Severity reportableSeverity) {
mShouldLog = mSeverity <= reportableSeverity;
mBuffer.setShouldLog(mShouldLog);
}
private:
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
static std::string severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
bool mShouldLog;
Severity mSeverity;
};
//! \class Logger
//!
//! \brief Class which manages logging of TensorRT tools and samples
//!
//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
//! and supports logging two types of messages:
//!
//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
//! - Test pass/fail messages
//!
//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
//!
//! In the future, this class could be extended to support dumping test results to a file in some standard format
//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
//!
//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
//! library and messages coming from the sample.
//!
//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
//! object.
class Logger : public nvinfer1::ILogger {
public:
Logger(Severity severity = Severity::kWARNING) : mReportableSeverity(severity) {}
//!
//! \enum TestResult
//! \brief Represents the state of a given test
//!
enum class TestResult {
kRUNNING, //!< The test is running
kPASSED, //!< The test passed
kFAILED, //!< The test failed
kWAIVED //!< The test was waived
};
//!
//! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
//! \return The nvinfer1::ILogger associated with this Logger
//!
//! TODO Once all samples are updated to use this method to register the logger with TensorRT,
//! we can eliminate the inheritance of Logger from ILogger
//!
nvinfer1::ILogger& getTRTLogger() { return *this; }
//!
//! \brief Implementation of the nvinfer1::ILogger::log() virtual method
//!
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) noexcept override {
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
//!
//! \brief Method for controlling the verbosity of logging output
//!
//! \param severity The logger will only emit messages that have severity of this level or higher.
//!
void setReportableSeverity(Severity severity) { mReportableSeverity = severity; }
//!
//! \brief Opaque handle that holds logging information for a particular test
//!
//! This object is an opaque handle to information used by the Logger to print test results.
//! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
//! with Logger::reportTest{Start,End}().
//!
class TestAtom {
public:
TestAtom(TestAtom&&) = default;
private:
friend class Logger;
TestAtom(bool started, const std::string& name, const std::string& cmdline)
: mStarted(started), mName(name), mCmdline(cmdline) {}
bool mStarted;
std::string mName;
std::string mCmdline;
};
//!
//! \brief Define a test for logging
//!
//! \param[in] name The name of the test. This should be a string starting with
//! "TensorRT" and containing dot-separated strings containing
//! the characters [A-Za-z0-9_].
//! For example, "TensorRT.sample_googlenet"
//! \param[in] cmdline The command line used to reproduce the test
//
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
//!
static TestAtom defineTest(const std::string& name, const std::string& cmdline) {
return TestAtom(false, name, cmdline);
}
//!
//! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
//! as input
//!
//! \param[in] name The name of the test
//! \param[in] argc The number of command-line arguments
//! \param[in] argv The array of command-line arguments (given as C strings)
//!
//! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) {
auto cmdline = genCmdlineString(argc, argv);
return defineTest(name, cmdline);
}
//!
//! \brief Report that a test has started.
//!
//! \pre reportTestStart() has not been called yet for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has started
//!
static void reportTestStart(TestAtom& testAtom) {
reportTestResult(testAtom, TestResult::kRUNNING);
assert(!testAtom.mStarted);
testAtom.mStarted = true;
}
//!
//! \brief Report that a test has ended.
//!
//! \pre reportTestStart() has been called for the given testAtom
//!
//! \param[in] testAtom The handle to the test that has ended
//! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
//! TestResult::kFAILED, TestResult::kWAIVED
//!
static void reportTestEnd(const TestAtom& testAtom, TestResult result) {
assert(result != TestResult::kRUNNING);
assert(testAtom.mStarted);
reportTestResult(testAtom, result);
}
static int reportPass(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kPASSED);
return EXIT_SUCCESS;
}
static int reportFail(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kFAILED);
return EXIT_FAILURE;
}
static int reportWaive(const TestAtom& testAtom) {
reportTestEnd(testAtom, TestResult::kWAIVED);
return EXIT_SUCCESS;
}
static int reportTest(const TestAtom& testAtom, bool pass) {
return pass ? reportPass(testAtom) : reportFail(testAtom);
}
Severity getReportableSeverity() const { return mReportableSeverity; }
private:
//!
//! \brief returns an appropriate string for prefixing a log message with the given severity
//!
static const char* severityPrefix(Severity severity) {
switch (severity) {
case Severity::kINTERNAL_ERROR:
return "[F] ";
case Severity::kERROR:
return "[E] ";
case Severity::kWARNING:
return "[W] ";
case Severity::kINFO:
return "[I] ";
case Severity::kVERBOSE:
return "[V] ";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate string for prefixing a test result message with the given result
//!
static const char* testResultString(TestResult result) {
switch (result) {
case TestResult::kRUNNING:
return "RUNNING";
case TestResult::kPASSED:
return "PASSED";
case TestResult::kFAILED:
return "FAILED";
case TestResult::kWAIVED:
return "WAIVED";
default:
assert(0);
return "";
}
}
//!
//! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
//!
static std::ostream& severityOstream(Severity severity) {
return severity >= Severity::kINFO ? std::cout : std::cerr;
}
//!
//! \brief method that implements logging test results
//!
static void reportTestResult(const TestAtom& testAtom, TestResult result) {
severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
<< testAtom.mCmdline << std::endl;
}
//!
//! \brief generate a command line string from the given (argc, argv) values
//!
static std::string genCmdlineString(int argc, char const* const* argv) {
std::stringstream ss;
for (int i = 0; i < argc; i++) {
if (i > 0)
ss << " ";
ss << argv[i];
}
return ss.str();
}
Severity mReportableSeverity;
};
namespace {
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
//!
//! Example usage:
//!
//! LOG_VERBOSE(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
//!
//! Example usage:
//!
//! LOG_INFO(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_INFO(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
//!
//! Example usage:
//!
//! LOG_WARN(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_WARN(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
//!
//! Example usage:
//!
//! LOG_ERROR(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_ERROR(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
}
//!
//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
// ("fatal" severity)
//!
//! Example usage:
//!
//! LOG_FATAL(logger) << "hello world" << std::endl;
//!
inline LogStreamConsumer LOG_FATAL(const Logger& logger) {
return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
}
} // anonymous namespace
#endif // TENSORRT_LOGGING_H

25
googlenet/CMakeLists.txt Normal file
View File

@ -0,0 +1,25 @@
cmake_minimum_required(VERSION 2.6)
project(googlenet)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
add_executable(googlenet ${PROJECT_SOURCE_DIR}/googlenet.cpp)
target_link_libraries(googlenet nvinfer)
target_link_libraries(googlenet cudart)
add_definitions(-O2 -pthread)

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