diff --git a/.gitignore b/.gitignore index 3323bab5..a8823a88 100644 --- a/.gitignore +++ b/.gitignore @@ -19,13 +19,6 @@ build # 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/ @@ -92,7 +85,6 @@ ppocrv5/*.bin ppocrv5/*.log ppocrv5/*.graph.json ppocrv5/inference.json -ppocrv5/**/__pycache__/ ppocrv5/**/*.pyc ppocrv5/.venv*/ ppocrv5/**/*_layers.txt @@ -105,3 +97,223 @@ ppocrv5/samples/validation/**/*.jpg ppocrv5/samples/validation/**/*.jpeg ppocrv5/samples/validation/**/*.png !ppocrv5/samples/validation/README.md +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +*.lcov +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi/* +!.pixi/config.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule* +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ +# Temporary file for partial code execution +tempCodeRunnerFile.py + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml diff --git a/alexnet/FindTensorRT.cmake b/alexnet/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/alexnet/FindTensorRT.cmake +++ b/alexnet/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/alexnet/alexnet.cc b/alexnet/alexnet.cc index 6fdfc409..b2c4c22b 100644 --- a/alexnet/alexnet.cc +++ b/alexnet/alexnet.cc @@ -1,6 +1,8 @@ #include #include #include +#include +#include #include #include #include "logging.h" @@ -16,7 +18,7 @@ constexpr const std::array NAMES = {"data", "prob"}; constexpr const char* ENGINE_PATH = "../models/alexnet.engine"; constexpr const char* WTS_PATH = "../models/alexnet.wts"; constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1) ? true : false; static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; @@ -39,9 +41,9 @@ static Logger gLogger; ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt) { WeightMap weightMap = loadWeights(WTS_PATH); -#if TRT_VERSION >= 11200 +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); @@ -140,7 +142,7 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil network->markOutput(*fc3_1->getOutput(0)); // Build engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); auto* host_mem = builder->buildSerializedNetwork(*network, *config); auto* engine = runtime->deserializeCudaEngine(host_mem->data(), host_mem->size()); @@ -173,7 +175,7 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { (*modelStream) = engine->serialize(); // Close everything down -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -206,7 +208,7 @@ std::vector> doInference(IExecutionContext& context, const st CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -214,12 +216,15 @@ std::vector> doInference(IExecutionContext& context, const st buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); std::size_t size = s * batchSize * SIZES[i]; CHECK(cudaMalloc(&buffers[i], size)); - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -232,10 +237,16 @@ std::vector> doInference(IExecutionContext& context, const st } } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; @@ -288,7 +299,7 @@ int main(int argc, char** argv) { auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -309,7 +320,7 @@ int main(int argc, char** argv) { return -1; } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -320,36 +331,28 @@ int main(int argc, char** argv) { assert(context != nullptr); const std::string img_path = "../assets/cats.jpg"; - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = doInference(*context, img_path, N); - auto _end = std::chrono::system_clock::now(); - auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "ms\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } - } - } + auto firstProb = doInference(*context, img_path, N); + printFirstOutputs("alexnet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + int _top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << labels[idx] + << "\n"; + } - if (i == 99) { - std::cout << "prediction result:\n"; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << "\n"; - } - } + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, img_path, N); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("alexnet", latencies, N); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; delete runtime; diff --git a/alexnet/macros.h b/alexnet/macros.h index 55284a01..8b0eacdb 100644 --- a/alexnet/macros.h +++ b/alexnet/macros.h @@ -17,10 +17,15 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/alexnet/utils.h b/alexnet/utils.h index 2d5089f9..da343a81 100644 --- a/alexnet/utils.h +++ b/alexnet/utils.h @@ -3,7 +3,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -27,7 +29,7 @@ constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -183,7 +185,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -228,7 +230,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; @@ -236,7 +238,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -253,3 +255,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/assets/IMG_1.jpg b/assets/IMG_1.jpg new file mode 100644 index 00000000..5d43d292 Binary files /dev/null and b/assets/IMG_1.jpg differ diff --git a/csrnet/CMakeLists.txt b/csrnet/CMakeLists.txt index 55bb75b3..70025e4b 100644 --- a/csrnet/CMakeLists.txt +++ b/csrnet/CMakeLists.txt @@ -1,26 +1,37 @@ -cmake_minimum_required(VERSION 3.10) +cmake_minimum_required(VERSION 3.14) -project(csrnet) +project( + csrnet + VERSION 0.1 + LANGUAGES C CXX) -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) +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90 100 120) +endif() -# cuda -include_directories(/usr/local/cuda/targets/x86_64-linux/include ) -link_directories(/usr/local/cuda/targets/x86_64-linux/lib) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_INCLUDE_CURRENT_DIR TRUE) -# tensorrt -include_directories(/usr/include/x86_64-linux-gnu/) -link_directories(/usr/lib/x86_64-linux-gnu/) +option(CUDA_USE_STATIC_CUDA_RUNTIME "Use static cudaruntime library" OFF) -# opencv -find_package(OpenCV) -include_directories(${OpenCV_INCLUDE_DIRS}) +find_package(Threads REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(OpenCV REQUIRED) -include_directories(${PROJECT_SOURCE_DIR}/) +if(NOT TARGET TensorRT::TensorRT) + include(FindTensorRT.cmake) +else() + message("TensorRT has been found, skipping for ${PROJECT_NAME}") +endif() -add_executable(csrnet csrnet.cpp) -target_link_libraries(csrnet nvinfer cudart ${OpenCV_LIBS}) \ No newline at end of file +add_executable(${PROJECT_NAME} ${PROJECT_NAME}.cpp) + +target_include_directories(${PROJECT_NAME} PUBLIC ${CMAKE_CURRENT_LIST_DIR} + ${OpenCV_INCLUDE_DIRS}) + +target_link_libraries(${PROJECT_NAME} PUBLIC Threads::Threads CUDA::cudart + TensorRT::TensorRT ${OpenCV_LIBS}) diff --git a/csrnet/FindTensorRT.cmake b/csrnet/FindTensorRT.cmake new file mode 100644 index 00000000..61203d17 --- /dev/null +++ b/csrnet/FindTensorRT.cmake @@ -0,0 +1,149 @@ +cmake_minimum_required(VERSION 3.17.0) + +function(_guess_path var_name required_files) + set(_result "") + + foreach(path_entry IN LISTS ARGN) + if(NOT EXISTS "${path_entry}") + message(DEBUG "skip non-existing path '${path_entry}'") + continue() + endif() + + set(_ok TRUE) + foreach(required_file IN LISTS required_files) + if(NOT EXISTS "${path_entry}/${required_file}") + set(_ok FALSE) + message(DEBUG "'${path_entry}' missing '${required_file}'") + break() + endif() + endforeach() + + if(_ok) + list(APPEND _result "${path_entry}") + message(DEBUG "accept '${path_entry}'") + else() + message(DEBUG "reject '${path_entry}'") + endif() + endforeach() + + if(_result STREQUAL "") + message( + FATAL_ERROR + "_guess_path(${var_name}) failed: no valid path found. required_files='${required_files}' candidates='${ARGN}'" + ) + endif() + + set(${var_name} + "${_result}" + PARENT_SCOPE) +endfunction() + +# add library +add_library(TensorRT IMPORTED INTERFACE) +add_library(TensorRT::TensorRT ALIAS TensorRT) + +set(TRT_VERSION + CACHE + STRING + "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", \"8.6.1.6.Windows10.x86_64.cuda-12.0\" etc" +) + +if(NOT TRT_VERSION STREQUAL "" AND NOT $ENV{TRT_VERSION} STREQUAL "") + message( + WARNING + "TRT_VERSION defined by cmake and environment variable both, using the later one" + ) +endif() + +if(NOT $ENV{TRT_VERSION} STREQUAL "") + set(TRT_VERSION $ENV{TRT_VERSION}) +endif() + +string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION}) +set(TRT_MAJOR_VERSION "${_match}") +unset(_match) + +if(WIN32) + set(TensorRT_DIR "C:/Program Files/TensorRT-${TRT_VERSION}") + if(NOT EXISTS "${TensorRT_DIR}") + message( + FATAL_ERROR + "TensorRT_DIR=${TensorRT_DIR} does not exist!" + ) + endif() + + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) + message(DEBUG "Using ${_modules}") + else() + set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch + nvinfer_lean) + endif() + + set(TensorRT_LIBRARY_DIR "${TensorRT_DIR}/lib") + set(TensorRT_INCLUDE_DIR "${TensorRT_DIR}/include") +elseif(UNIX) + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _trt_arch) + set(_trt_include_candidates) + if(_trt_arch MATCHES "^(aarch64|arm64|arch64)$") + set(_trt_include_candidates "/usr/include/aarch64-linux-gnu" "/usr/include" + "/usr/local/cuda/targets/aarch64-linux/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/aarch64-linux-gnu/lib" + "/usr/lib/aarch64-linux-gnu" "/usr/lib/aarch64-linux-gnu/tegra" + "/usr/lib") + elseif(_trt_arch MATCHES "^(x86_64|amd64)$") + set(_trt_include_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" + "/usr/include/x86_64-linux-gnu" "/usr/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" + "/usr/lib/x86_64-linux-gnu" "/usr/lib") + else() + message(FATAL_ERROR "Unknown architecture") + endif() + + set(_modules nvinfer nvinfer_plugin) + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 8) + list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean) + endif() + + _guess_path(TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so" + ${_trt_library_candidates}) + message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") + _guess_path(TensorRT_INCLUDE_DIR "NvInfer.h" ${_trt_include_candidates}) + message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") +endif() + +foreach(lib IN LISTS _modules) + find_library( + TensorRT_${lib}_LIBRARY + NAMES ${lib} + HINTS ${TensorRT_LIBRARY_DIR}) + list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY}) +endforeach() + +target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES}) + +message(STATUS "Found TensorRT libs: ${TensorRT_LIBRARIES}") + +set_target_properties( + TensorRT + PROPERTIES C_STANDARD 17 + CXX_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + SKIP_BUILD_RPATH TRUE + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN" + INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") + +unset(TRT_MAJOR_VERSION) +unset(_modules) +unset(_trt_lib_suffix) +unset(_trt_include_candidates) +unset(_trt_library_candidates) +unset(_trt_arch) diff --git a/csrnet/README.md b/csrnet/README.md index 54cfc4a5..f9c7aee8 100644 --- a/csrnet/README.md +++ b/csrnet/README.md @@ -1,58 +1,67 @@ # csrnet +## Overview + The Pytorch implementation is [leeyeehoo/CSRNet-pytorch](https://github.com/leeyeehoo/CSRNet-pytorch). This repo is a TensorRT implementation of CSRNet. +PartA model download url: https://drive.google.com/file/d/1Z-atzS5Y2pOd-nEWqZRVBDMYJDreGWHH/view + 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 +## Usage +1. download the pretrained model file and put it at `tensorrtx/models/PartAmodel_best.pth.tar` (skip if you already have it). -# how to run +2. generate `csrnet.wts` from `PartAmodel_best.pth.tar` ```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 +pushd tensorrtx/csrnet +python3 gen_wts.py +popd ``` +3. build C++ code + +```bash +pushd tensorrtx/csrnet +cmake -S . -B build -G Ninja --fresh +cmake --build build +popd +``` + +4. serialize to `tensorrtx/models/csrnet.engine` + +```bash +pushd tensorrtx/csrnet +./build/csrnet -s +popd +``` -# result +5. run inference (reads `tensorrtx/assets/IMG_1.jpg`, writes `tensorrtx/assets/csrnet_output.jpg`) + +```bash +pushd tensorrtx/csrnet +./build/csrnet -d +popd +``` + +# result + +output looks like: + +```bash +... +Execution time: 62956us +0.000543074 0.00035309 0.00351256 0.00194812 0.00201269 0.000398191 -0.000205946 -0.000301642 0.000487901 -9.49481e-05 +==== +Execution time: 62431us +0.000543074 0.00035309 0.00351256 0.00194812 0.00201269 0.000398191 -0.000205946 -0.000301642 0.000487901 -9.49481e-05 +==== +approximate people num: 21 +``` -inference people num: 22.9101 +you can also check the output image in `tensorrtx/assets/csrnet_output_tensorrt.jpg`: -

- -

+![output](../assets/csrnet_output_tensorrt.jpg) diff --git a/csrnet/config.h b/csrnet/config.h deleted file mode 100644 index 1bf0a17e..00000000 --- a/csrnet/config.h +++ /dev/null @@ -1,16 +0,0 @@ -#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; \ No newline at end of file diff --git a/csrnet/csrnet.cpp b/csrnet/csrnet.cpp index 7ef08bd4..93f42b65 100644 --- a/csrnet/csrnet.cpp +++ b/csrnet/csrnet.cpp @@ -1,536 +1,401 @@ -#include "NvInfer.h" -#include "cuda_runtime_api.h" +#include +#include #include -#include -#include -#include + +#include +#include +#include +#include +#include #include #include -#include -#include #include -#include -#include -using namespace nvinfer1; - -#define CHECK(status) \ - do { \ - auto ret = (status); \ - if (ret != 0) { \ - std::cerr << "Cuda failure: " << ret << std::endl; \ - abort(); \ - } \ - } while (0) +#include +#include +#include + +#include "NvInfer.h" +#include "utils.h" + +using LayerConfig = std::vector>; +using WeightMap = std::map; +using O = nvinfer1::OptProfileSelector; + +constexpr static const std::array NAMES = {"data", "prob"}; +constexpr static const char* WTS_PATH = "../models/csrnet.wts"; +constexpr static const char* ENGINE_PATH = "../models/csrnet.engine"; + +// for simplicity, always use BatchSize == 1 +constexpr static const int64_t maxBatchSize = 1; +constexpr static const int64_t MAX_INPUT_SIZE = 1440; +constexpr static const int64_t MIN_INPUT_SIZE = 608; + +/** @note: we use this value on purpose to simplify preprocess in this demo. + * To set your own H, W value, be sure to make them divisible by 32 */ +constexpr static const int64_t OPT_INPUT_W = 1024; +constexpr static const int64_t OPT_INPUT_H = 768; + +constexpr static int64_t kMaxInputImageSize = MAX_INPUT_SIZE * MAX_INPUT_SIZE * 3; +constexpr static int64_t kMaxOutputProbSize = (MAX_INPUT_SIZE * MAX_INPUT_SIZE) >> 6; + +static const LayerConfig frontend_cfg = {64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512}; +static const LayerConfig backend_cfg = {512, 512, 512, 256, 128, 64}; + +constexpr static const std::array mean = {0.406, 0.456, 0.485}; +constexpr static const std::array stdv = {0.225, 0.224, 0.229}; static Logger gLogger; -static char *kWTSFile = ""; -std::map loadWeights(const std::string file) { - std::cout << "Loading weights: " << file << std::endl; - std::map 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(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; - } +ILayer* addBatchNorm2d(INetworkDefinition* network, WeightMap& m, ITensor& input, const std::string& lname, + float eps = 1e-3) { + static Weights none{DataType::kFLOAT, nullptr, 0ll}; + const auto* gamma = reinterpret_cast(m[lname + ".weight"].values); + const auto* beta = reinterpret_cast(m[lname + ".bias"].values); + const auto* mean = reinterpret_cast(m[lname + ".running_mean"].values); + const auto* var = reinterpret_cast(m[lname + ".running_var"].values); + auto len = m[lname + ".running_var"].count; + + auto* scval = static_cast(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}; - return weightMap; + auto* shift_val = static_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + shift_val[i] = beta[i] - (mean[i] * scval[i]); + } + Weights shift{DataType::kFLOAT, shift_val, len}; + + m[lname + ".scale"] = scale; + m[lname + ".shift"] = shift; + m[lname + ".power"] = none; + auto* bn = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, none); + assert(bn); + bn->setName(lname.c_str()); + return bn; } -// 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( - 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])); + +static auto make_layers(INetworkDefinition* net, WeightMap& w, ITensor& input, const LayerConfig& cfg, + const std::string& name = "frontend.", bool use_bn = false, + bool use_dilation = false) -> ILayer* { + auto* tensor = &input; + int32_t idx = 0; + ILayer* ret = nullptr; + + for (const auto& _v : cfg) { + if (std::holds_alternative(_v)) { + auto* pool = net->addPoolingNd(*tensor, PoolingType::kMAX, DimsHW{2, 2}); + pool->setStrideNd(DimsHW{2, 2}); + pool->setName((name + std::to_string(idx++)).c_str()); + tensor = pool->getOutput(0); + ret = pool; + assert(ret); + } else { + assert(std::holds_alternative(_v)); + const int32_t v = std::get(_v); + auto _name = name + std::to_string(idx++); + auto* conv = net->addConvolutionNd(*tensor, v, DimsHW{3, 3}, w[_name + ".weight"], w[_name + ".bias"]); + auto d_rate = use_dilation ? DimsHW{2, 2} : DimsHW{1, 1}; + conv->setPaddingNd(d_rate); + conv->setDilationNd(d_rate); + conv->setName(_name.c_str()); + tensor = conv->getOutput(0); + if (use_bn) { + auto _name = name + std::to_string(idx++); + auto* bn = addBatchNorm2d(net, w, *tensor, _name, 1e-5f); + tensor = bn->getOutput(0); + assert(bn); + } + auto* relu = net->addActivation(*tensor, ActivationType::kRELU); + relu->setName((name + std::to_string(idx++)).c_str()); + tensor = relu->getOutput(0); + ret = relu; + assert(ret); + } + } + return ret; } -ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, - IBuilderConfig *config, DataType dt) { - - // INetworkDefinition *network = builder->createNetworkV2(0U); - const auto explicitBatch = - 1U << static_cast( - NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); - INetworkDefinition *network = builder->createNetworkV2(explicitBatch); - ITensor *data = network->addInput(kInputTensorName, dt, Dims4{1, 3, -1, -1}); - assert(data); - std::map 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); + +auto createEngine(int32_t maxBatchSize, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, + DataType dt) -> ICudaEngine* { + WeightMap w = loadWeights(WTS_PATH); + +#if TRT_VERSION_GE(10, 0, 0) + auto* network = builder->createNetworkV2(0); +#else + auto* network = builder->createNetworkV2(1u << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)); #endif - ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config); - printf("build engine successfully : %s\n", kEngineFile); - // Don't need the network any more - network->destroy(); + ITensor* data = network->addInput(NAMES[0], dt, Dims4{1, 3, -1, -1}); + assert(data); + + auto* frontend = make_layers(network, w, *data, frontend_cfg, "frontend.", false, false); + auto* backend = make_layers(network, w, *frontend->getOutput(0), backend_cfg, "backend.", false, true); + auto conv = network->addConvolutionNd(*backend->getOutput(0), 1, DimsHW{1, 1}, w["output_layer.weight"], + w["output_layer.bias"]); + assert(conv); + + conv->setStrideNd(DimsHW{1, 1}); + conv->getOutput(0)->setName(NAMES[1]); + network->markOutput(*conv->getOutput(0)); + + IOptimizationProfile* profile = builder->createOptimizationProfile(); + profile->setDimensions(NAMES[0], O::kMIN, Dims4(maxBatchSize, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE)); + profile->setDimensions(NAMES[0], O::kOPT, Dims4(maxBatchSize, 3, OPT_INPUT_H, OPT_INPUT_W)); + profile->setDimensions(NAMES[0], O::kMAX, Dims4(maxBatchSize, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE)); + config->addOptimizationProfile(profile); + +#if TRT_VERSION_GE(8, 0, 0) + config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); + auto* serialized = builder->buildSerializedNetwork(*network, *config); + auto* engine = runtime->deserializeCudaEngine(serialized->data(), serialized->size()); + delete serialized; + delete network; +#else + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(WORKSPACE_SIZE); + auto* engine = builder->buildEngineWithConfig(*network, *config); + network->destroy(); +#endif + std::cout << "build out\n"; - // Release host memory - for (auto &mem : weightMap) { - free((void *)(mem.second.values)); - } + for (auto& _mem : w) { + free(const_cast(_mem.second.values)); + } - return engine; + 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(); + +void APIToModel(int32_t batch_size, IRuntime* runtime, IHostMemory** model_stream) { + auto* builder = createInferBuilder(gLogger); + auto* config = builder->createBuilderConfig(); + + auto* engine = createEngine(batch_size, runtime, builder, config, DataType::kFLOAT); + assert(engine != nullptr); + + (*model_stream) = engine->serialize(); + +#if TRT_VERSION_GE(8, 0, 0) + delete engine; + delete config; + delete builder; +#else + engine->destroy(); + config->destroy(); + builder->destroy(); +#endif } -int read_files_in_dir(const char *p_dir_name, - std::vector &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); +std::vector doInference(IExecutionContext& context, void* input, int batchSize) { + const auto& engine = context.getEngine(); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + std::vector buffers; + +#if TRT_VERSION_GE(8, 0, 0) + const int32_t nIO = engine.getNbIOTensors(); +#else + const int32_t nIO = engine.getNbBindings(); +#endif + + buffers.resize(nIO); + Dims out_shape; + std::size_t out_size; + context.setOptimizationProfileAsync(0, stream); + for (auto i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) + const auto* tensor_name = engine.getIOTensorName(i); + auto s = getSize(engine.getTensorDataType(tensor_name)); + if (i == 0) { + std::size_t size = s * batchSize * 3 * OPT_INPUT_W * OPT_INPUT_H; + CHECK(cudaMalloc(&buffers[i], size)); + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + // since input config is dynamic, must set input shape to let TensorRT deduce the output shape + context.setInputShape(tensor_name, Dims4{batchSize, 3, OPT_INPUT_H, OPT_INPUT_W}); + context.setInputTensorAddress(tensor_name, buffers[i]); + } + + // we have already known nIO==2, so make it simple here + if (i > 0) { + out_shape = context.getTensorShape(tensor_name); + out_size = std::accumulate(out_shape.d, out_shape.d + out_shape.nbDims, 1ULL, std::multiplies<>()); + CHECK(cudaMalloc(&buffers[i], s * out_size)); +#if TRT_VERSION_GE(10, 0, 0) + if (!context.setOutputTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setOutputTensorAddress failed\n"; + std::abort(); + } +#else + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } +#endif + } +#else + const int32_t idx = engine.getBindingIndex(NAMES[i]); + auto s = getSize(engine.getBindingDataType(idx)); + assert(idx == i); + if (engine.bindingIsInput(idx)) { + std::size_t size = s * batchSize * 3 * OPT_INPUT_W * OPT_INPUT_H; + CHECK(cudaMalloc(&buffers[i], size)); + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + context.setBindingDimensions(idx, Dims4{batchSize, 3, OPT_INPUT_H, OPT_INPUT_W}); + } else { + assert(context.allInputDimensionsSpecified()); + out_shape = context.getBindingDimensions(idx); + out_size = std::accumulate(out_shape.d, out_shape.d + out_shape.nbDims, 1, std::multiplies()); + CHECK(cudaMalloc(&buffers[i], s * out_size)); + } +#endif } - } - closedir(p_dir); - return 0; + +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } +#else + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } +#endif + + std::vector prob; + prob.reserve(static_cast(nIO > 0 ? nIO - 1 : 0)); + for (int i = 1; i < nIO; ++i) { + prob.emplace_back(out_shape, DataType::kFLOAT, -1); + CHECK(cudaMemcpyAsync(prob.back().data, buffers[i], out_size * sizeof(float), cudaMemcpyDeviceToHost, stream)); + } + CHECK(cudaStreamSynchronize(stream)); + + for (auto& buffer : buffers) { + CHECK(cudaFree(buffer)); + } + CHECK(cudaStreamDestroy(stream)); + return prob; } -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; +static auto mainImpl(int argc, char** argv) -> int { + checkTrtEnv(); + if (argc != 2) { + std::cerr << "arguments not right!\n" << std::flush; + std::cerr << "./csrnet -s // serialize model to plan file\n" << std::flush; + std::cerr << "./csrnet -d // deserialize plan file and run inference\n" << std::flush; + return -1; + } + char* trtModelStream{nullptr}; + std::streamsize size{0}; + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + + if (std::string(argv[1]) == "-s") { + IHostMemory* modelStream{nullptr}; + APIToModel(maxBatchSize, runtime, &modelStream); + assert(modelStream != nullptr); + + std::ofstream p(ENGINE_PATH, std::ios::binary); + if (!p) { + std::cerr << "could not open plan output file\n"; + return -1; + } + if (modelStream->size() > static_cast(std::numeric_limits::max())) { + std::cerr << "this model is too large to serialize\n"; + return -1; + } + const auto* data_ptr = reinterpret_cast(modelStream->data()); + auto data_size = static_cast(modelStream->size()); + p.write(data_ptr, data_size); +#if TRT_VERSION_GE(8, 0, 0) + delete modelStream; +#else + modelStream->destroy(); +#endif + return 0; + } + if (std::string(argv[1]) == "-d") { + std::ifstream file(ENGINE_PATH, 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; + } + +#if TRT_VERSION_GE(8, 0, 0) + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); +#else + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); +#endif + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + + // read and preprocess input data + cv::Mat img_32f, rgb; + auto img = cv::imread("../assets/IMG_1.jpg", cv::IMREAD_COLOR); + img.convertTo(img_32f, CV_32FC3, 1.0 / 255.0, 0); + img_32f = (img_32f - cv::Scalar(mean[0], mean[1], mean[2])) / cv::Scalar(stdv[0], stdv[1], stdv[2]); + cv::cvtColor(img_32f, rgb, cv::COLOR_BGR2RGB); + static std::size_t img_size = OPT_INPUT_H * OPT_INPUT_W * img.channels(); + std::vector data(img_size); + for (int i = 0; i < maxBatchSize; ++i) { + // to NCHW (N == 1) + for (int y = 0; y < OPT_INPUT_H; ++y) { + for (int x = 0; x < OPT_INPUT_W; ++x) { + const cv::Vec3f v = img_32f.at(y, x); + data[i * img_size + 0 * OPT_INPUT_H * OPT_INPUT_W + y * OPT_INPUT_W + x] = v[0]; + data[i * img_size + 1 * OPT_INPUT_H * OPT_INPUT_W + y * OPT_INPUT_W + x] = v[1]; + data[i * img_size + 2 * OPT_INPUT_H * OPT_INPUT_W + y * OPT_INPUT_W + x] = v[2]; + } + } } - p.write(reinterpret_cast(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(); + + auto output = doInference(*context, data.data(), 1); + const auto* data_ptr = reinterpret_cast(output[0].data); + + DummyTensor& t = output[0]; + const auto out_h = t.dims.d[2]; + const auto out_w = t.dims.d[3]; + auto stride = static_cast(out_h) * out_w; + printFirstOutputs("csrnet", data_ptr, static_cast(stride)); + float num = std::accumulate(data_ptr, data_ptr + stride, 0.0f); + cv::Mat density((int)out_h, (int)out_w, CV_32FC1, t.data); + cv::Mat scaled, heatmap, save; + cv::normalize(density, scaled, 0, 255, cv::NORM_MINMAX, CV_8UC1); + cv::applyColorMap(scaled, heatmap, cv::COLORMAP_JET); + cv::resize(heatmap, heatmap, img.size(), 0., 0., cv::INTER_LINEAR); + cv::addWeighted(img, 0.7, heatmap, 0.3, 0, save); + cv::imwrite("../assets/csrnet_output_tensorrt.jpg", save); + std::cout << "approximate people num: " << std::ceil(num) << "\nsave to ../assets/csrnet_output_tensorrt.jpg\n"; + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, data.data(), 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } - } 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 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 mean_value{0.406, 0.456, 0.485}; // BGR - std::vector 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; - } + printBenchmark("csrnet", latencies); + + return 0; +} + +auto main(int argc, char** argv) -> int { + try { + return mainImpl(argc, argv); + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return -1; } - // 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(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; -} \ No newline at end of file +} diff --git a/csrnet/gen_wts.py b/csrnet/gen_wts.py index 68ba51fb..cd6e29f2 100644 --- a/csrnet/gen_wts.py +++ b/csrnet/gen_wts.py @@ -1,31 +1,185 @@ -from torch.nn.modules import module -from model import CSRNet -import torch -import os +""" +this file is modified from https://github.com/leeyeehoo/CSRNet-pytorch +to make it compatible with python3, and support wts file exporting +""" + import struct +import cv2 +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import models + + +class CSRNet(nn.Module): + def __init__(self, load_weights=False): + super(CSRNet, self).__init__() + self.seen = 0 + self.frontend_feat = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512] + self.backend_feat = [512, 512, 512, 256, 128, 64] + self.frontend = make_layers(self.frontend_feat) + self.backend = make_layers(self.backend_feat, in_channels=512, dilation=True) + self.output_layer = nn.Conv2d(64, 1, kernel_size=1) + if not load_weights: + mod = models.vgg16(pretrained=True) + self._initialize_weights() + for i in range(len(self.frontend.state_dict().items())): + self.frontend.state_dict().items()[i][1].data[:] = mod.state_dict().items()[i][1].data[:] + + def forward(self, x): + x = self.frontend(x) + x = self.backend(x) + x = self.output_layer(x) + return x + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, std=0.01) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + +def make_layers(cfg, in_channels=3, batch_norm=False, dilation=False): + if dilation: + d_rate = 2 + else: + d_rate = 1 + layers = [] + for v in cfg: + if v == "M": + layers += [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate, dilation=d_rate) + if batch_norm: + layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] + else: + layers += [conv2d, nn.ReLU(inplace=True)] + in_channels = v + return nn.Sequential(*layers) + + +def preprocess(img: np.array) -> torch.Tensor: + """ + a preprocess method align with CSRNet: no resize, no normalization + input resolution: (1024, 768, 3) + + Args: + img (np.array): input image + + Returns: + torch.Tensor: preprocessed image in `NCHW` layout + """ + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) + mean = np.array([92.8207477031, 95.2757037428, 104.877445883], dtype=np.float32) + img = img - mean + img = img.transpose(2, 0, 1)[None, ...] + return torch.from_numpy(img) + + +def visualize() -> None: + """this method is for anyone who wants to see the ground truth output + of the input included in this demo, just call this method directly. + """ + import scipy + + def gaussian_filter_density(gt): + density = np.zeros(gt.shape, dtype=np.float32) + gt_count = np.count_nonzero(gt) + + if gt_count == 0: + return density + + ys, xs = np.nonzero(gt) + pts = np.column_stack((xs, ys)) + + tree = scipy.spatial.cKDTree(pts, leafsize=2048) # build KDTree + distances, _ = tree.query(pts, k=min(4, gt_count)) + + for i, (x, y) in enumerate(pts): + pt2d = np.zeros(gt.shape, dtype=np.float32) + pt2d[y, x] = 1.0 + + if gt_count > 1: + sigma = np.sum(distances[i][1:]) * 0.1 + else: + sigma = np.mean(gt.shape) / 4.0 + + density += scipy.ndimage.gaussian_filter(pt2d, sigma, mode="constant") + + return density + + def visualize_density_cv2(img: np.array, density): + if density.max() > 0: + density_norm = density / density.max() + else: + density_norm = density + + density_uint8 = (density_norm * 255).astype(np.uint8) + heatmap = cv2.applyColorMap(density_uint8, cv2.COLORMAP_JET) + blended = cv2.addWeighted(img, 0.7, heatmap, 0.3, 0) + return blended + + img = cv2.imread("../assets/IMG_1.jpg", cv2.IMREAD_COLOR) + gt = scipy.io.loadmat("../assets/GT_IMG_1.mat") + gt = gt["image_info"][0, 0][0, 0][0] + k = np.zeros((img.shape[0], img.shape[1])) + for i in range(0, len(gt)): + if int(gt[i][1]) < img.shape[0] and int(gt[i][0]) < img.shape[1]: + k[int(gt[i][1]), int(gt[i][0])] = 1 + k = gaussian_filter_density(k) + img = visualize_density_cv2(img, k) + cv2.imwrite("../assets/csrnet_gt_vis.jpg", img) + + +def main(): + # load pth model + model_path = "../models/PartAmodel_best.pth.tar" # partBmodel_best.pth + model = CSRNet(load_weights=True) + checkpoint = torch.load(model_path, weights_only=False) + model.load_state_dict(checkpoint["state_dict"]) + model.eval() + with torch.inference_mode(): + img = cv2.imread("../assets/IMG_1.jpg", cv2.IMREAD_COLOR) + h, w, _ = img.shape + data = preprocess(img) + output = model(data) -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") + # this is the postprocess part + # 1. convert the output to heatmap + heatmap = F.interpolate(output, size=(h, w), mode="bilinear", align_corners=False) + heatmap = heatmap[0].detach().cpu().numpy().transpose(1, 2, 0) + heatmap = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX) + heatmap = cv2.applyColorMap(heatmap.astype(np.uint8)[..., None], cv2.COLORMAP_JET) + # 2. blend the heatmap with the original image + heatmap = cv2.resize(heatmap, (w, h), 0, 0, cv2.INTER_LINEAR) + overlay = cv2.addWeighted(img, 0.7, heatmap, 0.3, 0) + status = cv2.imwrite("../assets/csrnet_output_torch.jpg", overlay) + assert status + # save to wts + print("Writing into csrnet.wts") + with open("../models/csrnet.wts", "w") as f: + f.write("{}\n".format(len(model.state_dict().keys()))) + for k, v in model.state_dict().items(): + print(f"key: {k}\tvalue: {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") -# load model -model_path = "partBmodel_best.pth.tar" -model = CSRNet() -checkpoint = torch.load(model_path) -model.load_state_dict(checkpoint['state_dict']) +if __name__ == "__main__": + # Uncomment the following line to generate ground truth heatmap + # with the file in `../assets/GT_IMG_1.mat`, check the result + # at `../assets/csrnet_gt_vis.jpg` + # visualize() -# 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') \ No newline at end of file + main() diff --git a/csrnet/logging.h b/csrnet/logging.h index fb09a20a..4e2bd80d 100644 --- a/csrnet/logging.h +++ b/csrnet/logging.h @@ -17,486 +17,450 @@ #ifndef TENSORRT_LOGGING_H #define TENSORRT_LOGGING_H -#include "NvInferRuntimeCommon.h" -#include "macros.h" #include +#include #include #include #include #include #include #include +#include +#include "NvInferRuntime.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(); + public: + LogStreamConsumerBuffer(std::ostream& stream, std::string prefix, bool shouldLog) + : mOutput(stream), mPrefix(std::move(prefix)), mShouldLog(shouldLog) {} + + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) noexcept + : mOutput(other.mOutput), mPrefix(std::move(other.mPrefix)), mShouldLog(other.mShouldLog) {} + + ~LogStreamConsumerBuffer() override { + // 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(×tamp); - 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(); + + // 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 + int sync() override { + putOutput(); + return 0; + } + + void putOutput() { + if (mShouldLog) { + // prepend timestamp + std::time_t timestamp = std::time(nullptr); + tm* tm_local = std::localtime(×tamp); + 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; } + void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; } -private: - std::ostream &mOutput; - std::string mPrefix; - bool mShouldLog; + private: + std::ostream& mOutput; + std::string mPrefix; + bool mShouldLog; }; //! //! \class LogStreamConsumerBase -//! \brief Convenience object used to initialize LogStreamConsumerBuffer before -//! std::ostream in LogStreamConsumer +//! \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) {} + public: + LogStreamConsumerBase(std::ostream& stream, std::string prefix, bool shouldLog) + : mBuffer(stream, std::move(prefix), shouldLog) {} -protected: - LogStreamConsumerBuffer mBuffer; + protected: + LogStreamConsumerBuffer mBuffer; }; //! //! \class LogStreamConsumer -//! \brief Convenience object used to facilitate use of C++ stream syntax when -//! logging messages. +//! \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. +//! 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 ""; + 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) noexcept + : 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; + 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: +//! \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) +//! - 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. +//! 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). +//! 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 +//! 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 +//! 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 ""; + private: + struct TestInfo; + + public: + Logger(Severity severity = Severity::kWARNING) : mReportableSeverity(severity) {} + + //! + //! \enum TestResult + //! \brief Represents the state of a given test + //! + enum class TestResult : std::uint8_t { + 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) << '\n'; } - } - - //! - //! \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 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, TestInfo info) + : mStarted(started), mName(std::move(info.name)), mCmdline(std::move(info.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, TestInfo{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); } - } - - //! - //! \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]; + + 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); + } + + [[nodiscard]] Severity getReportableSeverity() const { return mReportableSeverity; } + + private: + struct TestInfo { + std::string name; + std::string cmdline; + }; + //! + //! \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 << '\n'; + } + + //! + //! \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(); } - return ss.str(); - } - Severity mReportableSeverity; + Severity mReportableSeverity; }; namespace { //! -//! \brief produces a LogStreamConsumer object that can be used to log messages -//! of severity kVERBOSE +//! \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); +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 +//! \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); +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 +//! \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); +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 +//! \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); +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 +//! \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); +inline LogStreamConsumer LOG_FATAL(const Logger& logger) { + return LogStreamConsumer{logger.getReportableSeverity(), Severity::kINTERNAL_ERROR}; } -} // anonymous namespace +} // anonymous namespace -#endif // TENSORRT_LOGGING_H +#endif // TENSORRT_LOGGING_H diff --git a/csrnet/macros.h b/csrnet/macros.h index 05551039..b8155b4c 100644 --- a/csrnet/macros.h +++ b/csrnet/macros.h @@ -1,12 +1,33 @@ -#ifndef __MACROS_H -#define __MACROS_H +#pragma once +#include -#if NV_TENSORRT_MAJOR >= 8 +#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 + +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) + +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else #define TRT_NOEXCEPT #define TRT_CONST_ENQUEUE #endif - -#endif // __MACROS_H diff --git a/csrnet/utils.h b/csrnet/utils.h new file mode 100644 index 00000000..5e55857f --- /dev/null +++ b/csrnet/utils.h @@ -0,0 +1,199 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "macros.h" + +using namespace nvinfer1; + +constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; + +#define CHECK(status) \ + do { \ + auto ret = (status); \ + if (ret != cudaSuccess) { \ + std::cerr << "Cuda failure: " << ret << "\n"; \ + std::abort(); \ + } \ + } while (0) + +static void checkTrtEnv(int device = 0) { +#if TRT_VERSION_LT(8, 0, 0) + CHECK(cudaGetDevice(&device)); + cudaDeviceProp prop{}; + CHECK(cudaGetDeviceProperties(&prop, device)); + const int sm = prop.major * 10 + prop.minor; + if (sm > 86) { + std::cerr << "TensorRT < 8 does not support SM > 86 on this GPU."; + std::abort(); + } +#endif +} + +/** + * @brief TensorRT weight files have a simple space delimited format: + * [type] [size] + * + * @param file input weight file path + * @return std::map + */ +static auto loadWeights(const std::string& file) { + std::cout << "Loading weights: " << file << "\n"; + std::map 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--) { + nvinfer1::Weights wt{nvinfer1::DataType::kFLOAT, nullptr, 0}; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> wt.count; + + // Load blob + auto* val = new uint32_t[wt.count]; + input >> std::hex; + for (auto x = 0ll; x < wt.count; ++x) { + input >> val[x]; + } + wt.values = val; + weightMap[name] = wt; + } + + return weightMap; +} + +static size_t getSize(DataType dt) { + switch (dt) { +#if TRT_VERSION_GE(8, 5, 1) + case DataType::kUINT8: +#endif + case DataType::kINT8: + return sizeof(int8_t); + case DataType::kFLOAT: + return sizeof(float); + case DataType::kHALF: + return sizeof(int16_t); + case DataType::kINT32: + return sizeof(int32_t); + default: { + std::cerr << "Unsupported data type\n"; + std::abort(); + } + } +} + +class DummyTensor final { + public: + DummyTensor(Dims dims, DataType dtype, int dev = -1) : dims(dims), dtype(dtype), dev(dev) { + total = std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies<>()); + if (dev == -1) { + data = malloc(getSize(dtype) * total); + } else { + std::cerr << "dev != -1 not implemented!"; + std::abort(); + } + } + DummyTensor(const DummyTensor& other) : DummyTensor(other.dims, other.dtype, other.dev) { + std::memcpy(data, other.data, getSize(dtype) * total); + } + DummyTensor& operator=(const DummyTensor& other) { + if (this == &other) { + return *this; + } + DummyTensor tmp(other); + std::swap(total, tmp.total); + std::swap(dev, tmp.dev); + std::swap(dtype, tmp.dtype); + std::swap(dims, tmp.dims); + std::swap(data, tmp.data); + return *this; + } + DummyTensor(DummyTensor&& other) noexcept + : total(other.total), dev(other.dev), dtype(other.dtype), dims(other.dims), data(other.data) { + other.total = 0; + other.data = nullptr; + } + DummyTensor& operator=(DummyTensor&& other) noexcept { + if (this == &other) { + return *this; + } + if (data) { + free(data); + } + total = other.total; + dev = other.dev; + dtype = other.dtype; + dims = other.dims; + data = other.data; + other.total = 0; + other.data = nullptr; + return *this; + } + ~DummyTensor() { + if (data) { + free(data); + } + } + std::size_t total; + int dev; + DataType dtype; + Dims dims; + void* data; +}; + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/googlenet/FindTensorRT.cmake b/googlenet/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/googlenet/FindTensorRT.cmake +++ b/googlenet/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/googlenet/googlenet.cpp b/googlenet/googlenet.cpp index da2494c8..3cc21062 100644 --- a/googlenet/googlenet.cpp +++ b/googlenet/googlenet.cpp @@ -2,6 +2,8 @@ #include #include #include +#include +#include #include #include #include "logging.h" @@ -20,7 +22,7 @@ static constexpr const int32_t INPUT_H = 224; static constexpr const int32_t INPUT_W = 224; static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; static constexpr const std::array NAMES = {"data", "prob"}; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1) ? true : false; static constexpr const char* WTS_PATH = "../models/googlenet.wts"; static constexpr const char* ENGINE_PATH = "../models/googlenet.engine"; static constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; @@ -128,9 +130,9 @@ IConcatenationLayer* inception(INetworkDefinition* network, WeightMap& weightMap ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt) { WeightMap weightMap = loadWeights(WTS_PATH); -#if TRT_VERSION >= 11200 +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); @@ -201,7 +203,7 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil fc1->getOutput(0)->setName(NAMES[1]); network->markOutput(*fc1->getOutput(0)); // Build engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); @@ -234,7 +236,7 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { // Serialize the engine (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -251,7 +253,7 @@ std::vector> doInference(IExecutionContext& context, void* in CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -260,7 +262,7 @@ std::vector> doInference(IExecutionContext& context, void* in buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); size = s * batchSize * SIZES[i]; @@ -268,7 +270,10 @@ std::vector> doInference(IExecutionContext& context, void* in if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -281,10 +286,16 @@ std::vector> doInference(IExecutionContext& context, void* in #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; @@ -335,7 +346,7 @@ int main(int argc, char** argv) { const auto* data_ptr = reinterpret_cast(modelStream->data()); auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -356,7 +367,7 @@ int main(int argc, char** argv) { return 1; } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -380,37 +391,29 @@ int main(int argc, char** argv) { } assert(input); - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = doInference(*context, input, 1); - auto _end = std::chrono::system_clock::now(); - auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "us\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } - } - } + auto firstProb = doInference(*context, input, 1); + printFirstOutputs("googlenet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + int _top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << labels[idx] + << "\n"; + } - if (i == 99) { - std::cout << "prediction result:\n"; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << "\n"; - } - } + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("googlenet", latencies); delete[] trtModelStream; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; delete runtime; diff --git a/googlenet/macros.h b/googlenet/macros.h index daae7bfb..6d851752 100644 --- a/googlenet/macros.h +++ b/googlenet/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/googlenet/utils.h b/googlenet/utils.h index fa723458..842b7131 100644 --- a/googlenet/utils.h +++ b/googlenet/utils.h @@ -2,7 +2,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -25,7 +28,7 @@ constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -181,7 +184,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -226,7 +229,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; @@ -234,7 +237,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -251,3 +254,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/lenet/FindTensorRT.cmake b/lenet/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/lenet/FindTensorRT.cmake +++ b/lenet/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/lenet/lenet.cpp b/lenet/lenet.cpp index 379edbbd..dc6c1fb4 100644 --- a/lenet/lenet.cpp +++ b/lenet/lenet.cpp @@ -2,8 +2,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -34,9 +36,9 @@ static Logger gLogger; * @return ICudaEngine* */ ICudaEngine* createLenetEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt) { -#if TRT_VERSION >= 11200 +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); @@ -129,7 +131,7 @@ ICudaEngine* createLenetEngine(int32_t N, IRuntime* runtime, IBuilder* builder, prob->getOutput(0)->setName(NAMES[1]); network->markOutput(*prob->getOutput(0)); -#if TRT_VERSION >= 8400 +#if TRT_VERSION_GE(8, 4, 0) config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); #else config->setMaxWorkspaceSize(WORKSPACE_SIZE); @@ -137,7 +139,7 @@ ICudaEngine* createLenetEngine(int32_t N, IRuntime* runtime, IBuilder* builder, #endif // Build engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) IHostMemory* serialized_mem = builder->buildSerializedNetwork(*network, *config); ICudaEngine* engine = runtime->deserializeCudaEngine(serialized_mem->data(), serialized_mem->size()); delete network; @@ -173,7 +175,7 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { // Serialize the engine (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -190,7 +192,7 @@ std::vector> doInference(IExecutionContext& context, void* in CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -199,7 +201,7 @@ std::vector> doInference(IExecutionContext& context, void* in buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); size = s * batchSize * SIZES[i]; @@ -207,7 +209,10 @@ std::vector> doInference(IExecutionContext& context, void* in if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -220,10 +225,16 @@ std::vector> doInference(IExecutionContext& context, void* in #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; @@ -275,7 +286,7 @@ int main(int argc, char** argv) { auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -305,7 +316,7 @@ int main(int argc, char** argv) { img = img / cv::Scalar(0.3081); assert(img.total() * img.elemSize() == SIZES[0] * sizeof(float)); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -314,36 +325,26 @@ int main(int argc, char** argv) { IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); - // Run inference - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = doInference(*context, img.data, 1); - auto _end = std::chrono::system_clock::now(); - auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "us\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 9) { - std::cout << "\n====\n"; - break; - } - } - } + auto firstProb = doInference(*context, img.data, 1); + printFirstOutputs("lenet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + int _top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << idx << "\n"; + } - if (i == 99) { - std::cout << "prediction result:\n"; - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << idx - << "\n"; - } - } + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, img.data, 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("lenet", latencies); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; delete runtime; diff --git a/lenet/macros.h b/lenet/macros.h index daae7bfb..6d851752 100644 --- a/lenet/macros.h +++ b/lenet/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/lenet/utils.h b/lenet/utils.h index 2dcf2436..e866807a 100644 --- a/lenet/utils.h +++ b/lenet/utils.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -27,7 +28,7 @@ enum : std::uint32_t { WORKSPACE_SIZE = 16 << 20 }; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -98,7 +99,7 @@ static std::vector> topk(const std::vector& v, int6 static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -115,3 +116,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/lprnet/FindTensorRT.cmake b/lprnet/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/lprnet/FindTensorRT.cmake +++ b/lprnet/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/lprnet/lprnet.cpp b/lprnet/lprnet.cpp index 2295a3a2..28990a63 100644 --- a/lprnet/lprnet.cpp +++ b/lprnet/lprnet.cpp @@ -3,12 +3,14 @@ #include #include #include +#include #include #include #include #include #include #include + #include "logging.h" #include "utils.h" #ifdef _WIN32 @@ -33,7 +35,7 @@ static constexpr const int32_t INPUT_H = 24; static constexpr const int32_t INPUT_W = 94; static constexpr const std::array NAMES = {"data", "prob"}; static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 18 * 68}; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1) ? true : false; static constexpr const std::array mean = {0.5f, 0.5f, 0.5f}; static constexpr const std::array stdv = {1.f, 1.f, 1.f}; @@ -114,9 +116,9 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil const int nc = 68; WeightMap w = loadWeights(WTS_PATH); -#if TRT_VERSION >= 11200 +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); @@ -250,7 +252,7 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil network->markOutput(*logits->getOutput(0)); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); @@ -280,7 +282,7 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -297,7 +299,7 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -306,7 +308,7 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); size = s * batchSize * SIZES[i]; @@ -314,7 +316,10 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -327,10 +332,16 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; @@ -385,7 +396,7 @@ int main(int argc, char** argv) { const auto* data_ptr = reinterpret_cast(modelStream->data()); auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -418,7 +429,7 @@ int main(int argc, char** argv) { input = data.data(); } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -427,44 +438,35 @@ int main(int argc, char** argv) { IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = doInference(*context, input, 1); - auto _end = std::chrono::system_clock::now(); - auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "us\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } - } - } - - if (i == 99) { - int prev = 67; - std::string str; - for (int t = 0; t < 18; ++t) { - std::array scores{}; - for (int c = 0; c < 68; ++c) { - scores[c] = prob[0][t + 18 * c]; - } - int best = - static_cast(std::distance(scores.begin(), std::max_element(scores.begin(), scores.end()))); - if (best != prev && best != 67) - str += alphabet[best]; - prev = best; - } - std::cout << "result: " << str << "\n"; + auto firstProb = doInference(*context, input, 1); + printFirstOutputs("lprnet", firstProb[0].data(), firstProb[0].size()); + int prev = 67; + std::string str; + for (int t = 0; t < 18; ++t) { + std::array scores{}; + for (int c = 0; c < 68; ++c) { + scores[c] = firstProb[0][t + 18 * c]; } + int best = static_cast(std::distance(scores.begin(), std::max_element(scores.begin(), scores.end()))); + if (best != prev && best != 67) + str += alphabet[best]; + prev = best; + } + std::cout << "result: " << str << "\n"; + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("lprnet", latencies); delete[] trtModelStream; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; delete runtime; diff --git a/lprnet/macros.h b/lprnet/macros.h index daae7bfb..6d851752 100644 --- a/lprnet/macros.h +++ b/lprnet/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/lprnet/utils.h b/lprnet/utils.h index a2cd4ee3..784646e3 100644 --- a/lprnet/utils.h +++ b/lprnet/utils.h @@ -2,7 +2,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -23,7 +26,7 @@ using namespace nvinfer1; } while (0) static inline void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -179,7 +182,7 @@ static inline ILayer* addTransformLayer(INetworkDefinition* network, ITensor& in ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -224,7 +227,7 @@ static inline ILayer* addTransformLayer(INetworkDefinition* network, ITensor& in auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; @@ -232,7 +235,7 @@ static inline ILayer* addTransformLayer(INetworkDefinition* network, ITensor& in static inline size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -249,3 +252,41 @@ static inline size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/mlp/FindTensorRT.cmake b/mlp/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/mlp/FindTensorRT.cmake +++ b/mlp/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/mlp/macros.h b/mlp/macros.h index daae7bfb..6d851752 100644 --- a/mlp/macros.h +++ b/mlp/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/mlp/mlp.cpp b/mlp/mlp.cpp index 3a2fe9ae..4bd8c9cd 100644 --- a/mlp/mlp.cpp +++ b/mlp/mlp.cpp @@ -1,8 +1,10 @@ #include #include +#include #include #include #include + #include "logging.h" #include "utils.h" @@ -34,7 +36,7 @@ ICudaEngine* createMLPEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IB std::map weightMap = loadWeights(WTS_PATH); // Create an empty network -#if TRT_VERSION >= 10000 +#if TRT_VERSION_GE(10, 0, 0) auto* network = builder->createNetworkV2(0); #else auto* network = builder->createNetworkV2(1u << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)); @@ -61,7 +63,7 @@ ICudaEngine* createMLPEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IB // mark the output network->markOutput(*output); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) IHostMemory* serialized_mem = builder->buildSerializedNetwork(*network, *config); ICudaEngine* engine = runtime->deserializeCudaEngine(serialized_mem->data(), serialized_mem->size()); delete network; @@ -100,7 +102,7 @@ void APIToModel(int32_t maxBatchSize, IRuntime* runtime, IHostMemory** modelStre // serialize the engine into binary stream (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -123,7 +125,7 @@ void doInference(IExecutionContext& ctx, void* input, float* output, int64_t bat // Get engine from the ctx const ICudaEngine& engine = ctx.getEngine(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) int32_t nIO = engine.getNbIOTensors(); const int inputIndex = 0; const int outputIndex = engine.getNbIOTensors() - 1; @@ -142,7 +144,7 @@ void doInference(IExecutionContext& ctx, void* input, float* output, int64_t bat std::vector buffers(nIO, nullptr); size_t inputSize = 0; size_t outputSize = batchSize * OUTPUT_SIZE * sizeof(float); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* input_name = engine.getIOTensorName(inputIndex); inputSize = batchSize * INPUT_SIZE * getSize(engine.getTensorDataType(input_name)); #else @@ -153,17 +155,23 @@ void doInference(IExecutionContext& ctx, void* input, float* output, int64_t bat CHECK(cudaMemcpyAsync(buffers[inputIndex], input, inputSize, cudaMemcpyHostToDevice, stream)); // execute inference using ctx provided by engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) for (int32_t i = 0; i < engine.getNbIOTensors(); i++) { auto const name = engine.getIOTensorName(i); - auto dims = ctx.getTensorShape(name); - auto total = std::accumulate(dims.d, dims.d + dims.nbDims, 1ll, std::multiplies<>()); - std::cout << name << "\t" << total << "\n"; - ctx.setTensorAddress(name, buffers[i]); + if (!ctx.setTensorAddress(name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } + } + if (!ctx.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); } - assert(ctx.enqueueV3(stream)); #else - assert(ctx.enqueueV2(buffers.data(), stream, nullptr)); + if (!ctx.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif CHECK(cudaMemcpyAsync(output, buffers[outputIndex], outputSize, cudaMemcpyDeviceToHost, stream)); @@ -206,7 +214,7 @@ int main(int argc, char** argv) { auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -227,7 +235,7 @@ int main(int argc, char** argv) { } } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -241,16 +249,21 @@ int main(int argc, char** argv) { std::array output = {-1.f}; std::array input = {12.0f}; - for (int i = 0; i < 100; i++) { - auto start = std::chrono::high_resolution_clock::now(); + doInference(*ctx, input.data(), output.data()); + printFirstOutputs("mlp", output.data(), output.size()); + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int i = 0; i < kBenchmarkRuns; i++) { + auto start = std::chrono::steady_clock::now(); doInference(*ctx, input.data(), output.data()); - auto end = std::chrono::high_resolution_clock::now(); + auto end = std::chrono::steady_clock::now(); auto time = std::chrono::duration_cast(end - start).count(); - std::cout << "Execution time: " << time << "us\n" - << "output: " << output[0] << "\n"; + latencies.push_back(static_cast(time) / 1000.0); } + printBenchmark("mlp", latencies); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete ctx; delete engine; delete runtime; diff --git a/mlp/utils.h b/mlp/utils.h index 5579b214..eb70a85e 100644 --- a/mlp/utils.h +++ b/mlp/utils.h @@ -1,11 +1,17 @@ #pragma once #include +#include #include +#include +#include #include +#include #include #include +#include #include #include +#include #include "macros.h" using namespace nvinfer1; @@ -22,7 +28,7 @@ constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -76,7 +82,7 @@ static auto loadWeights(const std::string& file) { static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -93,3 +99,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/mnasnet/FindTensorRT.cmake b/mnasnet/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/mnasnet/FindTensorRT.cmake +++ b/mnasnet/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/mnasnet/macros.h b/mnasnet/macros.h index daae7bfb..6d851752 100644 --- a/mnasnet/macros.h +++ b/mnasnet/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/mnasnet/mnasnet.cpp b/mnasnet/mnasnet.cpp index fcbcb3ce..486b32bc 100644 --- a/mnasnet/mnasnet.cpp +++ b/mnasnet/mnasnet.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -21,7 +22,7 @@ static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, static const std::string WTS_PATH = "../models/mnasnet0_5.wts"; static const std::string ENGINE_PATH = "../models/mnasnet0_5.engine"; static constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1) ? true : false; static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; @@ -158,9 +159,9 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IRuntime* runtime, IBuilder DataType dt) { auto weightMap = loadWeights(WTS_PATH); -#if TRT_VERSION >= 11200 +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); @@ -215,7 +216,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IRuntime* runtime, IBuilder network->markOutput(*fc1->getOutput(0)); // Build engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); auto* _serialized = builder->buildSerializedNetwork(*network, *config); auto* engine = runtime->deserializeCudaEngine(_serialized->data(), _serialized->size()); @@ -250,7 +251,7 @@ void APIToModel(unsigned int maxBatchSize, IRuntime* runtime, IHostMemory** mode (*modelStream) = engine->serialize(); // Close everything down -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -267,7 +268,7 @@ std::vector> do_inference(IExecutionContext& context, void* i CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -276,7 +277,7 @@ std::vector> do_inference(IExecutionContext& context, void* i buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); size = s * batch_size * SIZES[i]; @@ -284,7 +285,10 @@ std::vector> do_inference(IExecutionContext& context, void* i if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -297,10 +301,16 @@ std::vector> do_inference(IExecutionContext& context, void* i #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; @@ -352,7 +362,7 @@ int main(int argc, char** argv) { const auto* data_ptr = reinterpret_cast(modelStream->data()); auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); @@ -373,7 +383,7 @@ int main(int argc, char** argv) { return -1; } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* engine = runtime->deserializeCudaEngine(trt_model_stream, size); #else auto* engine = runtime->deserializeCudaEngine(trt_model_stream, size, nullptr); @@ -396,34 +406,26 @@ int main(int argc, char** argv) { input = flat_img.data(); } - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = do_inference(*context, input, 1); - auto _end = std::chrono::system_clock::now(); - auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "ms\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } - } - } + auto firstProb = do_inference(*context, input, 1); + printFirstOutputs("mnasnet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + int _top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << labels[idx] + << "\n"; + } - if (i == 99) { - std::cout << "prediction result:\n"; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << "\n"; - } - } + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)do_inference(*context, input, 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("mnasnet", latencies); delete[] trt_model_stream; return 0; diff --git a/mnasnet/utils.h b/mnasnet/utils.h index 75ff771a..df8e9085 100644 --- a/mnasnet/utils.h +++ b/mnasnet/utils.h @@ -2,7 +2,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -25,7 +28,7 @@ constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -181,7 +184,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -226,7 +229,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; @@ -234,7 +237,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -251,3 +254,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/mobilenet/CMakeLists.txt b/mobilenet/CMakeLists.txt new file mode 100644 index 00000000..0670cea2 --- /dev/null +++ b/mobilenet/CMakeLists.txt @@ -0,0 +1,41 @@ +cmake_minimum_required(VERSION 3.14) + +project( + mobilenet + VERSION 0.1 + LANGUAGES C CXX CUDA) + +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES + 60 + 70 + 72 + 75 + 80 + 86 + 89) +endif() + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_INCLUDE_CURRENT_DIR TRUE) + +option(CUDA_USE_STATIC_CUDA_RUNTIME "Use static cudaruntime library" OFF) + +find_package(Threads REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(OpenCV REQUIRED) + +if(NOT TARGET TensorRT::TensorRT) + include(FindTensorRT.cmake) +endif() + +add_executable(${PROJECT_NAME} ${PROJECT_NAME}.cpp) + +target_include_directories(${PROJECT_NAME} PRIVATE ${OpenCV_INCLUDE_DIRS}) + +target_link_libraries(${PROJECT_NAME} PRIVATE Threads::Threads CUDA::cudart + TensorRT::TensorRT ${OpenCV_LIBS}) diff --git a/mobilenet/FindTensorRT.cmake b/mobilenet/FindTensorRT.cmake new file mode 100644 index 00000000..c877d7df --- /dev/null +++ b/mobilenet/FindTensorRT.cmake @@ -0,0 +1,90 @@ +cmake_minimum_required(VERSION 3.17.0) + +set(TRT_VERSION + $ENV{TRT_VERSION} + CACHE STRING + "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\"") + +function(_guess_path var_name) + set(_result "") + + foreach(path_entry IN LISTS ARGN) + if(EXISTS "${path_entry}") + list(APPEND _result "${path_entry}") + endif() + endforeach() + + set(${var_name} + "${_result}" + PARENT_SCOPE) +endfunction() + +# find TensorRT include folder +if(NOT DEFINED TensorRT_INCLUDE_DIR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + _guess_path(TensorRT_INCLUDE_DIR + "/usr/local/cuda/targets/aarch64-linux/include") + else() + _guess_path( + TensorRT_INCLUDE_DIR + "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" + "/usr/include/x86_64-linux-gnu") + endif() + message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") +endif() + +# find TensorRT library folder +if(NOT TensorRT_LIBRARY_DIR) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + _guess_path(TensorRT_LIBRARY_DIR "/usr/lib/aarch64-linux-gnu/tegra") + else() + _guess_path(TensorRT_LIBRARY_DIR "/usr/lib/x86_64-linux-gnu" + "/usr/local/tensorrt/targets/x86_64-linux-gnu/lib") + endif() + message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") +endif() + +set(TensorRT_LIBRARIES) + +message(STATUS "Found TensorRT lib: ${TensorRT_LIBRARIES}") + +# process for different TensorRT version +if(DEFINED TRT_VERSION AND NOT TRT_VERSION STREQUAL "") + string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION}) + set(TRT_MAJOR_VERSION "${_match}") + set(_modules nvinfer nvinfer_plugin) + unset(_match) + + if(TRT_MAJOR_VERSION GREATER_EQUAL 8) + list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean) + endif() +else() + message(FATAL_ERROR "Please set a environment variable \"TRT_VERSION\"") +endif() + +# find and add all modules of TensorRT into list +foreach(lib IN LISTS _modules) + find_library( + TensorRT_${lib}_LIBRARY + NAMES ${lib} + HINTS ${TensorRT_LIBRARY_DIR}) + list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY}) +endforeach() + +# make the "TensorRT target" +add_library(TensorRT IMPORTED INTERFACE) +add_library(TensorRT::TensorRT ALIAS TensorRT) +target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES}) + +set_target_properties( + TensorRT + PROPERTIES C_STANDARD 17 + CXX_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + SKIP_BUILD_RPATH TRUE + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN" + INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") + +unset(TRT_MAJOR_VERSION) +unset(_modules) diff --git a/mobilenet/gen_wts.py b/mobilenet/gen_wts.py new file mode 100644 index 00000000..1c4fce37 --- /dev/null +++ b/mobilenet/gen_wts.py @@ -0,0 +1,87 @@ +import argparse +import struct +from pathlib import Path + +import cv2 +import numpy as np +import torch +from torchvision.models import MobileNet_V2_Weights, MobileNet_V3_Small_Weights +from torchvision.models.mobilenet import mobilenet_v2, mobilenet_v3_small + +SUPPORTED_MODELS = ("mobilenet_v2", "mobilenet_v3_small") + + +def read_imagenet_labels(path: Path) -> dict[int, str]: + clsid2label: dict[int, str] = {} + with path.open("r", encoding="utf-8") as f: + for line in f.readlines(): + k, v = line.split(": ") + clsid2label.setdefault(int(k), v[1:-3]) + return clsid2label + + +def preprocess(img: np.ndarray) -> torch.Tensor: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR) + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + img = (img - mean) / std + img = img.transpose(2, 0, 1)[None, ...] + return torch.from_numpy(img) + + +def create_model(name: str) -> torch.nn.Module: + if name == "mobilenet_v2": + return mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT) + if name == "mobilenet_v3_small": + return mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.DEFAULT) + raise ValueError(f"unsupported model: {name}") + + +def export_model(name: str, output_dir: Path, labels: dict[int, str], image: torch.Tensor) -> None: + print(f"Now dealing with model: {name}") + model = create_model(name).eval() + with torch.inference_mode(): + output = model(image) + for i, batch in enumerate(torch.topk(output, k=3).indices): + for j, idx in enumerate(batch): + print(f"\tBatch: {i}, Top: {j}, logits: {output[i][idx]:.4f}, label: {labels[int(idx)]}") + print(f"{'=' * 32}") + + output_dir.mkdir(parents=True, exist_ok=True) + with (output_dir / f"{name}.wts").open("w", encoding="utf-8") as f: + f.write(f"{len(model.state_dict().keys())}\n") + for key, value in model.state_dict().items(): + values = value.reshape(-1).cpu().numpy() + f.write(f"{key} {len(values)} ") + print(key, value.shape) + for item in values: + f.write(" ") + f.write(struct.pack(">f", float(item)).hex()) + f.write("\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Export torchvision MobileNet weights to tensorrtx .wts format.") + parser.add_argument("--model", choices=SUPPORTED_MODELS, action="append", help="model to export; repeatable") + parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parents[1] / "models") + parser.add_argument("--image", type=Path, default=Path(__file__).resolve().parents[1] / "assets" / "cats.jpg") + parser.add_argument( + "--labels", + type=Path, + default=Path(__file__).resolve().parents[1] / "assets" / "imagenet1000_clsidx_to_labels.txt", + ) + args = parser.parse_args() + + labels = read_imagenet_labels(args.labels) + img = cv2.imread(str(args.image), cv2.IMREAD_COLOR) + if img is None: + raise FileNotFoundError(args.image) + image = preprocess(img) + + for name in args.model or SUPPORTED_MODELS: + export_model(name, args.output_dir, labels, image) + + +if __name__ == "__main__": + main() diff --git a/mobilenet/logging.h b/mobilenet/logging.h new file mode 100644 index 00000000..f7f955a5 --- /dev/null +++ b/mobilenet/logging.h @@ -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 +#include +#include +#include +#include +#include +#include +#include "NvInferRuntime.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(×tamp); + 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 diff --git a/mobilenet/macros.h b/mobilenet/macros.h new file mode 100644 index 00000000..4f442e6a --- /dev/null +++ b/mobilenet/macros.h @@ -0,0 +1,35 @@ +#pragma once +#include + +#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 + +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) (TRT_VERSION >= TRT_VERSION_ENCODE((major), (minor), (patch), 0)) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) + +#if TRT_VERSION_LT(7, 2, 2) +#error "TensorRT >= 7.2.2 is required for this demo." +#endif + +#if TRT_VERSION_GE(8, 0, 0) +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif diff --git a/mobilenet/mobilenet.cpp b/mobilenet/mobilenet.cpp new file mode 100644 index 00000000..a5708d22 --- /dev/null +++ b/mobilenet/mobilenet.cpp @@ -0,0 +1,634 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" +#include "utils.h" + +struct MobileNetVariant { + std::string name; + std::string wts_path; + std::string engine_path; + bool is_v3_small; +}; + +struct V2BlockParams { + int32_t index; + int32_t in_channels; + int32_t out_channels; + int32_t stride; + int32_t expansion; +}; + +struct V3BlockParams { + int32_t index; + int32_t in_channels; + int32_t out_channels; + int32_t hidden_channels; + int32_t kernel_size; + int32_t stride; + bool use_se; + bool use_hs; +}; + +static constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; +static constexpr const int64_t N = 1; +static constexpr const int32_t INPUT_H = 224; +static constexpr const int32_t INPUT_W = 224; +static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; +static constexpr const std::array NAMES = {"data", "logits"}; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1); +static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; +static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; + +using namespace nvinfer1; +using WeightMap = std::map; +using M = MatrixOperation; +using E = ElementWiseOperation; +using NDCF = NetworkDefinitionCreationFlag; + +static Logger gLogger; + +static auto getVariantConfig(const std::string& name) -> MobileNetVariant { + if (name == "mobilenet_v2" || name == "v2") { + return {"mobilenet_v2", "models/mobilenet_v2.wts", "models/mobilenet_v2.engine", false}; + } + if (name == "mobilenet_v3_small" || name == "v3_small") { + return {"mobilenet_v3_small", "models/mobilenet_v3_small.wts", "models/mobilenet_v3_small.engine", true}; + } + std::cerr << "Unsupported MobileNet variant: " << name << "\n"; + std::cerr << "Choose one of: mobilenet_v2, mobilenet_v3_small\n"; + std::abort(); +} + +static auto emptyWeights() -> Weights { + return Weights{DataType::kFLOAT, nullptr, 0ll}; +} + +static auto getWeight(const WeightMap& weight_map, const std::string& key) -> Weights { + const auto iter = weight_map.find(key); + if (iter == weight_map.end()) { + std::cerr << "Missing MobileNet weight: " << key << "\n"; + std::abort(); + } + return iter->second; +} + +static auto addBatchNorm2d(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& lname, + float eps = 1e-5f) -> ILayer* { + const auto* gamma = static_cast(getWeight(weight_map, lname + ".weight").values); + const auto* beta = static_cast(getWeight(weight_map, lname + ".bias").values); + const auto* bn_mean = static_cast(getWeight(weight_map, lname + ".running_mean").values); + const auto* var = static_cast(getWeight(weight_map, lname + ".running_var").values); + const auto len = getWeight(weight_map, lname + ".running_var").count; + + auto* scale_values = static_cast(std::malloc(sizeof(float) * static_cast(len))); + auto* shift_values = static_cast(std::malloc(sizeof(float) * static_cast(len))); + if (scale_values == nullptr || shift_values == nullptr) { + std::cerr << "batchnorm weight allocation failed\n"; + std::abort(); + } + for (int64_t i = 0; i < len; ++i) { + scale_values[i] = gamma[i] / std::sqrt(var[i] + eps); + shift_values[i] = beta[i] - bn_mean[i] * gamma[i] / std::sqrt(var[i] + eps); + } + + Weights scale{DataType::kFLOAT, scale_values, len}; + Weights shift{DataType::kFLOAT, shift_values, len}; + static const Weights power{DataType::kFLOAT, nullptr, 0ll}; + + weight_map[lname + ".scale"] = scale; + weight_map[lname + ".shift"] = shift; + weight_map[lname + ".power"] = power; + auto* scale_layer = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_layer); + return scale_layer; +} + +static auto addHardSigmoid(INetworkDefinition* network, ITensor& input, const std::string& name) -> IActivationLayer* { + auto* hsigmoid = network->addActivation(input, ActivationType::kHARD_SIGMOID); + assert(hsigmoid); + hsigmoid->setAlpha(1.0f / 6.0f); + hsigmoid->setBeta(0.5f); + hsigmoid->setName(name.c_str()); + return hsigmoid; +} + +static auto addHardSwish(INetworkDefinition* network, ITensor& input, const std::string& name) -> ILayer* { + auto* hsigmoid = addHardSigmoid(network, input, name + ".hsigmoid"); + auto* hswish = network->addElementWise(input, *hsigmoid->getOutput(0), E::kPROD); + assert(hswish); + hswish->setName(name.c_str()); + return hswish; +} + +static auto addRelu6(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const std::string& name) -> ILayer* { + auto* relu = network->addActivation(input, ActivationType::kRELU); + assert(relu); + relu->setName((name + ".relu").c_str()); + + auto* six = static_cast(std::malloc(sizeof(float))); + if (six == nullptr) { + std::cerr << "relu6 constant allocation failed\n"; + std::abort(); + } + *six = 6.0f; + const std::string key = name + ".relu6.max"; + weight_map[key] = Weights{DataType::kFLOAT, six, 1ll}; + auto* c = network->addConstant(Dims4{1, 1, 1, 1}, weight_map[key]); + assert(c); + auto* clipped = network->addElementWise(*relu->getOutput(0), *c->getOutput(0), E::kMIN); + assert(clipped); + clipped->setName(name.c_str()); + return clipped; +} + +static auto addActivation(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& name, + bool hard_swish, bool relu6) -> ILayer* { + if (hard_swish) { + return addHardSwish(network, input, name); + } + if (relu6) { + return addRelu6(network, weight_map, input, name); + } + auto* relu = network->addActivation(input, ActivationType::kRELU); + assert(relu); + relu->setName(name.c_str()); + return relu; +} + +static auto addConvBnAct(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const std::string& conv_key, const std::string& bn_key, int32_t out_channels, + int32_t kernel_size, int32_t stride, int32_t padding, int32_t groups, bool with_act, + bool hard_swish = false, bool relu6 = false, float eps = 1e-5f) -> ILayer* { + auto* conv = network->addConvolutionNd(input, out_channels, DimsHW{kernel_size, kernel_size}, + getWeight(weight_map, conv_key), emptyWeights()); + assert(conv); + conv->setStrideNd(DimsHW{stride, stride}); + conv->setPaddingNd(DimsHW{padding, padding}); + conv->setNbGroups(groups); + conv->setName(conv_key.c_str()); + + auto* bn = addBatchNorm2d(network, weight_map, *conv->getOutput(0), bn_key, eps); + bn->setName((bn_key + ".bn").c_str()); + if (!with_act) { + return bn; + } + return addActivation(network, weight_map, *bn->getOutput(0), bn_key + ".act", hard_swish, relu6); +} + +static auto addLinear(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& weight_key, + const std::string& bias_key, int32_t in_features, int32_t out_features, + const std::string& name) -> ILayer* { + auto* fcw = network->addConstant(DimsHW{out_features, in_features}, getWeight(weight_map, weight_key)); + auto* fcb = network->addConstant(DimsHW{1, out_features}, getWeight(weight_map, bias_key)); + assert(fcw); + assert(fcb); + auto* fc_matmul = network->addMatrixMultiply(input, M::kNONE, *fcw->getOutput(0), M::kTRANSPOSE); + assert(fc_matmul); + auto* fc = network->addElementWise(*fc_matmul->getOutput(0), *fcb->getOutput(0), E::kSUM); + assert(fc); + fc->setName(name.c_str()); + return fc; +} + +static auto addSqueezeExcitation(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const std::string& lname, int32_t channels) -> ILayer* { + auto* pool = network->addReduce(input, ReduceOperation::kAVG, 0xc, true); + assert(pool); + pool->setName((lname + ".avgpool").c_str()); + + const auto fc1_weight = getWeight(weight_map, lname + ".fc1.weight"); + const auto squeeze_channels = toI32(fc1_weight.count / channels); + auto* fc1 = network->addConvolutionNd(*pool->getOutput(0), squeeze_channels, DimsHW{1, 1}, fc1_weight, + getWeight(weight_map, lname + ".fc1.bias")); + assert(fc1); + fc1->setName((lname + ".fc1").c_str()); + auto* relu = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU); + assert(relu); + relu->setName((lname + ".relu").c_str()); + auto* fc2 = network->addConvolutionNd(*relu->getOutput(0), channels, DimsHW{1, 1}, + getWeight(weight_map, lname + ".fc2.weight"), + getWeight(weight_map, lname + ".fc2.bias")); + assert(fc2); + fc2->setName((lname + ".fc2").c_str()); + auto* scale = addHardSigmoid(network, *fc2->getOutput(0), lname + ".scale"); + auto* out = network->addElementWise(input, *scale->getOutput(0), E::kPROD); + assert(out); + out->setName((lname + ".mul").c_str()); + return out; +} + +static auto addV2Block(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const V2BlockParams& params) -> ILayer* { + const std::string lname = "features." + std::to_string(params.index) + ".conv."; + const int32_t hidden = params.in_channels * params.expansion; + ILayer* layer = nullptr; + if (params.expansion != 1) { + layer = addConvBnAct(network, weight_map, input, lname + "0.0.weight", lname + "0.1", hidden, 1, 1, 0, 1, true, + false, true); + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), lname + "1.0.weight", lname + "1.1", hidden, 3, + params.stride, 1, hidden, true, false, true); + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), lname + "2.weight", lname + "3", + params.out_channels, 1, 1, 0, 1, false); + } else { + layer = addConvBnAct(network, weight_map, input, lname + "0.0.weight", lname + "0.1", hidden, 3, params.stride, + 1, hidden, true, false, true); + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), lname + "1.weight", lname + "2", + params.out_channels, 1, 1, 0, 1, false); + } + if (params.stride == 1 && params.in_channels == params.out_channels) { + auto* ew = network->addElementWise(input, *layer->getOutput(0), E::kSUM); + assert(ew); + ew->setName((lname + "residual").c_str()); + return ew; + } + return layer; +} + +static auto addV3Block(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const V3BlockParams& params) -> ILayer* { + const std::string lname = "features." + std::to_string(params.index) + ".block."; + ITensor* tensor = &input; + int32_t depthwise_index = 0; + int32_t project_index = 2; + if (params.in_channels != params.hidden_channels) { + auto* layer = addConvBnAct(network, weight_map, input, lname + "0.0.weight", lname + "0.1", + params.hidden_channels, 1, 1, 0, 1, true, params.use_hs, false, 1e-3f); + tensor = layer->getOutput(0); + depthwise_index = 1; + project_index = params.use_se ? 3 : 2; + } + + auto* layer = addConvBnAct(network, weight_map, *tensor, lname + std::to_string(depthwise_index) + ".0.weight", + lname + std::to_string(depthwise_index) + ".1", params.hidden_channels, + params.kernel_size, params.stride, params.kernel_size / 2, params.hidden_channels, true, + params.use_hs, false, 1e-3f); + if (params.use_se) { + layer = addSqueezeExcitation(network, weight_map, *layer->getOutput(0), + lname + std::to_string(depthwise_index + 1), params.hidden_channels); + } + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), lname + std::to_string(project_index) + ".0.weight", + lname + std::to_string(project_index) + ".1", params.out_channels, 1, 1, 0, 1, false, false, + false, 1e-3f); + if (params.stride == 1 && params.in_channels == params.out_channels) { + auto* ew = network->addElementWise(input, *layer->getOutput(0), E::kSUM); + assert(ew); + ew->setName((lname + "residual").c_str()); + return ew; + } + return layer; +} + +static auto buildMobileNetV2(INetworkDefinition* network, WeightMap& weight_map, ITensor& input) -> ITensor* { + auto* layer = addConvBnAct(network, weight_map, input, "features.0.0.weight", "features.0.1", 32, 3, 2, 1, 1, true, + false, true); + const std::array blocks = {{{1, 32, 16, 1, 1}, + {2, 16, 24, 2, 6}, + {3, 24, 24, 1, 6}, + {4, 24, 32, 2, 6}, + {5, 32, 32, 1, 6}, + {6, 32, 32, 1, 6}, + {7, 32, 64, 2, 6}, + {8, 64, 64, 1, 6}, + {9, 64, 64, 1, 6}, + {10, 64, 64, 1, 6}, + {11, 64, 96, 1, 6}, + {12, 96, 96, 1, 6}, + {13, 96, 96, 1, 6}, + {14, 96, 160, 2, 6}, + {15, 160, 160, 1, 6}, + {16, 160, 160, 1, 6}, + {17, 160, 320, 1, 6}}}; + for (const auto& block : blocks) { + layer = addV2Block(network, weight_map, *layer->getOutput(0), block); + } + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), "features.18.0.weight", "features.18.1", 1280, 1, 1, + 0, 1, true, false, true); + auto* pool = network->addReduce(*layer->getOutput(0), ReduceOperation::kAVG, 0xc, false); + assert(pool); + pool->setName("avgpool"); + return addLinear(network, weight_map, *pool->getOutput(0), "classifier.1.weight", "classifier.1.bias", 1280, 1000, + "classifier.1") + ->getOutput(0); +} + +static auto buildMobileNetV3Small(INetworkDefinition* network, WeightMap& weight_map, ITensor& input) -> ITensor* { + auto* layer = addConvBnAct(network, weight_map, input, "features.0.0.weight", "features.0.1", 16, 3, 2, 1, 1, true, + true, false, 1e-3f); + const std::array blocks = {{{1, 16, 16, 16, 3, 2, true, false}, + {2, 16, 24, 72, 3, 2, false, false}, + {3, 24, 24, 88, 3, 1, false, false}, + {4, 24, 40, 96, 5, 2, true, true}, + {5, 40, 40, 240, 5, 1, true, true}, + {6, 40, 40, 240, 5, 1, true, true}, + {7, 40, 48, 120, 5, 1, true, true}, + {8, 48, 48, 144, 5, 1, true, true}, + {9, 48, 96, 288, 5, 2, true, true}, + {10, 96, 96, 576, 5, 1, true, true}, + {11, 96, 96, 576, 5, 1, true, true}}}; + for (const auto& block : blocks) { + layer = addV3Block(network, weight_map, *layer->getOutput(0), block); + } + layer = addConvBnAct(network, weight_map, *layer->getOutput(0), "features.12.0.weight", "features.12.1", 576, 1, 1, + 0, 1, true, true, false, 1e-3f); + auto* pool = network->addReduce(*layer->getOutput(0), ReduceOperation::kAVG, 0xc, false); + assert(pool); + pool->setName("avgpool"); + layer = addLinear(network, weight_map, *pool->getOutput(0), "classifier.0.weight", "classifier.0.bias", 576, 1024, + "classifier.0"); + layer = addHardSwish(network, *layer->getOutput(0), "classifier.1"); + return addLinear(network, weight_map, *layer->getOutput(0), "classifier.3.weight", "classifier.3.bias", 1024, 1000, + "classifier.3") + ->getOutput(0); +} + +static auto createEngine(int32_t batch_size, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt, + const MobileNetVariant& variant) -> ICudaEngine* { + WeightMap weight_map = loadWeights(variant.wts_path); + +#if TRT_VERSION_GE(10, 12, 0) + auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); +#elif TRT_VERSION_GE(10, 0, 0) + auto flag = 0U; +#else + auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); +#endif + auto* network = builder->createNetworkV2(flag); + assert(network); + + ITensor* input = nullptr; + if constexpr (TRT_PREPROCESS) { + dt = DataType::kUINT8; + input = network->addInput(NAMES[0], dt, Dims4{batch_size, INPUT_H, INPUT_W, 3}); + auto* transform = addTransformLayer(network, *input, true, mean, stdv); + input = transform->getOutput(0); + } else { + input = network->addInput(NAMES[0], dt, Dims4{batch_size, 3, INPUT_H, INPUT_W}); + } + assert(input); + + ITensor* logits = variant.is_v3_small ? buildMobileNetV3Small(network, weight_map, *input) + : buildMobileNetV2(network, weight_map, *input); + logits->setName(NAMES[1]); + network->markOutput(*logits); + +#if TRT_VERSION_GE(8, 0, 0) + config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); + IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); + assert(mem); + ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); + delete mem; + delete network; +#else + builder->setMaxBatchSize(batch_size); + config->setMaxWorkspaceSize(WORKSPACE_SIZE); + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + network->destroy(); +#endif + std::cout << "build finished\n"; + + for (auto& mem : weight_map) { + std::free(const_cast(mem.second.values)); + } + return engine; +} + +static void APIToModel(int32_t batch_size, IRuntime* runtime, IHostMemory** model_stream, + const MobileNetVariant& variant) { + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + ICudaEngine* engine = createEngine(batch_size, runtime, builder, config, DataType::kFLOAT, variant); + assert(engine != nullptr); + + (*model_stream) = engine->serialize(); + +#if TRT_VERSION_GE(8, 0, 0) + delete engine; + delete config; + delete builder; +#else + engine->destroy(); + config->destroy(); + builder->destroy(); +#endif +} + +static auto doInference(IExecutionContext& context, void* input, + int64_t batch_size) -> std::vector> { + ICudaEngine const& engine = context.getEngine(); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + std::vector buffers; + +#if TRT_VERSION_GE(8, 0, 0) + const int32_t nIO = engine.getNbIOTensors(); +#else + const int32_t nIO = engine.getNbBindings(); +#endif + + buffers.resize(nIO); + for (auto i = 0; i < nIO; ++i) { + std::size_t size = 0; +#if TRT_VERSION_GE(8, 0, 0) + auto* tensor_name = engine.getIOTensorName(i); + const std::string name = tensor_name; + auto element_size = getSize(engine.getTensorDataType(tensor_name)); + size = element_size * static_cast(batch_size) * (name == NAMES[0] ? SIZES[0] : SIZES[1]); + CHECK(cudaMalloc(&buffers[i], size)); + if (name == NAMES[0]) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } +#else + const int32_t idx = engine.getBindingIndex(NAMES[i]); + auto element_size = getSize(engine.getBindingDataType(idx)); + assert(idx == i); + size = element_size * static_cast(batch_size) * SIZES[i]; + CHECK(cudaMalloc(&buffers[i], size)); + if (i == 0) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } +#endif + } + +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } +#else + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } +#endif + + std::vector> prob; + for (int i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) + const std::string name = engine.getIOTensorName(i); + if (name == NAMES[0]) { + continue; + } + constexpr auto output_size = SIZES[1]; +#else + if (i == 0) { + continue; + } + const auto output_size = SIZES[i]; +#endif + std::vector tmp(batch_size * output_size, std::nanf("")); + const auto size = static_cast(batch_size) * output_size * sizeof(float); + CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); + prob.emplace_back(std::move(tmp)); + } + CHECK(cudaStreamSynchronize(stream)); + + for (auto& buffer : buffers) { + CHECK(cudaFree(buffer)); + } + CHECK(cudaStreamDestroy(stream)); + return prob; +} + +auto main(int argc, char** argv) -> int { + checkTrtEnv(); + if (argc < 2 || argc > 3) { + std::cerr << "arguments not right!\n"; + std::cerr << "./mobilenet -s [model] // serialize model to plan file\n"; + std::cerr << "./mobilenet -d [model] // deserialize plan file and run inference\n"; + std::cerr << "model: mobilenet_v2 | mobilenet_v3_small\n"; + return -1; + } + + const auto variant = getVariantConfig(argc == 3 ? argv[2] : "mobilenet_v2"); + std::cout << "Using MobileNet variant: " << variant.name << "\n"; + + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + char* trt_model_stream{nullptr}; + std::streamsize size{0}; + + if (std::string(argv[1]) == "-s") { + IHostMemory* model_stream{nullptr}; + APIToModel(N, runtime, &model_stream, variant); + assert(model_stream != nullptr); + + std::ofstream plan(variant.engine_path, std::ios::binary | std::ios::trunc); + if (!plan) { + std::cerr << "could not open plan output file\n"; + return -1; + } + if (model_stream->size() > static_cast(std::numeric_limits::max())) { + std::cerr << "this model is too large to serialize\n"; + return -1; + } + const auto* data_ptr = reinterpret_cast(model_stream->data()); + const auto data_size = static_cast(model_stream->size()); + plan.write(data_ptr, data_size); +#if TRT_VERSION_GE(8, 0, 0) + delete model_stream; + delete runtime; +#else + model_stream->destroy(); + runtime->destroy(); +#endif + return 0; + } + if (std::string(argv[1]) == "-d") { + std::ifstream file(variant.engine_path, std::ios::binary); + if (file.good()) { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trt_model_stream = new char[size]; + assert(trt_model_stream); + file.read(trt_model_stream, size); + file.close(); + } + } else { + return -1; + } + +#if TRT_VERSION_GE(8, 0, 0) + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size); +#else + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size, nullptr); +#endif + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + delete[] trt_model_stream; + + void* input = nullptr; + std::vector flat_img; + cv::Mat img = cv::imread("assets/cats.jpg", cv::IMREAD_COLOR); + if (img.empty()) { + std::cerr << "failed to read image: assets/cats.jpg\n"; + return -1; + } + if constexpr (TRT_PREPROCESS) { + cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR); + input = static_cast(img.data); + } else { + flat_img = preprocess_img(img, true, mean, stdv, N, INPUT_H, INPUT_W); + input = flat_img.data(); + } + + auto firstProb = doInference(*context, input, N); + printFirstOutputs("mobilenet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + if (labels.empty()) { + std::cerr << "failed to load labels from " << LABELS_PATH << "\n"; + std::abort(); + } + int top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << top++ << " idx: " << idx << ", logits: " << std::setprecision(4) << logits + << ", label: " << labels[idx] << "\n"; + } + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int i = 0; i < kBenchmarkRuns; ++i) { + auto start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, N); + auto end = std::chrono::steady_clock::now(); + auto period = std::chrono::duration_cast(end - start); + latencies.push_back(static_cast(period.count()) / 1000.0); + } + printBenchmark("mobilenet", latencies, N); +#if TRT_VERSION_GE(8, 0, 0) + delete context; + delete engine; + delete runtime; +#else + context->destroy(); + engine->destroy(); + runtime->destroy(); +#endif + + return 0; +} diff --git a/mobilenet/utils.h b/mobilenet/utils.h new file mode 100644 index 00000000..4f111957 --- /dev/null +++ b/mobilenet/utils.h @@ -0,0 +1,311 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "macros.h" + +using namespace nvinfer1; + +#define CHECK(status) \ + do { \ + auto ret = (status); \ + if (ret != cudaSuccess) { \ + std::cerr << "Cuda failure: " << ret << "\n"; \ + std::abort(); \ + } \ + } while (0) + +static void checkTrtEnv(int device = 0) { +#if TRT_VERSION_LT(8, 0, 0) + CHECK(cudaGetDevice(&device)); + cudaDeviceProp prop{}; + CHECK(cudaGetDeviceProperties(&prop, device)); + const int sm = prop.major * 10 + prop.minor; + if (sm > 86) { + std::cerr << "TensorRT < 8 does not support SM > 86 on this GPU."; + std::abort(); + } +#endif +} + +template +static auto toI32(T value) -> int32_t { + static_assert(std::is_integral::value, "toI32 requires an integral type."); + static_assert(sizeof(T) > sizeof(int32_t), "toI32 is only for values wider than int32_t."); + if constexpr (std::is_signed::value) { + assert(value >= static_cast(std::numeric_limits::min())); + } + assert(value <= static_cast(std::numeric_limits::max())); + return static_cast(value); +} + +/** + * @brief TensorRT weight files have a simple space delimited format: + * [type] [size] + * + * @param file input weight file path + * @return std::map + */ +static std::map loadWeights(const std::string& file) { + std::cout << "Loading weights: " << file << "\n"; + std::map 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--) { + nvinfer1::Weights wt{nvinfer1::DataType::kFLOAT, nullptr, 0}; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> wt.count; + + // Load blob + auto* val = static_cast(std::malloc(sizeof(uint32_t) * static_cast(wt.count))); + if (val == nullptr) { + std::cerr << "weight allocation failed\n"; + std::abort(); + } + input >> std::hex; + for (auto x = 0ll; x < wt.count; ++x) { + input >> val[x]; + } + wt.values = val; + weightMap[name] = wt; + } + + return weightMap; +} + +/** + * @brief a preprocess function aligning with ImageNet preprocess in torchvision, only support 3-channel image + * + * @param img opencv image with BGR layout + * @param bgr2rgb whether to convert BGR to RGB + * @param mean subtract mean + * @param std divide std + * @param n batch size + * @param h resize height + * @param w resize width + * @return std::vector contiguous flatten image data in float32 type + */ +static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const std::array& mean, + const std::array& std, int n, int h, int w) { + const auto c = img.channels(); + const auto size = c * h * w; + if (c != 3) { + std::cerr << "this demo only supports 3 channel input image.\n"; + std::abort(); + } + if (bgr2rgb) { + cv::cvtColor(img, img, cv::COLOR_BGR2RGB); + } + cv::resize(img, img, cv::Size(w, h), 0, 0, cv::INTER_LINEAR); + img.convertTo(img, CV_32FC3, 1.f / 255); + img = (img - cv::Scalar(mean[0], mean[1], mean[2])) / cv::Scalar(std[0], std[1], std[2]); + std::vector chw(static_cast(n) * c * h * w, 0.f); + + // fill all batch with the same input image + for (int i = 0; i < n; ++i) { + for (int y = 0; y < h; ++y) { + for (int x = 0; x < w; ++x) { + const cv::Vec3f v = img.at(y, x); + chw[i * size + 0 * h * w + y * w + x] = v[0]; + chw[i * size + 1 * h * w + y * w + x] = v[1]; + chw[i * size + 2 * h * w + y * w + x] = v[2]; + } + } + } + return chw; +} + +static auto topk(const std::vector& v, int k) -> std::vector> { + if (k <= 0) + return {}; + auto stride = std::min(k, static_cast(v.size())); + + std::vector idx(v.size()); + std::iota(idx.begin(), idx.end(), 0); + + std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&](int a, int b) { return v[a] > v[b]; }); + + std::vector> out; + out.reserve(stride); + for (auto i = 0; i < stride; ++i) + out.emplace_back(idx[i], v[idx[i]]); + return out; +} + +static std::map loadImagenetLabelMap(const std::string& path) { + std::map labels; + std::ifstream in(path); + if (!in.is_open()) { + return labels; + } + std::string line; + while (std::getline(in, line)) { + auto colon = line.find(':'); + if (colon == std::string::npos) { + continue; + } + auto first_quote = line.find('\'', colon); + if (first_quote == std::string::npos) { + continue; + } + auto second_quote = line.find('\'', first_quote + 1); + if (second_quote == std::string::npos) { + continue; + } + int idx = std::stoi(line.substr(0, colon)); + labels[idx] = line.substr(first_quote + 1, second_quote - first_quote - 1); + } + return labels; +} + +static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bool bgr2rgb, + const std::array& mean, const std::array& std) { + struct ScaleParams { + std::array shift; + std::array scale; + }; + static std::vector> gScaleParams; + auto params = std::make_unique(); + params->shift = {-mean[0] / std[0], -mean[1] / std[1], -mean[2] / std[2]}; + params->scale = {1.f / (std[0] * 255.f), 1.f / (std[1] * 255.f), 1.f / (std[2] * 255.f)}; + + static const Weights empty{DataType::kFLOAT, nullptr, 0ll}; + const Weights shift{DataType::kFLOAT, params->shift.data(), 3ll}; + const Weights scale{DataType::kFLOAT, params->scale.data(), 3ll}; + + gScaleParams.emplace_back(std::move(params)); + + ITensor* in = &input; + if (input.getType() != DataType::kFLOAT) { +#if TRT_VERSION_GE(8, 0, 0) + auto* cast = network->addCast(input, DataType::kFLOAT); + assert(cast); + cast->setName("Cast to FP32"); + in = cast->getOutput(0); +#else + auto* identity = network->addIdentity(input); + assert(identity); + identity->setName("Convert to FP32"); + identity->setOutputType(0, DataType::kFLOAT); + in = identity->getOutput(0); +#endif + } + // Convert from NHWC to NCHW + auto* perm = network->addShuffle(*in); + assert(perm); + perm->setName("NHWC -> NCHW"); + perm->setFirstTranspose(Permutation{0, 3, 1, 2}); + + // Convert from BGR to RGB (optional) + ITensor* data{nullptr}; + if (bgr2rgb) { + auto add_slice = [&](int c, const char* name) -> ITensor* { + auto dims = perm->getOutput(0)->getDimensions(); + Dims4 start = {0, c, 0, 0}, stride = {1, 1, 1, 1}; + Dims4 size = {dims.d[0], 1, dims.d[2], dims.d[3]}; + auto* _slice = network->addSlice(*perm->getOutput(0), start, size, stride); + _slice->setName(name); + assert(_slice && _slice->getNbOutputs() == 1); + return _slice->getOutput(0); + }; + std::array channels = {add_slice(2, "R"), add_slice(1, "G"), add_slice(0, "B")}; + auto* cat = network->addConcatenation(channels.data(), 3); + assert(cat); + cat->setName("RGB"); + cat->setAxis(1); + data = cat->getOutput(0); + } else { + data = perm->getOutput(0); + } + + // Normalize + auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); + assert(trans); + trans->setName("mean & std"); +#if TRT_VERSION_GE(8, 0, 0) + trans->setChannelAxis(1); +#endif + return trans; +} + +static size_t getSize(DataType dt) { + switch (dt) { +#if TRT_VERSION_GE(8, 5, 1) + case DataType::kUINT8: +#endif + case DataType::kINT8: + return sizeof(int8_t); + case DataType::kFLOAT: + return sizeof(float); + case DataType::kHALF: + return sizeof(int16_t); + case DataType::kINT32: + return sizeof(int32_t); + default: { + std::cerr << "Unsupported data type\n"; + std::abort(); + } + } +} + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/resnet/CMakeLists.txt b/resnet/CMakeLists.txt index f7356098..6fa6b3ca 100644 --- a/resnet/CMakeLists.txt +++ b/resnet/CMakeLists.txt @@ -1,41 +1,42 @@ -cmake_minimum_required(VERSION 2.6) - -project(resnet) - -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(resnet18 ${PROJECT_SOURCE_DIR}/resnet18.cpp) -target_link_libraries(resnet18 nvinfer) -target_link_libraries(resnet18 cudart) - -add_executable(resnet34 ${PROJECT_SOURCE_DIR}/resnet34.cpp) -target_link_libraries(resnet34 nvinfer) -target_link_libraries(resnet34 cudart) - -add_executable(resnet50 ${PROJECT_SOURCE_DIR}/resnet50.cpp) -target_link_libraries(resnet50 nvinfer) -target_link_libraries(resnet50 cudart) - -add_executable(resnext50 ${PROJECT_SOURCE_DIR}/resnext50_32x4d.cpp) -target_link_libraries(resnext50 nvinfer) -target_link_libraries(resnext50 cudart) - -add_executable(wideresnet50 ${PROJECT_SOURCE_DIR}/wideresnet50.cpp) -target_link_libraries(wideresnet50 nvinfer) -target_link_libraries(wideresnet50 cudart) - -add_definitions(-O2 -pthread) - +cmake_minimum_required(VERSION 3.14) + +project( + resnet + VERSION 0.1 + LANGUAGES C CXX CUDA) + +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES + 60 + 70 + 72 + 75 + 80 + 86 + 89) +endif() + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_INCLUDE_CURRENT_DIR TRUE) + +option(CUDA_USE_STATIC_CUDA_RUNTIME "Use static cudaruntime library" OFF) + +find_package(Threads REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(OpenCV REQUIRED) + +if(NOT TARGET TensorRT::TensorRT) + include(FindTensorRT.cmake) +else() + message("TensorRT has been found, skipping for ${PROJECT_NAME}") +endif() + +add_executable(${PROJECT_NAME} ${PROJECT_NAME}.cpp) + +target_include_directories(${PROJECT_NAME} PRIVATE ${CMAKE_CURRENT_LIST_DIR} ${OpenCV_INCLUDE_DIRS}) + +target_link_libraries(${PROJECT_NAME} PRIVATE Threads::Threads CUDA::cudart TensorRT::TensorRT ${OpenCV_LIBS}) diff --git a/resnet/FindTensorRT.cmake b/resnet/FindTensorRT.cmake new file mode 100644 index 00000000..61203d17 --- /dev/null +++ b/resnet/FindTensorRT.cmake @@ -0,0 +1,149 @@ +cmake_minimum_required(VERSION 3.17.0) + +function(_guess_path var_name required_files) + set(_result "") + + foreach(path_entry IN LISTS ARGN) + if(NOT EXISTS "${path_entry}") + message(DEBUG "skip non-existing path '${path_entry}'") + continue() + endif() + + set(_ok TRUE) + foreach(required_file IN LISTS required_files) + if(NOT EXISTS "${path_entry}/${required_file}") + set(_ok FALSE) + message(DEBUG "'${path_entry}' missing '${required_file}'") + break() + endif() + endforeach() + + if(_ok) + list(APPEND _result "${path_entry}") + message(DEBUG "accept '${path_entry}'") + else() + message(DEBUG "reject '${path_entry}'") + endif() + endforeach() + + if(_result STREQUAL "") + message( + FATAL_ERROR + "_guess_path(${var_name}) failed: no valid path found. required_files='${required_files}' candidates='${ARGN}'" + ) + endif() + + set(${var_name} + "${_result}" + PARENT_SCOPE) +endfunction() + +# add library +add_library(TensorRT IMPORTED INTERFACE) +add_library(TensorRT::TensorRT ALIAS TensorRT) + +set(TRT_VERSION + CACHE + STRING + "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", \"8.6.1.6.Windows10.x86_64.cuda-12.0\" etc" +) + +if(NOT TRT_VERSION STREQUAL "" AND NOT $ENV{TRT_VERSION} STREQUAL "") + message( + WARNING + "TRT_VERSION defined by cmake and environment variable both, using the later one" + ) +endif() + +if(NOT $ENV{TRT_VERSION} STREQUAL "") + set(TRT_VERSION $ENV{TRT_VERSION}) +endif() + +string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION}) +set(TRT_MAJOR_VERSION "${_match}") +unset(_match) + +if(WIN32) + set(TensorRT_DIR "C:/Program Files/TensorRT-${TRT_VERSION}") + if(NOT EXISTS "${TensorRT_DIR}") + message( + FATAL_ERROR + "TensorRT_DIR=${TensorRT_DIR} does not exist!" + ) + endif() + + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) + message(DEBUG "Using ${_modules}") + else() + set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch + nvinfer_lean) + endif() + + set(TensorRT_LIBRARY_DIR "${TensorRT_DIR}/lib") + set(TensorRT_INCLUDE_DIR "${TensorRT_DIR}/include") +elseif(UNIX) + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _trt_arch) + set(_trt_include_candidates) + if(_trt_arch MATCHES "^(aarch64|arm64|arch64)$") + set(_trt_include_candidates "/usr/include/aarch64-linux-gnu" "/usr/include" + "/usr/local/cuda/targets/aarch64-linux/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/aarch64-linux-gnu/lib" + "/usr/lib/aarch64-linux-gnu" "/usr/lib/aarch64-linux-gnu/tegra" + "/usr/lib") + elseif(_trt_arch MATCHES "^(x86_64|amd64)$") + set(_trt_include_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" + "/usr/include/x86_64-linux-gnu" "/usr/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" + "/usr/lib/x86_64-linux-gnu" "/usr/lib") + else() + message(FATAL_ERROR "Unknown architecture") + endif() + + set(_modules nvinfer nvinfer_plugin) + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 8) + list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean) + endif() + + _guess_path(TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so" + ${_trt_library_candidates}) + message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") + _guess_path(TensorRT_INCLUDE_DIR "NvInfer.h" ${_trt_include_candidates}) + message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") +endif() + +foreach(lib IN LISTS _modules) + find_library( + TensorRT_${lib}_LIBRARY + NAMES ${lib} + HINTS ${TensorRT_LIBRARY_DIR}) + list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY}) +endforeach() + +target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES}) + +message(STATUS "Found TensorRT libs: ${TensorRT_LIBRARIES}") + +set_target_properties( + TensorRT + PROPERTIES C_STANDARD 17 + CXX_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + SKIP_BUILD_RPATH TRUE + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN" + INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") + +unset(TRT_MAJOR_VERSION) +unset(_modules) +unset(_trt_lib_suffix) +unset(_trt_include_candidates) +unset(_trt_library_candidates) +unset(_trt_arch) diff --git a/resnet/gen_wts.py b/resnet/gen_wts.py new file mode 100644 index 00000000..d9e78703 --- /dev/null +++ b/resnet/gen_wts.py @@ -0,0 +1,107 @@ +import argparse +import struct +from pathlib import Path + +import cv2 +import numpy as np +import torch +from torchvision.models import ( + ResNeXt50_32X4D_Weights, + ResNet18_Weights, + ResNet34_Weights, + ResNet50_Weights, + Wide_ResNet50_2_Weights, + resnet18, + resnet34, + resnet50, + resnext50_32x4d, + wide_resnet50_2, +) + +SUPPORTED_MODELS = ("resnet18", "resnet34", "resnet50", "resnext50_32x4d", "wide_resnet50_2") + + +def read_imagenet_labels(path: Path) -> dict[int, str]: + labels: dict[int, str] = {} + with path.open("r", encoding="utf-8") as file: + for line in file: + key, value = line.split(": ", maxsplit=1) + labels[int(key)] = value.strip()[1:-2] + return labels + + +def preprocess(image_path: Path) -> torch.Tensor: + image = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + if image is None: + raise FileNotFoundError(f"failed to read image: {image_path}") + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_LINEAR) + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + image = (image - mean) / std + image = image.transpose(2, 0, 1)[None, ...] + return torch.from_numpy(image) + + +def create_model(name: str) -> torch.nn.Module: + if name == "resnet18": + return resnet18(weights=ResNet18_Weights.DEFAULT) + if name == "resnet34": + return resnet34(weights=ResNet34_Weights.DEFAULT) + if name == "resnet50": + return resnet50(weights=ResNet50_Weights.DEFAULT) + if name == "resnext50_32x4d": + return resnext50_32x4d(weights=ResNeXt50_32X4D_Weights.DEFAULT) + if name == "wide_resnet50_2": + return wide_resnet50_2(weights=Wide_ResNet50_2_Weights.DEFAULT) + raise ValueError(f"unsupported model: {name}") + + +def export_model(name: str, output_dir: Path, labels: dict[int, str], image: torch.Tensor) -> None: + print(f"Now dealing with model: {name}") + model = create_model(name).eval() + with torch.inference_mode(): + output = model(image) + for batch_idx, batch in enumerate(torch.topk(output, k=3).indices): + for top_idx, label_idx in enumerate(batch): + print( + f"\tBatch: {batch_idx}, Top: {top_idx}, " + f"logits: {output[batch_idx][label_idx]:.4f}, label: {labels[int(label_idx)]}" + ) + print("=" * 32) + + output_dir.mkdir(parents=True, exist_ok=True) + with (output_dir / f"{name}.wts").open("w", encoding="utf-8") as file: + file.write(f"{len(model.state_dict().keys())}\n") + for key, value in model.state_dict().items(): + values = value.reshape(-1).cpu().numpy() + file.write(f"{key} {len(values)} ") + print(key, value.shape) + for item in values: + file.write(" ") + file.write(struct.pack(">f", float(item)).hex()) + file.write("\n") + + +def parse_args() -> argparse.Namespace: + repo_root = Path(__file__).resolve().parents[1] + parser = argparse.ArgumentParser(description="Export torchvision ResNet weights to TensorRT .wts files.") + parser.add_argument("--model", choices=SUPPORTED_MODELS, action="append", help="model to export; repeatable") + parser.add_argument("--output-dir", type=Path, default=repo_root / "models") + parser.add_argument("--image", type=Path, default=repo_root / "assets" / "cats.jpg") + parser.add_argument( + "--labels", type=Path, default=repo_root / "assets" / "imagenet1000_clsidx_to_labels.txt" + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + labels = read_imagenet_labels(args.labels) + image = preprocess(args.image) + for name in args.model or SUPPORTED_MODELS: + export_model(name, args.output_dir, labels, image) + + +if __name__ == "__main__": + main() diff --git a/resnet/logging.h b/resnet/logging.h index 602b69fb..f7f955a5 100644 --- a/resnet/logging.h +++ b/resnet/logging.h @@ -17,7 +17,6 @@ #ifndef TENSORRT_LOGGING_H #define TENSORRT_LOGGING_H -#include "NvInferRuntimeCommon.h" #include #include #include @@ -25,32 +24,24 @@ #include #include #include +#include "NvInferRuntime.h" +#include "macros.h" using Severity = nvinfer1::ILogger::Severity; -class LogStreamConsumerBuffer : public std::stringbuf -{ -public: +class LogStreamConsumerBuffer : public std::stringbuf { + public: LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) - : mOutput(stream) - , mPrefix(prefix) - , mShouldLog(shouldLog) - { - } + : mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {} - LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) - : mOutput(other.mOutput) - { - } + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) : mOutput(other.mOutput) {} - ~LogStreamConsumerBuffer() - { + ~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()) - { + if (pbase() != pptr()) { putOutput(); } } @@ -58,16 +49,13 @@ class LogStreamConsumerBuffer : public std::stringbuf // 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() - { + virtual int sync() { putOutput(); return 0; } - void putOutput() - { - if (mShouldLog) - { + void putOutput() { + if (mShouldLog) { // prepend timestamp std::time_t timestamp = std::time(nullptr); tm* tm_local = std::localtime(×tamp); @@ -88,12 +76,9 @@ class LogStreamConsumerBuffer : public std::stringbuf } } - void setShouldLog(bool shouldLog) - { - mShouldLog = shouldLog; - } + void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; } -private: + private: std::ostream& mOutput; std::string mPrefix; bool mShouldLog; @@ -103,15 +88,12 @@ class LogStreamConsumerBuffer : public std::stringbuf //! \class LogStreamConsumerBase //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer //! -class LogStreamConsumerBase -{ -public: +class LogStreamConsumerBase { + public: LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) - : mBuffer(stream, prefix, shouldLog) - { - } + : mBuffer(stream, prefix, shouldLog) {} -protected: + protected: LogStreamConsumerBuffer mBuffer; }; @@ -124,49 +106,49 @@ class LogStreamConsumerBase //! 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: +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) - { - } + : 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) - { - } + : 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) - { + void setReportableSeverity(Severity reportableSeverity) { mShouldLog = mSeverity <= reportableSeverity; mBuffer.setShouldLog(mShouldLog); } -private: - static std::ostream& severityOstream(Severity severity) - { + 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 ""; + 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 ""; } } @@ -198,24 +180,19 @@ class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream //! 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) - { - } +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 + enum class TestResult { + kRUNNING, //!< The test is running + kPASSED, //!< The test passed + kFAILED, //!< The test failed + kWAIVED //!< The test was waived }; //! @@ -225,10 +202,7 @@ class Logger : public nvinfer1::ILogger //! 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; - } + nvinfer1::ILogger& getTRTLogger() { return *this; } //! //! \brief Implementation of the nvinfer1::ILogger::log() virtual method @@ -236,8 +210,7 @@ class Logger : public nvinfer1::ILogger //! 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 - { + void log(Severity severity, const char* msg) TRT_NOEXCEPT override { LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; } @@ -246,10 +219,7 @@ class Logger : public nvinfer1::ILogger //! //! \param severity The logger will only emit messages that have severity of this level or higher. //! - void setReportableSeverity(Severity severity) - { - mReportableSeverity = severity; - } + void setReportableSeverity(Severity severity) { mReportableSeverity = severity; } //! //! \brief Opaque handle that holds logging information for a particular test @@ -258,20 +228,15 @@ class Logger : public nvinfer1::ILogger //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used //! with Logger::reportTest{Start,End}(). //! - class TestAtom - { - public: + class TestAtom { + public: TestAtom(TestAtom&&) = default; - private: + private: friend class Logger; TestAtom(bool started, const std::string& name, const std::string& cmdline) - : mStarted(started) - , mName(name) - , mCmdline(cmdline) - { - } + : mStarted(started), mName(name), mCmdline(cmdline) {} bool mStarted; std::string mName; @@ -289,8 +254,7 @@ class Logger : public nvinfer1::ILogger // //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). //! - static TestAtom defineTest(const std::string& name, const std::string& cmdline) - { + static TestAtom defineTest(const std::string& name, const std::string& cmdline) { return TestAtom(false, name, cmdline); } @@ -303,8 +267,7 @@ class Logger : public nvinfer1::ILogger //! \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) - { + static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) { auto cmdline = genCmdlineString(argc, argv); return defineTest(name, cmdline); } @@ -316,8 +279,7 @@ class Logger : public nvinfer1::ILogger //! //! \param[in] testAtom The handle to the test that has started //! - static void reportTestStart(TestAtom& testAtom) - { + static void reportTestStart(TestAtom& testAtom) { reportTestResult(testAtom, TestResult::kRUNNING); assert(!testAtom.mStarted); testAtom.mStarted = true; @@ -332,86 +294,85 @@ class Logger : public nvinfer1::ILogger //! \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) - { + static void reportTestEnd(const TestAtom& testAtom, TestResult result) { assert(result != TestResult::kRUNNING); assert(testAtom.mStarted); reportTestResult(testAtom, result); } - static int reportPass(const TestAtom& testAtom) - { + static int reportPass(const TestAtom& testAtom) { reportTestEnd(testAtom, TestResult::kPASSED); return EXIT_SUCCESS; } - static int reportFail(const TestAtom& testAtom) - { + static int reportFail(const TestAtom& testAtom) { reportTestEnd(testAtom, TestResult::kFAILED); return EXIT_FAILURE; } - static int reportWaive(const TestAtom& testAtom) - { + static int reportWaive(const TestAtom& testAtom) { reportTestEnd(testAtom, TestResult::kWAIVED); return EXIT_SUCCESS; } - static int reportTest(const TestAtom& testAtom, bool pass) - { + static int reportTest(const TestAtom& testAtom, bool pass) { return pass ? reportPass(testAtom) : reportFail(testAtom); } - Severity getReportableSeverity() const - { - return mReportableSeverity; - } + Severity getReportableSeverity() const { return mReportableSeverity; } -private: + 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 ""; + 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 ""; + 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) - { + 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) - { + static void reportTestResult(const TestAtom& testAtom, TestResult result) { severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " << testAtom.mCmdline << std::endl; } @@ -419,11 +380,9 @@ class Logger : public nvinfer1::ILogger //! //! \brief generate a command line string from the given (argc, argv) values //! - static std::string genCmdlineString(int argc, char const* const* argv) - { + static std::string genCmdlineString(int argc, char const* const* argv) { std::stringstream ss; - for (int i = 0; i < argc; i++) - { + for (int i = 0; i < argc; i++) { if (i > 0) ss << " "; ss << argv[i]; @@ -434,8 +393,7 @@ class Logger : public nvinfer1::ILogger Severity mReportableSeverity; }; -namespace -{ +namespace { //! //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE @@ -444,8 +402,7 @@ namespace //! //! LOG_VERBOSE(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) -{ +inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) { return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); } @@ -456,8 +413,7 @@ inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) //! //! LOG_INFO(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_INFO(const Logger& logger) -{ +inline LogStreamConsumer LOG_INFO(const Logger& logger) { return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); } @@ -468,8 +424,7 @@ inline LogStreamConsumer LOG_INFO(const Logger& logger) //! //! LOG_WARN(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_WARN(const Logger& logger) -{ +inline LogStreamConsumer LOG_WARN(const Logger& logger) { return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); } @@ -480,8 +435,7 @@ inline LogStreamConsumer LOG_WARN(const Logger& logger) //! //! LOG_ERROR(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_ERROR(const Logger& logger) -{ +inline LogStreamConsumer LOG_ERROR(const Logger& logger) { return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); } @@ -493,11 +447,10 @@ inline LogStreamConsumer LOG_ERROR(const Logger& logger) //! //! LOG_FATAL(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_FATAL(const Logger& logger) -{ +inline LogStreamConsumer LOG_FATAL(const Logger& logger) { return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); } -} // anonymous namespace +} // anonymous namespace -#endif // TENSORRT_LOGGING_H +#endif // TENSORRT_LOGGING_H diff --git a/resnet/macros.h b/resnet/macros.h new file mode 100644 index 00000000..4f442e6a --- /dev/null +++ b/resnet/macros.h @@ -0,0 +1,35 @@ +#pragma once +#include + +#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 + +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) (TRT_VERSION >= TRT_VERSION_ENCODE((major), (minor), (patch), 0)) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) + +#if TRT_VERSION_LT(7, 2, 2) +#error "TensorRT >= 7.2.2 is required for this demo." +#endif + +#if TRT_VERSION_GE(8, 0, 0) +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif diff --git a/resnet/resnet.cpp b/resnet/resnet.cpp new file mode 100644 index 00000000..29133fe3 --- /dev/null +++ b/resnet/resnet.cpp @@ -0,0 +1,514 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" +#include "utils.h" + +static constexpr int32_t INPUT_H = 224; +static constexpr int32_t INPUT_W = 224; +static constexpr int32_t OUTPUT_SIZE = 1000; +static constexpr int32_t N = 1; +static constexpr std::size_t WORKSPACE_SIZE = 16 << 20; +static constexpr std::array NAMES = {"data", "prob"}; +static constexpr std::array SIZES = {3 * INPUT_H * INPUT_W, OUTPUT_SIZE}; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; +static constexpr bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1); +static constexpr std::array mean = {0.485f, 0.456f, 0.406f}; +static constexpr std::array stdv = {0.229f, 0.224f, 0.225f}; + +using namespace nvinfer1; +using WeightMap = std::map; +using M = nvinfer1::MatrixOperation; +using NDCF = nvinfer1::NetworkDefinitionCreationFlag; + +static Logger gLogger; + +struct ResNetVariant { + const char* name; + const char* wts_path; + const char* engine_path; + std::array layers; + bool bottleneck; + int32_t groups; + int32_t width_per_group; +}; + +static auto getVariantConfig(const std::string& name) -> ResNetVariant { + if (name == "resnet18") { + return {"resnet18", "models/resnet18.wts", "models/resnet18.engine", {2, 2, 2, 2}, false, 1, 64}; + } + if (name == "resnet34") { + return {"resnet34", "models/resnet34.wts", "models/resnet34.engine", {3, 4, 6, 3}, false, 1, 64}; + } + if (name == "resnet50") { + return {"resnet50", "models/resnet50.wts", "models/resnet50.engine", {3, 4, 6, 3}, true, 1, 64}; + } + if (name == "resnext50_32x4d") { + return {"resnext50_32x4d", + "models/resnext50_32x4d.wts", + "models/resnext50_32x4d.engine", + {3, 4, 6, 3}, + true, + 32, + 4}; + } + if (name == "wide_resnet50_2") { + return {"wide_resnet50_2", + "models/wide_resnet50_2.wts", + "models/wide_resnet50_2.engine", + {3, 4, 6, 3}, + true, + 1, + 128}; + } + std::cerr << "unsupported resnet variant: " << name << "\n"; + std::cerr << "model: resnet18 | resnet34 | resnet50 | resnext50_32x4d | wide_resnet50_2\n"; + std::abort(); +} + +static auto addBatchNorm2d(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& lname, + float eps) -> IScaleLayer* { + const auto* gamma = static_cast(weight_map[lname + ".weight"].values); + const auto* beta = static_cast(weight_map[lname + ".bias"].values); + const auto* mean_ptr = static_cast(weight_map[lname + ".running_mean"].values); + const auto* var = static_cast(weight_map[lname + ".running_var"].values); + const auto len = weight_map[lname + ".running_var"].count; + + auto* scval = static_cast(std::malloc(sizeof(float) * static_cast(len))); + auto* shval = static_cast(std::malloc(sizeof(float) * static_cast(len))); + auto* pval = static_cast(std::malloc(sizeof(float) * static_cast(len))); + for (int64_t i = 0; i < len; ++i) { + scval[i] = gamma[i] / std::sqrt(var[i] + eps); + shval[i] = beta[i] - mean_ptr[i] * gamma[i] / std::sqrt(var[i] + eps); + pval[i] = 1.0f; + } + + Weights scale{DataType::kFLOAT, scval, len}; + Weights shift{DataType::kFLOAT, shval, len}; + Weights power{DataType::kFLOAT, pval, len}; + weight_map[lname + ".scale"] = scale; + weight_map[lname + ".shift"] = shift; + weight_map[lname + ".power"] = power; + + auto* layer = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(layer); +#if TRT_VERSION_GE(8, 0, 0) + layer->setChannelAxis(1); +#endif + return layer; +} + +static auto addConv2d(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& name, + int32_t out_channels, int32_t kernel, int32_t stride, int32_t padding, + int32_t groups = 1) -> IConvolutionLayer* { + static constexpr Weights empty{DataType::kFLOAT, nullptr, 0}; + auto* conv = + network->addConvolutionNd(input, out_channels, DimsHW{kernel, kernel}, weight_map[name + ".weight"], empty); + assert(conv); + conv->setStrideNd(DimsHW{stride, stride}); + conv->setPaddingNd(DimsHW{padding, padding}); + conv->setNbGroups(groups); + conv->setName(name.c_str()); + return conv; +} + +static auto addRelu(INetworkDefinition* network, ITensor& input) -> IActivationLayer* { + auto* relu = network->addActivation(input, ActivationType::kRELU); + assert(relu); + return relu; +} + +static auto addBasicBlock(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, int32_t in_channels, + int32_t planes, int32_t stride, const std::string& lname) -> ITensor* { + auto* conv1 = addConv2d(network, weight_map, input, lname + "conv1", planes, 3, stride, 1); + auto* bn1 = addBatchNorm2d(network, weight_map, *conv1->getOutput(0), lname + "bn1", 1e-5f); + auto* relu1 = addRelu(network, *bn1->getOutput(0)); + + auto* conv2 = addConv2d(network, weight_map, *relu1->getOutput(0), lname + "conv2", planes, 3, 1, 1); + auto* bn2 = addBatchNorm2d(network, weight_map, *conv2->getOutput(0), lname + "bn2", 1e-5f); + + ITensor* shortcut = &input; + if (stride != 1 || in_channels != planes) { + auto* conv3 = addConv2d(network, weight_map, input, lname + "downsample.0", planes, 1, stride, 0); + auto* bn3 = addBatchNorm2d(network, weight_map, *conv3->getOutput(0), lname + "downsample.1", 1e-5f); + shortcut = bn3->getOutput(0); + } + + auto* sum = network->addElementWise(*shortcut, *bn2->getOutput(0), ElementWiseOperation::kSUM); + assert(sum); + return addRelu(network, *sum->getOutput(0))->getOutput(0); +} + +static auto addBottleneck(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, int32_t in_channels, + int32_t planes, int32_t stride, int32_t groups, int32_t width_per_group, + const std::string& lname) -> ITensor* { + constexpr int32_t expansion = 4; + const int32_t width = planes * width_per_group / 64 * groups; + const int32_t out_channels = planes * expansion; + + auto* conv1 = addConv2d(network, weight_map, input, lname + "conv1", width, 1, 1, 0); + auto* bn1 = addBatchNorm2d(network, weight_map, *conv1->getOutput(0), lname + "bn1", 1e-5f); + auto* relu1 = addRelu(network, *bn1->getOutput(0)); + + auto* conv2 = addConv2d(network, weight_map, *relu1->getOutput(0), lname + "conv2", width, 3, stride, 1, groups); + auto* bn2 = addBatchNorm2d(network, weight_map, *conv2->getOutput(0), lname + "bn2", 1e-5f); + auto* relu2 = addRelu(network, *bn2->getOutput(0)); + + auto* conv3 = addConv2d(network, weight_map, *relu2->getOutput(0), lname + "conv3", out_channels, 1, 1, 0); + auto* bn3 = addBatchNorm2d(network, weight_map, *conv3->getOutput(0), lname + "bn3", 1e-5f); + + ITensor* shortcut = &input; + if (stride != 1 || in_channels != out_channels) { + auto* conv4 = addConv2d(network, weight_map, input, lname + "downsample.0", out_channels, 1, stride, 0); + auto* bn4 = addBatchNorm2d(network, weight_map, *conv4->getOutput(0), lname + "downsample.1", 1e-5f); + shortcut = bn4->getOutput(0); + } + + auto* sum = network->addElementWise(*shortcut, *bn3->getOutput(0), ElementWiseOperation::kSUM); + assert(sum); + return addRelu(network, *sum->getOutput(0))->getOutput(0); +} + +static auto addLayer(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, int32_t& in_channels, + const ResNetVariant& variant, int32_t layer_index, int32_t planes, int32_t blocks, + int32_t stride) -> ITensor* { + ITensor* x = &input; + for (int32_t i = 0; i < blocks; ++i) { + const int32_t block_stride = i == 0 ? stride : 1; + const std::string lname = "layer" + std::to_string(layer_index) + "." + std::to_string(i) + "."; + if (variant.bottleneck) { + x = addBottleneck(network, weight_map, *x, in_channels, planes, block_stride, variant.groups, + variant.width_per_group, lname); + in_channels = planes * 4; + } else { + x = addBasicBlock(network, weight_map, *x, in_channels, planes, block_stride, lname); + in_channels = planes; + } + } + return x; +} + +static auto addLinear(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& lname, + int32_t in_features, int32_t out_features) -> ITensor* { + auto* reshape = network->addShuffle(input); + assert(reshape); + reshape->setReshapeDimensions(Dims2{N, in_features}); + + auto* kernel = network->addConstant(Dims2{out_features, in_features}, weight_map[lname + ".weight"]); + assert(kernel); + auto* matmul = network->addMatrixMultiply(*reshape->getOutput(0), M::kNONE, *kernel->getOutput(0), M::kTRANSPOSE); + assert(matmul); + + auto* bias = network->addConstant(Dims2{1, out_features}, weight_map[lname + ".bias"]); + assert(bias); + auto* sum = network->addElementWise(*matmul->getOutput(0), *bias->getOutput(0), ElementWiseOperation::kSUM); + assert(sum); + return sum->getOutput(0); +} + +static auto buildResNet(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const ResNetVariant& variant) -> ITensor* { + auto* conv1 = addConv2d(network, weight_map, input, "conv1", 64, 7, 2, 3); + auto* bn1 = addBatchNorm2d(network, weight_map, *conv1->getOutput(0), "bn1", 1e-5f); + auto* relu1 = addRelu(network, *bn1->getOutput(0)); + + auto* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3}); + assert(pool1); + pool1->setStrideNd(DimsHW{2, 2}); + pool1->setPaddingNd(DimsHW{1, 1}); + + int32_t in_channels = 64; + ITensor* x = addLayer(network, weight_map, *pool1->getOutput(0), in_channels, variant, 1, 64, variant.layers[0], 1); + x = addLayer(network, weight_map, *x, in_channels, variant, 2, 128, variant.layers[1], 2); + x = addLayer(network, weight_map, *x, in_channels, variant, 3, 256, variant.layers[2], 2); + x = addLayer(network, weight_map, *x, in_channels, variant, 4, 512, variant.layers[3], 2); + + auto* pool2 = network->addPoolingNd(*x, PoolingType::kAVERAGE, DimsHW{7, 7}); + assert(pool2); + pool2->setStrideNd(DimsHW{1, 1}); + return addLinear(network, weight_map, *pool2->getOutput(0), "fc", in_channels, OUTPUT_SIZE); +} + +static auto createEngine(int32_t batch_size, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt, + const ResNetVariant& variant) -> ICudaEngine* { + WeightMap weight_map = loadWeights(variant.wts_path); + +#if TRT_VERSION_GE(10, 12, 0) + auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); +#elif TRT_VERSION_GE(10, 0, 0) + auto flag = 0U; +#else + auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); +#endif + auto* network = builder->createNetworkV2(flag); + assert(network); + + ITensor* input = nullptr; + if constexpr (TRT_PREPROCESS) { + dt = DataType::kUINT8; + input = network->addInput(NAMES[0], dt, Dims4{batch_size, INPUT_H, INPUT_W, 3}); + auto* transform = addTransformLayer(network, *input, true, mean, stdv); + input = transform->getOutput(0); + } else { + input = network->addInput(NAMES[0], dt, Dims4{batch_size, 3, INPUT_H, INPUT_W}); + } + assert(input); + + ITensor* logits = buildResNet(network, weight_map, *input, variant); + logits->setName(NAMES[1]); + network->markOutput(*logits); + +#if TRT_VERSION_GE(8, 0, 0) + config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); + IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); + assert(mem); + ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); + delete mem; + delete network; +#else + builder->setMaxBatchSize(batch_size); + config->setMaxWorkspaceSize(WORKSPACE_SIZE); + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + network->destroy(); +#endif + std::cout << "build finished\n"; + + for (auto& mem : weight_map) { + std::free(const_cast(mem.second.values)); + } + return engine; +} + +static void APIToModel(int32_t batch_size, IRuntime* runtime, IHostMemory** model_stream, + const ResNetVariant& variant) { + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + ICudaEngine* engine = createEngine(batch_size, runtime, builder, config, DataType::kFLOAT, variant); + assert(engine != nullptr); + + (*model_stream) = engine->serialize(); + +#if TRT_VERSION_GE(8, 0, 0) + delete engine; + delete config; + delete builder; +#else + engine->destroy(); + config->destroy(); + builder->destroy(); +#endif +} + +static auto doInference(IExecutionContext& context, void* input, + int64_t batch_size) -> std::vector> { + ICudaEngine const& engine = context.getEngine(); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + std::vector buffers; + +#if TRT_VERSION_GE(8, 0, 0) + const int32_t nIO = engine.getNbIOTensors(); +#else + const int32_t nIO = engine.getNbBindings(); +#endif + + buffers.resize(nIO); + for (auto i = 0; i < nIO; ++i) { + std::size_t size = 0; +#if TRT_VERSION_GE(8, 0, 0) + auto* tensor_name = engine.getIOTensorName(i); + const std::string name = tensor_name; + auto element_size = getSize(engine.getTensorDataType(tensor_name)); + size = element_size * static_cast(batch_size) * (name == NAMES[0] ? SIZES[0] : SIZES[1]); + CHECK(cudaMalloc(&buffers[i], size)); + if (name == NAMES[0]) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } +#else + const int32_t idx = engine.getBindingIndex(NAMES[i]); + auto element_size = getSize(engine.getBindingDataType(idx)); + assert(idx == i); + size = element_size * static_cast(batch_size) * SIZES[i]; + CHECK(cudaMalloc(&buffers[i], size)); + if (i == 0) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } +#endif + } + +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } +#else + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } +#endif + + std::vector> prob; + for (int32_t i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) + const std::string name = engine.getIOTensorName(i); + if (name == NAMES[0]) { + continue; + } + constexpr auto output_size = SIZES[1]; +#else + if (i == 0) { + continue; + } + const auto output_size = SIZES[i]; +#endif + std::vector tmp(static_cast(batch_size) * output_size, std::nanf("")); + const auto size = static_cast(batch_size) * output_size * sizeof(float); + CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); + prob.emplace_back(std::move(tmp)); + } + CHECK(cudaStreamSynchronize(stream)); + + for (auto& buffer : buffers) { + CHECK(cudaFree(buffer)); + } + CHECK(cudaStreamDestroy(stream)); + return prob; +} + +auto main(int argc, char** argv) -> int { + checkTrtEnv(); + if (argc < 2 || argc > 3) { + std::cerr << "arguments not right!\n"; + std::cerr << "./resnet -s [model] // serialize model to plan file\n"; + std::cerr << "./resnet -d [model] // deserialize plan file and run inference\n"; + std::cerr << "model: resnet18 | resnet34 | resnet50 | resnext50_32x4d | wide_resnet50_2\n"; + return -1; + } + + const auto variant = getVariantConfig(argc == 3 ? argv[2] : "resnet18"); + std::cout << "Using ResNet variant: " << variant.name << "\n"; + + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + char* trt_model_stream{nullptr}; + std::streamsize size{0}; + + if (std::string(argv[1]) == "-s") { + IHostMemory* model_stream{nullptr}; + APIToModel(N, runtime, &model_stream, variant); + assert(model_stream != nullptr); + + std::ofstream plan(variant.engine_path, std::ios::binary | std::ios::trunc); + if (!plan) { + std::cerr << "could not open plan output file\n"; + return -1; + } + if (model_stream->size() > static_cast(std::numeric_limits::max())) { + std::cerr << "this model is too large to serialize\n"; + return -1; + } + const auto* data_ptr = reinterpret_cast(model_stream->data()); + const auto data_size = static_cast(model_stream->size()); + plan.write(data_ptr, data_size); +#if TRT_VERSION_GE(8, 0, 0) + delete model_stream; + delete runtime; +#else + model_stream->destroy(); + runtime->destroy(); +#endif + return 0; + } + if (std::string(argv[1]) == "-d") { + std::ifstream file(variant.engine_path, std::ios::binary); + if (file.good()) { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trt_model_stream = new char[size]; + assert(trt_model_stream); + file.read(trt_model_stream, size); + file.close(); + } + } else { + return -1; + } + +#if TRT_VERSION_GE(8, 0, 0) + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size); +#else + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size, nullptr); +#endif + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + delete[] trt_model_stream; + + void* input = nullptr; + std::vector flat_img; + cv::Mat img = cv::imread("assets/cats.jpg", cv::IMREAD_COLOR); + if (img.empty()) { + std::cerr << "failed to read image: assets/cats.jpg\n"; + return -1; + } + if constexpr (TRT_PREPROCESS) { + cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR); + input = static_cast(img.data); + } else { + flat_img = preprocess_img(img, true, mean, stdv, N, INPUT_H, INPUT_W); + input = flat_img.data(); + } + + auto first_prob = doInference(*context, input, N); + printFirstOutputs("resnet", first_prob[0].data(), first_prob[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + if (labels.empty()) { + std::cerr << "failed to load labels from " << LABELS_PATH << "\n"; + std::abort(); + } + int32_t top = 0; + for (auto& [idx, logits] : topk(first_prob[0], 3)) { + std::cout << "Top: " << top++ << " idx: " << idx << ", logits: " << std::setprecision(4) << logits + << ", label: " << labels[idx] << "\n"; + } + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, N); + auto end = std::chrono::steady_clock::now(); + auto period = std::chrono::duration_cast(end - start); + latencies.push_back(static_cast(period.count()) / 1000.0); + } + printBenchmark("resnet", latencies, N); +#if TRT_VERSION_GE(8, 0, 0) + delete context; + delete engine; + delete runtime; +#else + context->destroy(); + engine->destroy(); + runtime->destroy(); +#endif + + return 0; +} diff --git a/resnet/utils.h b/resnet/utils.h new file mode 100644 index 00000000..4f111957 --- /dev/null +++ b/resnet/utils.h @@ -0,0 +1,311 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "macros.h" + +using namespace nvinfer1; + +#define CHECK(status) \ + do { \ + auto ret = (status); \ + if (ret != cudaSuccess) { \ + std::cerr << "Cuda failure: " << ret << "\n"; \ + std::abort(); \ + } \ + } while (0) + +static void checkTrtEnv(int device = 0) { +#if TRT_VERSION_LT(8, 0, 0) + CHECK(cudaGetDevice(&device)); + cudaDeviceProp prop{}; + CHECK(cudaGetDeviceProperties(&prop, device)); + const int sm = prop.major * 10 + prop.minor; + if (sm > 86) { + std::cerr << "TensorRT < 8 does not support SM > 86 on this GPU."; + std::abort(); + } +#endif +} + +template +static auto toI32(T value) -> int32_t { + static_assert(std::is_integral::value, "toI32 requires an integral type."); + static_assert(sizeof(T) > sizeof(int32_t), "toI32 is only for values wider than int32_t."); + if constexpr (std::is_signed::value) { + assert(value >= static_cast(std::numeric_limits::min())); + } + assert(value <= static_cast(std::numeric_limits::max())); + return static_cast(value); +} + +/** + * @brief TensorRT weight files have a simple space delimited format: + * [type] [size] + * + * @param file input weight file path + * @return std::map + */ +static std::map loadWeights(const std::string& file) { + std::cout << "Loading weights: " << file << "\n"; + std::map 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--) { + nvinfer1::Weights wt{nvinfer1::DataType::kFLOAT, nullptr, 0}; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> wt.count; + + // Load blob + auto* val = static_cast(std::malloc(sizeof(uint32_t) * static_cast(wt.count))); + if (val == nullptr) { + std::cerr << "weight allocation failed\n"; + std::abort(); + } + input >> std::hex; + for (auto x = 0ll; x < wt.count; ++x) { + input >> val[x]; + } + wt.values = val; + weightMap[name] = wt; + } + + return weightMap; +} + +/** + * @brief a preprocess function aligning with ImageNet preprocess in torchvision, only support 3-channel image + * + * @param img opencv image with BGR layout + * @param bgr2rgb whether to convert BGR to RGB + * @param mean subtract mean + * @param std divide std + * @param n batch size + * @param h resize height + * @param w resize width + * @return std::vector contiguous flatten image data in float32 type + */ +static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const std::array& mean, + const std::array& std, int n, int h, int w) { + const auto c = img.channels(); + const auto size = c * h * w; + if (c != 3) { + std::cerr << "this demo only supports 3 channel input image.\n"; + std::abort(); + } + if (bgr2rgb) { + cv::cvtColor(img, img, cv::COLOR_BGR2RGB); + } + cv::resize(img, img, cv::Size(w, h), 0, 0, cv::INTER_LINEAR); + img.convertTo(img, CV_32FC3, 1.f / 255); + img = (img - cv::Scalar(mean[0], mean[1], mean[2])) / cv::Scalar(std[0], std[1], std[2]); + std::vector chw(static_cast(n) * c * h * w, 0.f); + + // fill all batch with the same input image + for (int i = 0; i < n; ++i) { + for (int y = 0; y < h; ++y) { + for (int x = 0; x < w; ++x) { + const cv::Vec3f v = img.at(y, x); + chw[i * size + 0 * h * w + y * w + x] = v[0]; + chw[i * size + 1 * h * w + y * w + x] = v[1]; + chw[i * size + 2 * h * w + y * w + x] = v[2]; + } + } + } + return chw; +} + +static auto topk(const std::vector& v, int k) -> std::vector> { + if (k <= 0) + return {}; + auto stride = std::min(k, static_cast(v.size())); + + std::vector idx(v.size()); + std::iota(idx.begin(), idx.end(), 0); + + std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&](int a, int b) { return v[a] > v[b]; }); + + std::vector> out; + out.reserve(stride); + for (auto i = 0; i < stride; ++i) + out.emplace_back(idx[i], v[idx[i]]); + return out; +} + +static std::map loadImagenetLabelMap(const std::string& path) { + std::map labels; + std::ifstream in(path); + if (!in.is_open()) { + return labels; + } + std::string line; + while (std::getline(in, line)) { + auto colon = line.find(':'); + if (colon == std::string::npos) { + continue; + } + auto first_quote = line.find('\'', colon); + if (first_quote == std::string::npos) { + continue; + } + auto second_quote = line.find('\'', first_quote + 1); + if (second_quote == std::string::npos) { + continue; + } + int idx = std::stoi(line.substr(0, colon)); + labels[idx] = line.substr(first_quote + 1, second_quote - first_quote - 1); + } + return labels; +} + +static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bool bgr2rgb, + const std::array& mean, const std::array& std) { + struct ScaleParams { + std::array shift; + std::array scale; + }; + static std::vector> gScaleParams; + auto params = std::make_unique(); + params->shift = {-mean[0] / std[0], -mean[1] / std[1], -mean[2] / std[2]}; + params->scale = {1.f / (std[0] * 255.f), 1.f / (std[1] * 255.f), 1.f / (std[2] * 255.f)}; + + static const Weights empty{DataType::kFLOAT, nullptr, 0ll}; + const Weights shift{DataType::kFLOAT, params->shift.data(), 3ll}; + const Weights scale{DataType::kFLOAT, params->scale.data(), 3ll}; + + gScaleParams.emplace_back(std::move(params)); + + ITensor* in = &input; + if (input.getType() != DataType::kFLOAT) { +#if TRT_VERSION_GE(8, 0, 0) + auto* cast = network->addCast(input, DataType::kFLOAT); + assert(cast); + cast->setName("Cast to FP32"); + in = cast->getOutput(0); +#else + auto* identity = network->addIdentity(input); + assert(identity); + identity->setName("Convert to FP32"); + identity->setOutputType(0, DataType::kFLOAT); + in = identity->getOutput(0); +#endif + } + // Convert from NHWC to NCHW + auto* perm = network->addShuffle(*in); + assert(perm); + perm->setName("NHWC -> NCHW"); + perm->setFirstTranspose(Permutation{0, 3, 1, 2}); + + // Convert from BGR to RGB (optional) + ITensor* data{nullptr}; + if (bgr2rgb) { + auto add_slice = [&](int c, const char* name) -> ITensor* { + auto dims = perm->getOutput(0)->getDimensions(); + Dims4 start = {0, c, 0, 0}, stride = {1, 1, 1, 1}; + Dims4 size = {dims.d[0], 1, dims.d[2], dims.d[3]}; + auto* _slice = network->addSlice(*perm->getOutput(0), start, size, stride); + _slice->setName(name); + assert(_slice && _slice->getNbOutputs() == 1); + return _slice->getOutput(0); + }; + std::array channels = {add_slice(2, "R"), add_slice(1, "G"), add_slice(0, "B")}; + auto* cat = network->addConcatenation(channels.data(), 3); + assert(cat); + cat->setName("RGB"); + cat->setAxis(1); + data = cat->getOutput(0); + } else { + data = perm->getOutput(0); + } + + // Normalize + auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); + assert(trans); + trans->setName("mean & std"); +#if TRT_VERSION_GE(8, 0, 0) + trans->setChannelAxis(1); +#endif + return trans; +} + +static size_t getSize(DataType dt) { + switch (dt) { +#if TRT_VERSION_GE(8, 5, 1) + case DataType::kUINT8: +#endif + case DataType::kINT8: + return sizeof(int8_t); + case DataType::kFLOAT: + return sizeof(float); + case DataType::kHALF: + return sizeof(int16_t); + case DataType::kINT32: + return sizeof(int32_t); + default: { + std::cerr << "Unsupported data type\n"; + std::abort(); + } + } +} + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/shufflenetv2/FindTensorRT.cmake b/shufflenetv2/FindTensorRT.cmake index af75bfe2..61203d17 100644 --- a/shufflenetv2/FindTensorRT.cmake +++ b/shufflenetv2/FindTensorRT.cmake @@ -73,8 +73,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -139,7 +142,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/shufflenetv2/gen_wts.py b/shufflenetv2/gen_wts.py index db632d81..37416ea0 100644 --- a/shufflenetv2/gen_wts.py +++ b/shufflenetv2/gen_wts.py @@ -1,83 +1,85 @@ +import argparse +import os import struct +from pathlib import Path -import cv2 -import numpy as np import torch -from torchvision.models.shufflenetv2 import ( - shufflenet_v2_x0_5, - shufflenet_v2_x1_0, - shufflenet_v2_x1_5, - shufflenet_v2_x2_0, +from torchvision import models + + +MODEL_NAMES = ( + "shufflenet_v2_x0_5", + "shufflenet_v2_x1_0", + "shufflenet_v2_x1_5", + "shufflenet_v2_x2_0", ) +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +MODELS_DIR = REPO_ROOT / "models" -def read_imagenet_labels() -> dict[int, str]: - """ - read ImageNet 1000 labels +def require_cache_env() -> None: + missing = [name for name in ("TORCH_HOME", "HF_HOME") if not os.environ.get(name)] + if missing: + raise RuntimeError( + f"Please set required cache environment variables: {', '.join(missing)}" + ) + print(f"Using TORCH_HOME={os.environ['TORCH_HOME']}") + print(f"Using HF_HOME={os.environ['HF_HOME']}") - Returns: - dict[int, str]: labels dict - """ - clsid2label = {} - with open("../assets/imagenet1000_clsidx_to_labels.txt", "r") as f: - for i in f.readlines(): - k, v = i.split(": ") - clsid2label.setdefault(int(k), v[1:-3]) - return clsid2label +def build_model(name: str) -> torch.nn.Module: + suffix = name.removeprefix("shufflenet_").upper() + weights_cls = getattr(models, f"ShuffleNet_{suffix}_Weights") + model_fn = getattr(models, name) + model = model_fn(weights=weights_cls.DEFAULT) + model.eval() + return model -def preprocess(img: np.array) -> torch.Tensor: - """ - a preprocess method align with ImageNet dataset - Args: - img (np.array): input image +def write_wts(model: torch.nn.Module, output_path: Path) -> None: + state_dict = model.state_dict() + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w") as f: + f.write(f"{len(state_dict.keys())}\n") + for key, value in state_dict.items(): + values = value.reshape(-1).cpu().numpy() + f.write(f"{key} {len(values)}") + for item in values: + f.write(" ") + f.write(struct.pack(">f", float(item)).hex()) + f.write("\n") + size_mib = output_path.stat().st_size / 1024 / 1024 + print(f"[ok] wrote {output_path} ({size_mib:.1f} MiB, {len(state_dict)} tensors)") - Returns: - torch.Tensor: preprocessed image in `NCHW` layout - """ - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 - img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR) - mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) - std = np.array([0.229, 0.224, 0.225], dtype=np.float32) - img = (img - mean) / std - img = img.transpose(2, 0, 1)[None, ...] - return torch.from_numpy(img) +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export torchvision ShuffleNetV2 weights to tensorrtx .wts files." + ) + parser.add_argument( + "--model", + choices=MODEL_NAMES, + action="append", + help="ShuffleNetV2 variant to export. Repeat for multiple variants. Defaults to all variants.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=MODELS_DIR, + help="Directory for generated .wts files.", + ) + return parser.parse_args() -if __name__ == "__main__": - labels = read_imagenet_labels() - img = cv2.imread("../assets/cats.jpg", cv2.IMREAD_COLOR) - img = preprocess(img) - """ - NOTE: comment out the model you don't want - """ - models = [ - ("shufflenet_v2_x0_5", shufflenet_v2_x0_5(pretrained=True)), - ("shufflenet_v2_x1_0", shufflenet_v2_x1_0(pretrained=True)), - ("shufflenet_v2_x1_5", shufflenet_v2_x1_5(pretrained=True)), - ("shufflenet_v2_x2_0", shufflenet_v2_x2_0(pretrained=True)), - ] +def main() -> None: + args = parse_args() + require_cache_env() + model_names = args.model or MODEL_NAMES + for model_name in model_names: + print(f"writing {model_name}.wts") + write_wts(build_model(model_name), args.output_dir / f"{model_name}.wts") - for name, model in models: - model.eval() - with torch.inference_mode(): - output = model(img) - print(f"{name} result:") - for i, batch in enumerate(torch.topk(output, k=3).indices): - for j, idx in enumerate(batch): - print(f"\tBatch: {i}, Top: {j}, logits: {output[i][idx]:.4f}, label: {labels[int(idx)]}") - print(f"{'=' * 32}") - with open(f"../models/{name}.wts", "w") as f: - 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") +if __name__ == "__main__": + main() diff --git a/shufflenetv2/macros.h b/shufflenetv2/macros.h index daae7bfb..6d851752 100644 --- a/shufflenetv2/macros.h +++ b/shufflenetv2/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/shufflenetv2/shufflenetv2.cpp b/shufflenetv2/shufflenetv2.cpp index cfc28449..1773526a 100644 --- a/shufflenetv2/shufflenetv2.cpp +++ b/shufflenetv2/shufflenetv2.cpp @@ -1,46 +1,41 @@ #include +#include #include #include +#include +#include +#include +#include +#include #include +#include #include #include #include +#include #include + #include "logging.h" #include "utils.h" -struct ShuffleNetV2Params { +struct ShuffleNetV2Variant { + std::string name; std::array repeat; std::array output_chn; + std::string wts_path; + std::string engine_path; }; -/** - * @brief choose one below as the model to be built - * @param v2_x0_5 - * @param v2_x1_0 - * @param v2_x1_5 - * @param v2_x2_0 - */ -[[maybe_unused]] static constexpr ShuffleNetV2Params v2_x0_5 = {{4, 8, 4}, {24, 48, 96, 192, 1024}}; -[[maybe_unused]] static constexpr ShuffleNetV2Params v2_x1_0 = {{4, 8, 4}, {24, 116, 232, 464, 1024}}; -[[maybe_unused]] static constexpr ShuffleNetV2Params v2_x1_5 = {{4, 8, 4}, {24, 176, 352, 704, 1024}}; -[[maybe_unused]] static constexpr ShuffleNetV2Params v2_x2_0 = {{4, 8, 4}, {24, 244, 488, 976, 2048}}; - -constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; - -// stuff we know about shufflenet-v2 -constexpr const int64_t N = 1; -constexpr const int32_t INPUT_H = 224; -constexpr const int32_t INPUT_W = 224; -constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; -constexpr const std::array NAMES = {"data", "logits"}; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; +static constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; +static constexpr const int64_t N = 1; +static constexpr const int32_t INPUT_H = 224; +static constexpr const int32_t INPUT_W = 224; +static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; +static constexpr const std::array NAMES = {"data", "logits"}; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1); static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; - -static constexpr const char* WTS_PATH = "../models/shufflenet_v2_x0_5.wts"; -static constexpr const char* ENGINE_PATH = "../models/shufflenet.engine"; -static constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; using namespace nvinfer1; using WeightMap = std::map; @@ -49,265 +44,229 @@ using NDCF = nvinfer1::NetworkDefinitionCreationFlag; static Logger gLogger; -Dims debug_shape(const ILayer* l) { - Dims dims = l->getOutput(0)->getDimensions(); - std::cout << l->getOutput(0)->getName() << ":\t["; - for (int i = 0; i < dims.nbDims; i++) { - std::cout << dims.d[i] << ", "; +static auto getVariantConfig(const std::string& name) -> ShuffleNetV2Variant { + static const std::map, std::array>> variants = { + {"shufflenet_v2_x0_5", {{4, 8, 4}, {24, 48, 96, 192, 1024}}}, + {"shufflenet_v2_x1_0", {{4, 8, 4}, {24, 116, 232, 464, 1024}}}, + {"shufflenet_v2_x1_5", {{4, 8, 4}, {24, 176, 352, 704, 1024}}}, + {"shufflenet_v2_x2_0", {{4, 8, 4}, {24, 244, 488, 976, 2048}}}, + }; + + const auto iter = variants.find(name); + if (iter == variants.end()) { + std::cerr << "Unsupported ShuffleNetV2 variant: " << name << "\n"; + std::cerr << "Choose one of: shufflenet_v2_x0_5, shufflenet_v2_x1_0, shufflenet_v2_x1_5, shufflenet_v2_x2_0\n"; + std::abort(); } - std::cout << "]\n"; - return dims; + + return ShuffleNetV2Variant{name, iter->second.first, iter->second.second, "models/" + name + ".wts", + "models/" + name + ".engine"}; } -ILayer* addBatchNorm2d(INetworkDefinition* network, WeightMap& weightMap, ITensor& input, const std::string& lname, - float eps = 1e-3f) { - 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; - auto len = weightMap[lname + ".running_var"].count; - std::cout << lname << " running_var len: " << len << "\n"; - - auto* scval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { - scval[i] = gamma[i] / sqrt(var[i] + eps); +static auto addBatchNorm2d(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& lname, + float eps = 1e-3f) -> ILayer* { + const auto* gamma = static_cast(weight_map[lname + ".weight"].values); + const auto* beta = static_cast(weight_map[lname + ".bias"].values); + const auto* bn_mean = static_cast(weight_map[lname + ".running_mean"].values); + const auto* var = static_cast(weight_map[lname + ".running_var"].values); + const auto len = weight_map[lname + ".running_var"].count; + + auto* scale_values = reinterpret_cast(std::malloc(sizeof(float) * static_cast(len))); + auto* shift_values = reinterpret_cast(std::malloc(sizeof(float) * static_cast(len))); + if (scale_values == nullptr || shift_values == nullptr) { + std::cerr << "batchnorm weight allocation failed\n"; + std::abort(); } - Weights scale{DataType::kFLOAT, scval, len}; - - auto* shval = static_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { - shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); + for (int64_t i = 0; i < len; i++) { + scale_values[i] = gamma[i] / std::sqrt(var[i] + eps); + shift_values[i] = beta[i] - bn_mean[i] * gamma[i] / std::sqrt(var[i] + eps); } - Weights shift{DataType::kFLOAT, shval, len}; + + Weights scale{DataType::kFLOAT, scale_values, len}; + Weights shift{DataType::kFLOAT, shift_values, len}; static const Weights power{DataType::kFLOAT, nullptr, 0ll}; - 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; + weight_map[lname + ".scale"] = scale; + weight_map[lname + ".shift"] = shift; + weight_map[lname + ".power"] = power; + IScaleLayer* scale_layer = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_layer); + return scale_layer; } -/** - * @brief a basic convolution+bn layer with an optional relu layer - * - * @param network network definition - * @param m weight map - * @param input input tensor - * @param lname layer name - * @param ch output channels - * @param k kernel - * @param s stride - * @param p padding - * @param g groups - * @param with_relu true if with relu - * @return ILayer* - */ -ILayer* CBR(INetworkDefinition* network, WeightMap& m, ITensor& input, const std::string& lname, int ch, int k, - int s = 1, int p = 0, int g = 1, bool with_relu = true, int start_index = 0) { - static const Weights emptywts{DataType::kFLOAT, nullptr, 0ll}; - auto conv_name = lname + "." + std::to_string(start_index++); - auto* conv = network->addConvolutionNd(input, ch, DimsHW{k, k}, m[conv_name + ".weight"], emptywts); - +static auto addConvBnRelu(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, const std::string& lname, + int32_t out_channels, int32_t kernel_size, int32_t stride = 1, int32_t padding = 0, + int32_t groups = 1, bool with_relu = true, int32_t start_index = 0) -> ILayer* { + static const Weights empty_weights{DataType::kFLOAT, nullptr, 0ll}; + const auto conv_name = lname + "." + std::to_string(start_index++); + auto* conv = network->addConvolutionNd(input, out_channels, DimsHW{kernel_size, kernel_size}, + weight_map[conv_name + ".weight"], empty_weights); assert(conv); - conv->setStrideNd(DimsHW{s, s}); - conv->setPaddingNd(DimsHW{p, p}); - conv->setNbGroups(g); + conv->setStrideNd(DimsHW{stride, stride}); + conv->setPaddingNd(DimsHW{padding, padding}); + conv->setNbGroups(groups); conv->setName(conv_name.c_str()); - auto bn_name = lname + "." + std::to_string(start_index++); - auto* bn = addBatchNorm2d(network, m, *conv->getOutput(0), bn_name, 1e-5f); + const auto bn_name = lname + "." + std::to_string(start_index++); + auto* bn = addBatchNorm2d(network, weight_map, *conv->getOutput(0), bn_name, 1e-5f); bn->setName((bn_name + ".bn").c_str()); - if (with_relu) { - auto* relu = network->addActivation(*bn->getOutput(0), ActivationType::kRELU); - auto relu_name = lname + "." + std::to_string(start_index) + ".relu"; - assert(relu); - relu->setName(relu_name.c_str()); - return relu; + if (!with_relu) { + return bn; } - return bn; + + auto* relu = network->addActivation(*bn->getOutput(0), ActivationType::kRELU); + assert(relu); + relu->setName((lname + "." + std::to_string(start_index) + ".relu").c_str()); + return relu; } -/** - * @brief invered residual block - * - * @param network network definition - * @param m weight map - * @param input input tensor - * @param lname layer name - * @param inch input channels - * @param outch output channels - * @param s stride - * @return ILayer* - */ -ILayer* invertedRes(INetworkDefinition* net, WeightMap& m, ITensor& input, const std::string& lname, int inch, - int outch, int s) { - if (s < 1 || s > 3) { +static auto addInvertedResidual(INetworkDefinition* network, WeightMap& weight_map, ITensor& input, + const std::string& lname, int32_t in_channels, int32_t out_channels, + int32_t stride) -> ILayer* { + if (stride < 1 || stride > 3) { std::cerr << "stride must be in [1, 3]\n"; std::abort(); } - int32_t bf /* branch features */ = outch / 2; - ITensor *x1{nullptr}, *x2{nullptr}; - - if (s == 1) { - auto d = input.getDimensions(); - Dims4 stride{1, 1, 1, 1}; - Dims4 half{d.d[0], d.d[1] / 2, d.d[2], d.d[3]}; - auto* s1 = net->addSlice(input, Dims4{0, 0, 0, 0}, half, stride); - auto* s2 = net->addSlice(input, Dims4{0, d.d[1] / 2, 0, 0}, half, stride); - debug_shape(s2); - x1 = s1->getOutput(0); - x2 = s2->getOutput(0); + + const int32_t branch_features = out_channels / 2; + ITensor* branch1_output = nullptr; + ITensor* branch2_input = nullptr; + + if (stride == 1) { + const auto dims = input.getDimensions(); + const Dims4 half{dims.d[0], dims.d[1] / 2, dims.d[2], dims.d[3]}; + auto* slice1 = network->addSlice(input, Dims4{0, 0, 0, 0}, half, Dims4{1, 1, 1, 1}); + auto* slice2 = network->addSlice(input, Dims4{0, dims.d[1] / 2, 0, 0}, half, Dims4{1, 1, 1, 1}); + assert(slice1); + assert(slice2); + branch1_output = slice1->getOutput(0); + branch2_input = slice2->getOutput(0); } else { - if (s > 1) { - auto* b1 = CBR(net, m, input, lname + ".branch1", inch, 3, s, 1, inch, false, 0); - b1 = CBR(net, m, *b1->getOutput(0), lname + ".branch1", inch, 1, 1, 0, 1, true, 2); - x1 = b1->getOutput(0); - debug_shape(b1); - } else { - x1 = &input; - } - x2 = &input; + auto* branch1 = addConvBnRelu(network, weight_map, input, lname + ".branch1", in_channels, 3, stride, 1, + in_channels, false, 0); + branch1 = addConvBnRelu(network, weight_map, *branch1->getOutput(0), lname + ".branch1", branch_features, 1, 1, + 0, 1, true, 2); + branch1_output = branch1->getOutput(0); + branch2_input = &input; } - auto* b2 = CBR(net, m, *x2, lname + ".branch2", bf, 1, 1, 0, 1, true, 0); - b2 = CBR(net, m, *b2->getOutput(0), lname + ".branch2", bf, 3, s, 1, bf, false, 3); - b2 = CBR(net, m, *b2->getOutput(0), lname + ".branch2", bf, 1, 1, 0, 1, true, 5); - debug_shape(b2); + auto* branch2 = addConvBnRelu(network, weight_map, *branch2_input, lname + ".branch2", branch_features, 1, 1, 0, 1, + true, 0); + branch2 = addConvBnRelu(network, weight_map, *branch2->getOutput(0), lname + ".branch2", branch_features, 3, stride, + 1, branch_features, false, 3); + branch2 = addConvBnRelu(network, weight_map, *branch2->getOutput(0), lname + ".branch2", branch_features, 1, 1, 0, + 1, true, 5); - std::array cat_tensors = {x1, b2->getOutput(0)}; - auto* cat = net->addConcatenation(cat_tensors.data(), 2); - auto cat_name = lname + ".cat"; + std::array cat_tensors = {branch1_output, branch2->getOutput(0)}; + auto* cat = network->addConcatenation(cat_tensors.data(), toI32(cat_tensors.size())); assert(cat); - cat->setName(cat_name.c_str()); + cat->setName((lname + ".cat").c_str()); cat->setAxis(1); - static_cast(debug_shape(cat)); - - auto* sf1 = net->addShuffle(*cat->getOutput(0)); - assert(sf1); - sf1->setName((lname + ".shuffle.1").c_str()); - auto d = cat->getOutput(0)->getDimensions(); - auto dim_sf1 = Dims{5, {d.d[0], 2, d.d[1] / 2, d.d[2], d.d[3]}}; - sf1->setReshapeDimensions(dim_sf1); - sf1->setSecondTranspose({0, 2, 1, 3, 4}); - - auto* sf2 = net->addShuffle(*sf1->getOutput(0)); - assert(sf2); - sf2->setName((lname + ".shuffle.2").c_str()); - sf2->setReshapeDimensions(d); - - return sf2; + + auto* shuffle1 = network->addShuffle(*cat->getOutput(0)); + assert(shuffle1); + shuffle1->setName((lname + ".shuffle.1").c_str()); + const auto dims = cat->getOutput(0)->getDimensions(); + shuffle1->setReshapeDimensions(Dims{5, {dims.d[0], 2, dims.d[1] / 2, dims.d[2], dims.d[3]}}); + shuffle1->setSecondTranspose({0, 2, 1, 3, 4}); + + auto* shuffle2 = network->addShuffle(*shuffle1->getOutput(0)); + assert(shuffle2); + shuffle2->setName((lname + ".shuffle.2").c_str()); + shuffle2->setReshapeDimensions(dims); + return shuffle2; } -/** - * @brief Create a Engine object - * - * @param N max batch size - * @param runtime runtime - * @param builder builder - * @param config config - * @param dt data type - * @param param the type of model to be built - * @return ICudaEngine* - */ -ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt, - ShuffleNetV2Params param = v2_x0_5) { - WeightMap m = loadWeights(WTS_PATH); - -#if TRT_VERSION >= 11200 +static auto createEngine(int32_t batch_size, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt, + const ShuffleNetV2Variant& variant) -> ICudaEngine* { + WeightMap weight_map = loadWeights(variant.wts_path); + +#if TRT_VERSION_GE(10, 12, 0) auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); -#elif TRT_VERSION >= 10000 +#elif TRT_VERSION_GE(10, 0, 0) auto flag = 0U; #else auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); #endif - auto* net = builder->createNetworkV2(flag); + auto* network = builder->createNetworkV2(flag); + assert(network); - int32_t in_ch = 3; - ITensor* input{nullptr}; + ITensor* input = nullptr; if constexpr (TRT_PREPROCESS) { - // for simplicity, resize image on cpu side dt = DataType::kUINT8; - input = net->addInput(NAMES[0], dt, Dims4{N, INPUT_H, INPUT_W, in_ch}); - auto* trans = addTransformLayer(net, *input, true, mean, stdv); - input = trans->getOutput(0); + input = network->addInput(NAMES[0], dt, Dims4{batch_size, INPUT_H, INPUT_W, 3}); + auto* transform = addTransformLayer(network, *input, true, mean, stdv); + input = transform->getOutput(0); } else { - input = net->addInput(NAMES[0], dt, Dims4{N, in_ch, INPUT_H, INPUT_W}); + input = network->addInput(NAMES[0], dt, Dims4{batch_size, 3, INPUT_H, INPUT_W}); } assert(input); - /** conv1 and maxpool */ - auto* cbr1 = CBR(net, m, *input, "conv1", param.output_chn[0], 3, 2, 1); - auto* pool1 = net->addPoolingNd(*cbr1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3}); + auto* conv1 = addConvBnRelu(network, weight_map, *input, "conv1", variant.output_chn[0], 3, 2, 1); + auto* pool1 = network->addPoolingNd(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3}); assert(pool1); pool1->setStrideNd(DimsHW{2, 2}); pool1->setPaddingNd(DimsHW{1, 1}); - debug_shape(pool1); - - /** stage 2, 3, 4 */ - ILayer* _layer = pool1; - in_ch = param.output_chn[0]; - for (int stage = 2; stage < 5; ++stage) { - int32_t out_ch = param.output_chn[stage - 1]; - std::string lname = "stage" + std::to_string(stage); - std::cout << "================ " << lname << " ================\n"; - _layer = invertedRes(net, m, *_layer->getOutput(0), lname + ".0", in_ch, out_ch, 2); - debug_shape(_layer); - for (int j = 1; j < param.repeat[stage - 2]; ++j) { - _layer = invertedRes(net, m, *_layer->getOutput(0), lname + "." + std::to_string(j), out_ch, out_ch, 1); + + ILayer* layer = pool1; + int32_t in_channels = variant.output_chn[0]; + for (int32_t stage = 2; stage < 5; ++stage) { + const int32_t out_channels = variant.output_chn[stage - 1]; + const std::string lname = "stage" + std::to_string(stage); + layer = addInvertedResidual(network, weight_map, *layer->getOutput(0), lname + ".0", in_channels, out_channels, + 2); + for (int32_t block = 1; block < variant.repeat[stage - 2]; ++block) { + layer = addInvertedResidual(network, weight_map, *layer->getOutput(0), lname + "." + std::to_string(block), + out_channels, out_channels, 1); } - in_ch = out_ch; + in_channels = out_channels; } - /** conv5, mean and fully connected layer */ - auto* conv5 = CBR(net, m, *_layer->getOutput(0), "conv5", param.output_chn[4], 1, 1, 0); - auto* mean = net->addReduce(*conv5->getOutput(0), ReduceOperation::kAVG, 0xc, false); - mean->setName("global_pool(mean)"); - auto* fcw = net->addConstant(DimsHW{1000, 1024}, m["fc.weight"]); - auto* fcb = net->addConstant(DimsHW{1, 1000}, m["fc.bias"]); - auto* _fc = net->addMatrixMultiply(*mean->getOutput(0), M::kNONE, *fcw->getOutput(0), M::kTRANSPOSE); - auto* fc = net->addElementWise(*_fc->getOutput(0), *fcb->getOutput(0), ElementWiseOperation::kSUM); + auto* conv5 = addConvBnRelu(network, weight_map, *layer->getOutput(0), "conv5", variant.output_chn[4], 1, 1, 0); + auto* global_pool = network->addReduce(*conv5->getOutput(0), ReduceOperation::kAVG, 0xc, false); + assert(global_pool); + global_pool->setName("global_pool"); + auto* fcw = network->addConstant(DimsHW{1000, variant.output_chn[4]}, weight_map["fc.weight"]); + auto* fcb = network->addConstant(DimsHW{1, 1000}, weight_map["fc.bias"]); + auto* fc_matmul = + network->addMatrixMultiply(*global_pool->getOutput(0), M::kNONE, *fcw->getOutput(0), M::kTRANSPOSE); + auto* fc = network->addElementWise(*fc_matmul->getOutput(0), *fcb->getOutput(0), ElementWiseOperation::kSUM); + assert(fc); fc->getOutput(0)->setName(NAMES[1]); - debug_shape(fc); + network->markOutput(*fc->getOutput(0)); - net->markOutput(*fc->getOutput(0)); - -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); - IHostMemory* mem = builder->buildSerializedNetwork(*net, *config); + IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); + assert(mem); ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); - delete net; + delete mem; + delete network; #else - builder->setMaxBatchSize(N); + builder->setMaxBatchSize(batch_size); config->setMaxWorkspaceSize(WORKSPACE_SIZE); - ICudaEngine* engine = builder->buildEngineWithConfig(*net, *config); - net->destroy(); + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + network->destroy(); #endif std::cout << "build finished\n"; - // Release host memory - for (auto& mem : m) { - free((void*)(mem.second.values)); + for (auto& mem : weight_map) { + std::free(const_cast(mem.second.values)); } - return engine; } -void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { - // Create builder +static void APIToModel(int32_t batch_size, IRuntime* runtime, IHostMemory** model_stream, + const ShuffleNetV2Variant& variant) { 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(N, runtime, builder, config, DataType::kFLOAT); + ICudaEngine* engine = createEngine(batch_size, runtime, builder, config, DataType::kFLOAT, variant); assert(engine != nullptr); - // Serialize the engine - (*modelStream) = engine->serialize(); + (*model_stream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -318,13 +277,14 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { #endif } -auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> std::vector> { +static auto doInference(IExecutionContext& context, void* input, + int64_t batch_size) -> std::vector> { ICudaEngine const& engine = context.getEngine(); cudaStream_t stream; CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -333,20 +293,24 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); - auto s = getSize(engine.getTensorDataType(tensor_name)); - size = s * batchSize * SIZES[i]; + const std::string name = tensor_name; + auto element_size = getSize(engine.getTensorDataType(tensor_name)); + size = element_size * static_cast(batch_size) * (name == NAMES[0] ? SIZES[0] : SIZES[1]); CHECK(cudaMalloc(&buffers[i], size)); - if (i == 0) { + if (name == NAMES[0]) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); - auto s = getSize(engine.getBindingDataType(idx)); + auto element_size = getSize(engine.getBindingDataType(idx)); assert(idx == i); - size = s * batchSize * SIZES[i]; + size = element_size * static_cast(batch_size) * SIZES[i]; CHECK(cudaMalloc(&buffers[i], size)); if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); @@ -354,134 +318,150 @@ auto doInference(IExecutionContext& context, void* input, int64_t batchSize) -> #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; - for (int i = 1; i < nIO; ++i) { - std::vector tmp(batchSize * SIZES[i], std::nanf("")); - std::size_t size = batchSize * SIZES[i] * sizeof(float); + for (int i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) + const std::string name = engine.getIOTensorName(i); + if (name == NAMES[0]) { + continue; + } + constexpr auto output_size = SIZES[1]; +#else + if (i == 0) { + continue; + } + const auto output_size = SIZES[i]; +#endif + std::vector tmp(batch_size * output_size, std::nanf("")); + const auto size = static_cast(batch_size) * output_size * sizeof(float); CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); - prob.emplace_back(tmp); + prob.emplace_back(std::move(tmp)); } CHECK(cudaStreamSynchronize(stream)); - // Release stream and buffers - CHECK(cudaStreamDestroy(stream)); + for (auto& buffer : buffers) { CHECK(cudaFree(buffer)); } + CHECK(cudaStreamDestroy(stream)); return prob; } -int main(int argc, char** argv) { +auto main(int argc, char** argv) -> int { checkTrtEnv(); - if (argc != 2) { + if (argc < 2 || argc > 3) { std::cerr << "arguments not right!\n"; - std::cerr << "./shufflenet -s // serialize model to plan file\n"; - std::cerr << "./shufflenet -d // deserialize plan file and run inference\n"; + std::cerr << "./shufflenetv2 -s [model] // serialize model to plan file\n"; + std::cerr << "./shufflenetv2 -d [model] // deserialize plan file and run inference\n"; return -1; } - // create a model using the API directly and serialize it to a stream + const auto variant = getVariantConfig(argc == 3 ? argv[2] : "shufflenet_v2_x0_5"); + std::cout << "Using ShuffleNetV2 variant: " << variant.name << "\n"; + IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); - char* trtModelStream{nullptr}; + char* trt_model_stream{nullptr}; std::streamsize size{0}; if (std::string(argv[1]) == "-s") { - IHostMemory* modelStream{nullptr}; - APIToModel(1, runtime, &modelStream); - assert(modelStream != nullptr); + IHostMemory* model_stream{nullptr}; + APIToModel(1, runtime, &model_stream, variant); + assert(model_stream != nullptr); - std::ofstream p(ENGINE_PATH, std::ios::binary | std::ios::trunc); - if (!p) { + std::ofstream plan(variant.engine_path, std::ios::binary | std::ios::trunc); + if (!plan) { std::cerr << "could not open plan output file\n"; return -1; } - if (modelStream->size() > static_cast(std::numeric_limits::max())) { + if (model_stream->size() > static_cast(std::numeric_limits::max())) { std::cerr << "this model is too large to serialize\n"; return -1; } - const auto* data_ptr = reinterpret_cast(modelStream->data()); - auto data_size = static_cast(modelStream->size()); - p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 - delete modelStream; + const auto* data_ptr = reinterpret_cast(model_stream->data()); + const auto data_size = static_cast(model_stream->size()); + plan.write(data_ptr, data_size); +#if TRT_VERSION_GE(8, 0, 0) + delete model_stream; + delete runtime; #else - modelStream->destroy(); + model_stream->destroy(); + runtime->destroy(); #endif return 0; - } else if (std::string(argv[1]) == "-d") { - std::ifstream file(ENGINE_PATH, std::ios::binary); + } + if (std::string(argv[1]) == "-d") { + std::ifstream file(variant.engine_path, 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); + trt_model_stream = new char[size]; + assert(trt_model_stream); + file.read(trt_model_stream, size); file.close(); } } else { return -1; } -#if TRT_VERSION >= 8000 - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); +#if TRT_VERSION_GE(8, 0, 0) + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size); #else - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, size, nullptr); #endif assert(engine != nullptr); IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); - delete[] trtModelStream; + delete[] trt_model_stream; - // Run inference void* input = nullptr; std::vector flat_img; - cv::Mat img; + cv::Mat img = cv::imread("assets/cats.jpg", cv::IMREAD_COLOR); if constexpr (TRT_PREPROCESS) { - // for simplicity, resize image on cpu side - img = cv::imread("../assets/cats.jpg", cv::IMREAD_COLOR); cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR); input = static_cast(img.data); } else { - img = cv::imread("../assets/cats.jpg", cv::IMREAD_COLOR); flat_img = preprocess_img(img, true, mean, stdv, N, INPUT_H, INPUT_W); input = flat_img.data(); } - for (int i = 0; i < 100; ++i) { - auto start = std::chrono::system_clock::now(); - auto prob = doInference(*context, input, N); - auto end = std::chrono::system_clock::now(); - auto period = std::chrono::duration_cast(end - start); - std::cout << period.count() << "us\n"; - - for (auto& vector : prob) { - int idx = 0; - for (auto& v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } - } - } - if (i == 99) { - std::cout << "prediction result:\n"; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << "\n"; - } - } + auto firstProb = doInference(*context, input, N); + printFirstOutputs("shufflenetv2", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + if (labels.empty()) { + std::cerr << "failed to load labels from " << LABELS_PATH << "\n"; + std::abort(); + } + int top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << top++ << " idx: " << idx << ", logits: " << std::setprecision(4) << logits + << ", label: " << labels[idx] << "\n"; + } + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int i = 0; i < kBenchmarkRuns; ++i) { + auto start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, N); + auto end = std::chrono::steady_clock::now(); + auto period = std::chrono::duration_cast(end - start); + latencies.push_back(static_cast(period.count()) / 1000.0); } -#if TRT_VERSION >= 8000 + printBenchmark("shufflenetv2", latencies, N); +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; delete runtime; diff --git a/shufflenetv2/utils.h b/shufflenetv2/utils.h index fa6010a0..4f111957 100644 --- a/shufflenetv2/utils.h +++ b/shufflenetv2/utils.h @@ -2,12 +2,19 @@ #include #include #include +#include +#include +#include #include +#include #include +#include #include #include #include #include +#include +#include #include #include "macros.h" @@ -23,7 +30,7 @@ using namespace nvinfer1; } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -35,6 +42,17 @@ static void checkTrtEnv(int device = 0) { #endif } +template +static auto toI32(T value) -> int32_t { + static_assert(std::is_integral::value, "toI32 requires an integral type."); + static_assert(sizeof(T) > sizeof(int32_t), "toI32 is only for values wider than int32_t."); + if constexpr (std::is_signed::value) { + assert(value >= static_cast(std::numeric_limits::min())); + } + assert(value <= static_cast(std::numeric_limits::max())); + return static_cast(value); +} + /** * @brief TensorRT weight files have a simple space delimited format: * [type] [size] @@ -63,7 +81,11 @@ static std::map loadWeights(const std::string& f input >> name >> std::dec >> wt.count; // Load blob - auto* val = new uint32_t[wt.count]; + auto* val = static_cast(std::malloc(sizeof(uint32_t) * static_cast(wt.count))); + if (val == nullptr) { + std::cerr << "weight allocation failed\n"; + std::abort(); + } input >> std::hex; for (auto x = 0ll; x < wt.count; ++x) { input >> val[x]; @@ -179,7 +201,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -224,7 +246,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; @@ -232,7 +254,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo static size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -249,3 +271,41 @@ static size_t getSize(DataType dt) { } } } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/squeezenet/FindTensorRT.cmake b/squeezenet/FindTensorRT.cmake index 6d542776..114bea76 100644 --- a/squeezenet/FindTensorRT.cmake +++ b/squeezenet/FindTensorRT.cmake @@ -1,11 +1,5 @@ cmake_minimum_required(VERSION 3.17.0) -set(TRT_VERSION - $ENV{TRT_VERSION} - CACHE - STRING - "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", etc") - function(_guess_path var_name required_files) set(_result "") @@ -44,56 +38,98 @@ function(_guess_path var_name required_files) PARENT_SCOPE) endfunction() -# find TensorRT include folder -if(NOT DEFINED TensorRT_INCLUDE_DIR) - if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") - _guess_path( - TensorRT_INCLUDE_DIR "NvInfer.h" "/usr/include/aarch64-linux-gnu" - "/usr/include" "/usr/local/cuda/targets/aarch64-linux/include") - else() - _guess_path( - TensorRT_INCLUDE_DIR "NvInfer.h" - "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" - "/usr/include/x86_64-linux-gnu" "/usr/include") - endif() - message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") +# add library +add_library(TensorRT IMPORTED INTERFACE) +add_library(TensorRT::TensorRT ALIAS TensorRT) + +set(TRT_VERSION + "$ENV{TRT_VERSION}" + CACHE + STRING + "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", \"8.6.1.6.Windows10.x86_64.cuda-12.0\" etc" +) + +if(NOT "${TRT_VERSION}" STREQUAL "" AND NOT "$ENV{TRT_VERSION}" STREQUAL "" + AND NOT "${TRT_VERSION}" STREQUAL "$ENV{TRT_VERSION}") + message( + WARNING + "TRT_VERSION defined by cmake and environment variable both, using the environment variable" + ) endif() -# find TensorRT library folder -if(NOT TensorRT_LIBRARY_DIR) - if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") - _guess_path( - TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so" - "/usr/lib/aarch64-linux-gnu;/usr/lib/aarch64-linux-gnu/tegra" "/usr/lib") - else() - _guess_path( - TensorRT_LIBRARY_DIR - "libnvinfer.so;libnvinfer_plugin.so" - "/usr/lib/x86_64-linux-gnu;/usr/local/tensorrt/targets/x86_64-linux-gnu/lib;/usr/lib" +if(NOT "$ENV{TRT_VERSION}" STREQUAL "") + set(TRT_VERSION + "$ENV{TRT_VERSION}" + CACHE STRING "TensorRT version" FORCE) +endif() + +if("${TRT_VERSION}" STREQUAL "") + message(FATAL_ERROR "Please set the TRT_VERSION environment variable") +endif() + +string(REGEX MATCH "^[0-9]+" _match "${TRT_VERSION}") +set(TRT_MAJOR_VERSION "${_match}") +unset(_match) +if("${TRT_MAJOR_VERSION}" STREQUAL "") + message(FATAL_ERROR "Failed to parse TensorRT major version from ${TRT_VERSION}") +endif() + +if(WIN32) + set(TensorRT_DIR "C:/Program Files/TensorRT-${TRT_VERSION}") + if(NOT EXISTS "${TensorRT_DIR}") + message( + FATAL_ERROR + "TensorRT_DIR=${TensorRT_DIR} does not exist!" ) endif() - message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") -endif() -set(TensorRT_LIBRARIES) + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) + message(DEBUG "Using ${_modules}") + else() + set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch + nvinfer_lean) + endif() -message(STATUS "Found TensorRT lib: ${TensorRT_LIBRARIES}") + set(TensorRT_LIBRARY_DIR "${TensorRT_DIR}/lib") + set(TensorRT_INCLUDE_DIR "${TensorRT_DIR}/include") +elseif(UNIX) + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _trt_arch) + set(_trt_include_candidates) + if(_trt_arch MATCHES "^(aarch64|arm64|arch64)$") + set(_trt_include_candidates "/usr/include/aarch64-linux-gnu" "/usr/include" + "/usr/local/cuda/targets/aarch64-linux/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/aarch64-linux-gnu/lib" + "/usr/lib/aarch64-linux-gnu" "/usr/lib/aarch64-linux-gnu/tegra" + "/usr/lib") + elseif(_trt_arch MATCHES "^(x86_64|amd64)$") + set(_trt_include_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" + "/usr/include/x86_64-linux-gnu" "/usr/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" + "/usr/lib/x86_64-linux-gnu" "/usr/lib") + else() + message(FATAL_ERROR "Unknown architecture") + endif() -# process for different TensorRT version -if(DEFINED TRT_VERSION AND NOT TRT_VERSION STREQUAL "") - string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION}) - set(TRT_MAJOR_VERSION "${_match}") set(_modules nvinfer nvinfer_plugin) - unset(_match) - - if(TRT_MAJOR_VERSION GREATER_EQUAL 8) + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 8) list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean) endif() -else() - message(FATAL_ERROR "Please set a environment variable \"TRT_VERSION\"") + + _guess_path(TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so" + ${_trt_library_candidates}) + message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") + _guess_path(TensorRT_INCLUDE_DIR "NvInfer.h" ${_trt_include_candidates}) + message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") endif() -# find and add all modules of TensorRT into list foreach(lib IN LISTS _modules) find_library( TensorRT_${lib}_LIBRARY @@ -102,11 +138,10 @@ foreach(lib IN LISTS _modules) list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY}) endforeach() -# make the "TensorRT target" -add_library(TensorRT IMPORTED INTERFACE) -add_library(TensorRT::TensorRT ALIAS TensorRT) target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES}) +message(STATUS "Found TensorRT libs: ${TensorRT_LIBRARIES}") + set_target_properties( TensorRT PROPERTIES C_STANDARD 17 @@ -118,4 +153,8 @@ set_target_properties( INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") unset(TRT_MAJOR_VERSION) -unset(_modules) +unset(_modules) +unset(_trt_lib_suffix) +unset(_trt_include_candidates) +unset(_trt_library_candidates) +unset(_trt_arch) diff --git a/squeezenet/gen_wts.py b/squeezenet/gen_wts.py index 8ed0ec44..816cd54a 100644 --- a/squeezenet/gen_wts.py +++ b/squeezenet/gen_wts.py @@ -1,72 +1,44 @@ +import os import struct +from pathlib import Path -import cv2 -import numpy as np import torch -import torchvision +from torchvision.models import SqueezeNet1_1_Weights, squeezenet1_1 -def read_imagenet_labels() -> dict[int, str]: - """ - read ImageNet 1000 labels +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +MODELS_DIR = REPO_ROOT / "models" +WTS_PATH = MODELS_DIR / "squeezenet.wts" - Returns: - dict[int, str]: labels dict - """ - clsid2label = {} - with open("../assets/imagenet1000_clsidx_to_labels.txt", "r") as f: - for i in f.readlines(): - k, v = i.split(": ") - clsid2label.setdefault(int(k), v[1:-3]) - return clsid2label +def print_cache_env() -> None: + for key in ("TORCH_HOME", "HF_HOME"): + value = os.environ.get(key) + if value: + print(f"{key}={value}") -def preprocess(img: np.array) -> torch.Tensor: - """ - a preprocess method align with ImageNet dataset - - Args: - img (np.array): input image - - Returns: - torch.Tensor: preprocessed image in `NCHW` layout - """ - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 - img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR) - mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) - std = np.array([0.229, 0.224, 0.225], dtype=np.float32) - img = (img - mean) / std - img = img.transpose(2, 0, 1)[None, ...] - return torch.from_numpy(img) +def export_wts(model: torch.nn.Module, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + state_dict = model.state_dict() + with open(path, "w") as f: + f.write(f"{len(state_dict)}\n") + for key, value in state_dict.items(): + values = value.reshape(-1).cpu().numpy() + f.write(f"{key} {len(values)}") + print(key, value.shape) + for item in values: + f.write(" ") + f.write(struct.pack(">f", float(item)).hex()) + f.write("\n") -def main(): - labels = read_imagenet_labels() - model = torchvision.models.squeezenet1_1(pretrained=True) +def main() -> None: + print_cache_env() + model = squeezenet1_1(weights=SqueezeNet1_1_Weights.DEFAULT) model = model.eval() - - img = cv2.imread("../assets/cats.jpg", cv2.IMREAD_COLOR) - img = preprocess(img) - - with torch.inference_mode(): - output = model(img) - for i, batch in enumerate(torch.topk(output, k=3).indices): - for j, idx in enumerate(batch): - print(f"\tBatch: {i}, Top: {j}, logits: {output[i][idx]:.4f}, label: {labels[int(idx)]}") - print(f"{'=' * 32}") - - with open("../models/squeezenet.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))) - print(k, v.shape) - for vv in vr: - f.write(" ") - f.write(struct.pack(">f", float(vv)).hex()) - f.write("\n") - f.close() + export_wts(model, WTS_PATH) if __name__ == "__main__": diff --git a/squeezenet/macros.h b/squeezenet/macros.h index dc9a4d57..6d851752 100644 --- a/squeezenet/macros.h +++ b/squeezenet/macros.h @@ -1,4 +1,5 @@ #pragma once +#include #ifdef API_EXPORTS #if defined(_MSC_VER) @@ -15,10 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION >= 8000 +#if TRT_VERSION_LT(7, 2, 2) +#error "TensorRT >= 7.2.2 is required for this demo." +#endif + +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/squeezenet/squeezenet.cpp b/squeezenet/squeezenet.cpp index 381718c0..30daf500 100644 --- a/squeezenet/squeezenet.cpp +++ b/squeezenet/squeezenet.cpp @@ -1,85 +1,98 @@ #include +#include #include #include +#include +#include +#include #include +#include #include +#include #include #include +#include #include #include "logging.h" #include "utils.h" -// stuff we know about squeezenet -static constexpr const int N = 1; -static constexpr const int INPUT_H = 224; -static constexpr const int INPUT_W = 224; -static constexpr const int SIZES[] = {3 * INPUT_H * INPUT_W, N * 1000}; -static constexpr const char* NAMES[] = {"data", "prob"}; -static constexpr const bool TRT_PREPROCESS = TRT_VERSION >= 8510 ? true : false; -static constexpr const float mean[3] = {0.485f, 0.456f, 0.406f}; -static constexpr const float stdv[3] = {0.229f, 0.224f, 0.225f}; +static constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; -static constexpr const char* WTS_PATH = "../models/squeezenet.wts"; -static constexpr const char* ENGINE_PATH = "../models/squeezenet.engine"; -static constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; +static constexpr const int64_t N = 1; +static constexpr const int32_t INPUT_H = 224; +static constexpr const int32_t INPUT_W = 224; +static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, N * 1000}; +static constexpr const std::array NAMES = {"data", "prob"}; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1) ? true : false; +static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; +static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; + +static constexpr const char* WTS_PATH = "models/squeezenet.wts"; +static constexpr const char* ENGINE_PATH = "models/squeezenet.engine"; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; using namespace nvinfer1; using WeightMap = std::map; +using NDCF = nvinfer1::NetworkDefinitionCreationFlag; static Logger gLogger; -ILayer* fire(INetworkDefinition* network, WeightMap& m, ITensor& input, const std::string& lname, - int32_t squeeze_planes, int32_t e1x1_planes, int32_t e3x3_planes) { - auto* conv1 = network->addConvolutionNd(input, squeeze_planes, DimsHW{1, 1}, m[lname + "squeeze.weight"], - m[lname + "squeeze.bias"]); +static auto fire(INetworkDefinition* network, WeightMap& weights, ITensor& input, const std::string& lname, + int32_t squeeze_planes, int32_t e1x1_planes, int32_t e3x3_planes) -> ILayer* { + auto* conv1 = network->addConvolutionNd(input, squeeze_planes, DimsHW{1, 1}, weights.at(lname + "squeeze.weight"), + weights.at(lname + "squeeze.bias")); assert(conv1); - auto* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU)->getOutput(0); + auto* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + assert(relu1); std::string _c = lname + "expand1x1"; - auto* conv2 = network->addConvolutionNd(*relu1, e1x1_planes, DimsHW{1, 1}, m[_c + ".weight"], m[_c + ".bias"]); + auto* conv2 = network->addConvolutionNd(*relu1->getOutput(0), e1x1_planes, DimsHW{1, 1}, weights.at(_c + ".weight"), + weights.at(_c + ".bias")); assert(conv2); auto* relu2 = network->addActivation(*conv2->getOutput(0), ActivationType::kRELU); assert(relu2); _c = lname + "expand3x3"; - auto* conv3 = network->addConvolutionNd(*relu1, e3x3_planes, DimsHW{3, 3}, m[_c + ".weight"], m[_c + ".bias"]); + auto* conv3 = network->addConvolutionNd(*relu1->getOutput(0), e3x3_planes, DimsHW{3, 3}, weights.at(_c + ".weight"), + weights.at(_c + ".bias")); assert(conv3); conv3->setPaddingNd(DimsHW{1, 1}); auto* relu3 = network->addActivation(*conv3->getOutput(0), ActivationType::kRELU); assert(relu3); - ITensor* inputTensors[] = {relu2->getOutput(0), relu3->getOutput(0)}; - auto* concat = network->addConcatenation(inputTensors, 2); + std::array inputTensors = {relu2->getOutput(0), relu3->getOutput(0)}; + auto* concat = network->addConcatenation(inputTensors.data(), 2); assert(concat); return concat; } -// Creat the engine using only the API and not any parser. -ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt) { +// Create the engine using only the API and not any parser. +static auto createEngine(int32_t batch_size, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, + DataType dt) -> ICudaEngine* { auto weightMap = loadWeights(WTS_PATH); -#if TRT_VERSION >= 10000 - auto* network = builder->createNetworkV2(0); + +#if TRT_VERSION_GE(10, 12, 0) + auto flag = 1U << static_cast(NDCF::kSTRONGLY_TYPED); +#elif TRT_VERSION_GE(10, 0, 0) + auto flag = 0U; #else - auto* network = builder->createNetworkV2(1u << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)); + auto flag = 1U << static_cast(NDCF::kEXPLICIT_BATCH); #endif + auto* network = builder->createNetworkV2(flag); ITensor* data{nullptr}; if constexpr (TRT_PREPROCESS) { -#if TRT_VERSION > 8510 dt = DataType::kUINT8; -#else - dt = DataType::kINT8; -#endif - data = network->addInput(NAMES[0], dt, Dims4{N, INPUT_H, INPUT_W, 3}); + data = network->addInput(NAMES[0], dt, Dims4{batch_size, INPUT_H, INPUT_W, 3}); auto* trans = addTransformLayer(network, *data, true, mean, stdv); data = trans->getOutput(0); } else { - data = network->addInput(NAMES[0], dt, Dims4{N, 3, INPUT_H, INPUT_W}); + data = network->addInput(NAMES[0], dt, Dims4{batch_size, 3, INPUT_H, INPUT_W}); } assert(data); - auto* conv1 = network->addConvolutionNd(*data, 64, DimsHW{3, 3}, weightMap["features.0.weight"], - weightMap["features.0.bias"]); + auto* conv1 = network->addConvolutionNd(*data, 64, DimsHW{3, 3}, weightMap.at("features.0.weight"), + weightMap.at("features.0.bias")); assert(conv1); conv1->setStrideNd(DimsHW{2, 2}); auto* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); @@ -113,8 +126,8 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil cat1 = fire(network, weightMap, *cat1->getOutput(0), "features.12.", 64, 256, 256); // classifier - auto* conv2 = network->addConvolutionNd(*cat1->getOutput(0), 1000, DimsHW{1, 1}, weightMap["classifier.1.weight"], - weightMap["classifier.1.bias"]); + auto* conv2 = network->addConvolutionNd(*cat1->getOutput(0), 1000, DimsHW{1, 1}, + weightMap.at("classifier.1.weight"), weightMap.at("classifier.1.bias")); assert(conv2); auto* relu2 = network->addActivation(*conv2->getOutput(0), ActivationType::kRELU); assert(relu2); @@ -125,40 +138,54 @@ ICudaEngine* createEngine(int32_t N, IRuntime* runtime, IBuilder* builder, IBuil network->markOutput(*pool4->getOutput(0)); // Build engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); IHostMemory* mem = builder->buildSerializedNetwork(*network, *config); auto* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); + delete mem; delete network; #else - builder->setMaxBatchSize(N); + builder->setMaxBatchSize(batch_size); config->setMaxWorkspaceSize(WORKSPACE_SIZE); auto* engine = builder->buildEngineWithConfig(*network, *config); network->destroy(); #endif - std::cout << "build out" << std::endl; + std::cout << "build finished\n"; // Release host memory for (auto& mem : weightMap) { - free((void*)(mem.second.values)); + delete[] static_cast(mem.second.values); } return engine; } -void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { - // Create builder +static void destroyRuntime(IRuntime* runtime) { +#if TRT_VERSION_GE(8, 0, 0) + delete runtime; +#else + runtime->destroy(); +#endif +} + +static void destroyHostMemory(IHostMemory* memory) { +#if TRT_VERSION_GE(8, 0, 0) + delete memory; +#else + memory->destroy(); +#endif +} + +static void APIToModel(int32_t batch_size, IRuntime* runtime, IHostMemory** modelStream) { auto* builder = createInferBuilder(gLogger); auto* config = builder->createBuilderConfig(); - // Create model to populate the network, then set the outputs and create an engine - auto* engine = createEngine(N, runtime, builder, config, DataType::kFLOAT); + auto* engine = createEngine(batch_size, runtime, builder, config, DataType::kFLOAT); assert(engine != nullptr); - // Serialize the engine (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -169,13 +196,14 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { #endif } -std::vector> doInference(IExecutionContext& context, void* input, int32_t batch_size) { +static auto doInference(IExecutionContext& context, void* input, + int64_t batch_size) -> std::vector> { const auto& engine = context.getEngine(); cudaStream_t stream; CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); @@ -184,7 +212,7 @@ std::vector> doInference(IExecutionContext& context, void* in buffers.resize(nIO); for (auto i = 0; i < nIO; ++i) { std::size_t size = 0; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const auto* tensor_name = engine.getIOTensorName(i); auto s = getSize(engine.getTensorDataType(tensor_name)); size = s * batch_size * SIZES[i]; @@ -192,7 +220,10 @@ std::vector> doInference(IExecutionContext& context, void* in if (i == 0) { CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #else const int32_t idx = engine.getBindingIndex(NAMES[i]); auto s = getSize(engine.getBindingDataType(idx)); @@ -205,22 +236,28 @@ std::vector> doInference(IExecutionContext& context, void* in #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; for (int i = 1; i < nIO; ++i) { - std::vector tmp(batch_size * SIZES[i], std::nan("")); + std::vector tmp(batch_size * SIZES[i], std::nanf("")); std::size_t size = batch_size * SIZES[i] * sizeof(float); CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); prob.emplace_back(tmp); } CHECK(cudaStreamSynchronize(stream)); - cudaStreamDestroy(stream); + CHECK(cudaStreamDestroy(stream)); for (auto i = 0; i < nIO; ++i) { CHECK(cudaFree(buffers[i])); } @@ -230,9 +267,9 @@ std::vector> doInference(IExecutionContext& context, void* in int main(int argc, char** argv) { checkTrtEnv(); if (argc != 2) { - std::cerr << "arguments not right!" << std::endl; - std::cerr << "./squeezenet -s // serialize model to plan file" << std::endl; - std::cerr << "./squeezenet -d // deserialize plan file and run inference" << std::endl; + std::cerr << "arguments not right!\n"; + std::cerr << "./squeezenet -s // serialize model to plan file\n"; + std::cerr << "./squeezenet -d // deserialize plan file and run inference\n"; return -1; } @@ -240,7 +277,7 @@ int main(int argc, char** argv) { auto* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); char* trtModelStream{nullptr}; - size_t size{0}; + std::streamsize size{0}; if (std::string(argv[1]) == "-s") { IHostMemory* modelStream{nullptr}; @@ -249,15 +286,22 @@ int main(int argc, char** argv) { std::ofstream p(ENGINE_PATH, std::ios::binary | std::ios::trunc); if (!p) { - std::cerr << "could not open plan output file" << std::endl; + std::cerr << "could not open plan output file\n"; + destroyHostMemory(modelStream); + destroyRuntime(runtime); return -1; } - p.write(reinterpret_cast(modelStream->data()), modelStream->size()); -#if TRT_VERSION >= 8000 - delete modelStream; -#else - modelStream->destroy(); -#endif + if (modelStream->size() > static_cast(std::numeric_limits::max())) { + std::cerr << "this model is too large to serialize\n"; + destroyHostMemory(modelStream); + destroyRuntime(runtime); + return -1; + } + const auto* data_ptr = reinterpret_cast(modelStream->data()); + auto data_size = static_cast(modelStream->size()); + p.write(data_ptr, data_size); + destroyHostMemory(modelStream); + destroyRuntime(runtime); return 0; } else if (std::string(argv[1]) == "-d") { std::ifstream file(ENGINE_PATH, std::ios::binary); @@ -269,12 +313,17 @@ int main(int argc, char** argv) { assert(trtModelStream); file.read(trtModelStream, size); file.close(); + } else { + std::cerr << "could not open engine file\n"; + destroyRuntime(runtime); + return -1; } } else { + destroyRuntime(runtime); return -1; } -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* engine = runtime->deserializeCudaEngine(trtModelStream, size); #else auto* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); @@ -288,54 +337,47 @@ int main(int argc, char** argv) { cv::Mat img; if constexpr (TRT_PREPROCESS) { // for simplicity, resize image on cpu side - img = cv::imread("../assets/cats.jpg", cv::IMREAD_COLOR); + img = cv::imread("assets/cats.jpg", cv::IMREAD_COLOR); + assert(!img.empty()); cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR); input = static_cast(img.data); } else { - img = cv::imread("../assets/cats.jpg", cv::IMREAD_COLOR); + img = cv::imread("assets/cats.jpg", cv::IMREAD_COLOR); + assert(!img.empty()); flat_img = preprocess_img(img, true, mean, stdv, N, INPUT_H, INPUT_W); input = flat_img.data(); } + assert(input); + + auto firstProb = doInference(*context, input, N); + printFirstOutputs("squeezenet", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + int _top = 0; + for (const auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << labels[idx] + << "\n"; + } - for (int32_t i = 0; i < 100; ++i) { - auto _start = std::chrono::system_clock::now(); - auto prob = doInference(*context, input, N); - auto _end = std::chrono::system_clock::now(); + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, N); + auto _end = std::chrono::steady_clock::now(); auto _time = std::chrono::duration_cast(_end - _start).count(); - std::cout << "Execution time: " << _time << "us" << std::endl; - - for (auto vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====" << std::endl; - break; - } - } - } - - if (i == 99) { - std::cout << "prediction result: " << std::endl; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << std::endl; - } - } + latencies.push_back(static_cast(_time) / 1000.0); } + printBenchmark("squeezenet", latencies, N); delete[] trtModelStream; - // Destroy the engine -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete context; delete engine; - delete runtime; #else context->destroy(); engine->destroy(); - runtime->destroy(); #endif + destroyRuntime(runtime); return 0; } diff --git a/squeezenet/utils.h b/squeezenet/utils.h index 61681b48..bdfd66ea 100644 --- a/squeezenet/utils.h +++ b/squeezenet/utils.h @@ -1,41 +1,42 @@ #pragma once -#include #include #include +#include #include +#include +#include +#include #include +#include #include #include +#include #include #include -#include #include #include +#include "macros.h" using namespace nvinfer1; -#define WORKSPACE_SIZE (16 << 20) - -#define CHECK(status) \ - do { \ - auto ret = (status); \ - if (ret != cudaSuccess) { \ - std::cerr << "Cuda failure: " << ret << std::endl; \ - abort(); \ - } \ +#define CHECK(status) \ + do { \ + auto ret = (status); \ + if (ret != cudaSuccess) { \ + std::cerr << "Cuda failure: " << ret << "\n"; \ + std::abort(); \ + } \ } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 7220 -#error "TensorRT >= 7.2.2 is required for this demo." -#endif -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); const int sm = prop.major * 10 + prop.minor; if (sm > 86) { - throw std::runtime_error("TensorRT < 8 does not support SM > 86 on this GPU."); + std::cerr << "TensorRT < 8 does not support SM > 86 on this GPU."; + std::abort(); } #endif } @@ -48,7 +49,7 @@ static void checkTrtEnv(int device = 0) { * @return std::map */ static std::map loadWeights(const std::string& file) { - std::cout << "Loading weights: " << file << std::endl; + std::cout << "Loading weights: " << file << "\n"; std::map weightMap; // Open weights file @@ -68,9 +69,10 @@ static std::map loadWeights(const std::string& f input >> name >> std::dec >> wt.count; // Load blob - uint32_t* val = reinterpret_cast(malloc(sizeof(val) * wt.count)); - for (uint32_t x = 0; x < wt.count; ++x) { - input >> std::hex >> val[x]; + auto* val = new uint32_t[static_cast(wt.count)]; + input >> std::hex; + for (auto x = 0ll; x < wt.count; ++x) { + input >> val[x]; } wt.values = val; weightMap[name] = wt; @@ -91,12 +93,13 @@ static std::map loadWeights(const std::string& f * @param w resize width * @return std::vector contiguous flatten image data in float32 type */ -static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const float mean[3], const float std[3], int n, - int h, int w) { - const int c = img.channels(); - const std::size_t size = c * h * w; +static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const std::array& mean, + const std::array& std, int n, int h, int w) { + const auto c = img.channels(); + const auto size = c * h * w; if (c != 3) { - throw std::runtime_error("this demo only supports 3 channel input image."); + std::cerr << "this demo only supports 3 channel input image.\n"; + std::abort(); } if (bgr2rgb) { cv::cvtColor(img, img, cv::COLOR_BGR2RGB); @@ -104,7 +107,7 @@ static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const float cv::resize(img, img, cv::Size(w, h), 0, 0, cv::INTER_LINEAR); img.convertTo(img, CV_32FC3, 1.f / 255); img = (img - cv::Scalar(mean[0], mean[1], mean[2])) / cv::Scalar(std[0], std[1], std[2]); - std::vector chw(n * c * h * w, 0.f); + std::vector chw(static_cast(n) * c * h * w, 0.f); // fill all batch with the same input image for (int i = 0; i < n; ++i) { @@ -112,27 +115,27 @@ static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const float for (int x = 0; x < w; ++x) { const cv::Vec3f v = img.at(y, x); chw[i * size + 0 * h * w + y * w + x] = v[0]; - chw[i * size + 1 * h * w + y * h + x] = v[1]; - chw[i * size + 2 * h * w + y * h + x] = v[2]; + chw[i * size + 1 * h * w + y * w + x] = v[1]; + chw[i * size + 2 * h * w + y * w + x] = v[2]; } } } return chw; } -static std::vector> topk(const std::vector& v, int k) { +static auto topk(const std::vector& v, int k) -> std::vector> { if (k <= 0) return {}; - k = std::min(k, v.size()); + auto stride = std::min(k, static_cast(v.size())); std::vector idx(v.size()); std::iota(idx.begin(), idx.end(), 0); - std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&](int a, int b) { return v[a] > v[b]; }); + std::partial_sort(idx.begin(), idx.begin() + stride, idx.end(), [&](int a, int b) { return v[a] > v[b]; }); std::vector> out; - out.reserve(k); - for (int i = 0; i < k; ++i) + out.reserve(static_cast(stride)); + for (auto i = 0; i < stride; ++i) out.emplace_back(idx[i], v[idx[i]]); return out; } @@ -163,8 +166,8 @@ static std::map loadImagenetLabelMap(const std::string& path) return labels; } -static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bool bgr2rgb, const float mean[3], - const float std[3]) { +static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bool bgr2rgb, + const std::array& mean, const std::array& std) { struct ScaleParams { std::array shift; std::array scale; @@ -182,7 +185,7 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo ITensor* in = &input; if (input.getType() != DataType::kFLOAT) { -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) auto* cast = network->addCast(input, DataType::kFLOAT); assert(cast); cast->setName("Cast to FP32"); @@ -195,7 +198,6 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo in = identity->getOutput(0); #endif } - // Convert from NHWC to NCHW auto* perm = network->addShuffle(*in); assert(perm); @@ -212,11 +214,10 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* _slice = network->addSlice(*perm->getOutput(0), start, size, stride); _slice->setName(name); assert(_slice && _slice->getNbOutputs() == 1); - auto d = _slice->getOutput(0)->getDimensions(); return _slice->getOutput(0); }; - ITensor* channels[] = {add_slice(2, "R"), add_slice(1, "G"), add_slice(0, "B")}; - auto* cat = network->addConcatenation(channels, 3); + std::array channels = {add_slice(2, "R"), add_slice(1, "G"), add_slice(0, "B")}; + auto* cat = network->addConcatenation(channels.data(), 3); assert(cat); cat->setName("RGB"); cat->setAxis(1); @@ -229,15 +230,15 @@ static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bo auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); assert(trans); trans->setName("mean & std"); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) trans->setChannelAxis(1); #endif return trans; } -static size_t getSize(DataType dt) { +static std::size_t getSize(DataType dt) { switch (dt) { -#if TRT_VERSION >= 8510 +#if TRT_VERSION_GE(8, 5, 1) case DataType::kUINT8: #endif case DataType::kINT8: @@ -248,7 +249,47 @@ static size_t getSize(DataType dt) { return sizeof(int16_t); case DataType::kINT32: return sizeof(int32_t); - default: - throw std::runtime_error("Unsupported data type"); + default: { + std::cerr << "Unsupported data type\n"; + std::abort(); + } + } +} + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; } + std::cout << '\n'; } diff --git a/vgg/CMakeLists.txt b/vgg/CMakeLists.txt index 4c336c72..28ad39af 100644 --- a/vgg/CMakeLists.txt +++ b/vgg/CMakeLists.txt @@ -1,25 +1,37 @@ -cmake_minimum_required(VERSION 2.6) +cmake_minimum_required(VERSION 3.14) -project(vgg) +project( + vgg + VERSION 0.1 + LANGUAGES C CXX CUDA) -add_definitions(-std=c++11) +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90 100 120) +endif() -option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) -set(CMAKE_CXX_STANDARD 11) -set(CMAKE_BUILD_TYPE Debug) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_INCLUDE_CURRENT_DIR TRUE) -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/) +option(CUDA_USE_STATIC_CUDA_RUNTIME "Use static cudaruntime library" OFF) -add_executable(vgg ${PROJECT_SOURCE_DIR}/vgg11.cpp) -target_link_libraries(vgg nvinfer) -target_link_libraries(vgg cudart) +find_package(Threads REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(OpenCV REQUIRED) -add_definitions(-O2 -pthread) +if(NOT TARGET TensorRT::TensorRT) + include(FindTensorRT.cmake) +else() + message("TensorRT has been found, skipping for ${PROJECT_NAME}") +endif() +add_executable(${PROJECT_NAME} vgg.cc) + +target_include_directories(${PROJECT_NAME} PUBLIC ${CMAKE_CURRENT_LIST_DIR} + ${OpenCV_INCLUDE_DIRS}) + +target_link_libraries(${PROJECT_NAME} PUBLIC Threads::Threads CUDA::cudart + TensorRT::TensorRT ${OpenCV_LIBS}) diff --git a/vgg/FindTensorRT.cmake b/vgg/FindTensorRT.cmake new file mode 100644 index 00000000..61203d17 --- /dev/null +++ b/vgg/FindTensorRT.cmake @@ -0,0 +1,149 @@ +cmake_minimum_required(VERSION 3.17.0) + +function(_guess_path var_name required_files) + set(_result "") + + foreach(path_entry IN LISTS ARGN) + if(NOT EXISTS "${path_entry}") + message(DEBUG "skip non-existing path '${path_entry}'") + continue() + endif() + + set(_ok TRUE) + foreach(required_file IN LISTS required_files) + if(NOT EXISTS "${path_entry}/${required_file}") + set(_ok FALSE) + message(DEBUG "'${path_entry}' missing '${required_file}'") + break() + endif() + endforeach() + + if(_ok) + list(APPEND _result "${path_entry}") + message(DEBUG "accept '${path_entry}'") + else() + message(DEBUG "reject '${path_entry}'") + endif() + endforeach() + + if(_result STREQUAL "") + message( + FATAL_ERROR + "_guess_path(${var_name}) failed: no valid path found. required_files='${required_files}' candidates='${ARGN}'" + ) + endif() + + set(${var_name} + "${_result}" + PARENT_SCOPE) +endfunction() + +# add library +add_library(TensorRT IMPORTED INTERFACE) +add_library(TensorRT::TensorRT ALIAS TensorRT) + +set(TRT_VERSION + CACHE + STRING + "TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", \"8.6.1.6.Windows10.x86_64.cuda-12.0\" etc" +) + +if(NOT TRT_VERSION STREQUAL "" AND NOT $ENV{TRT_VERSION} STREQUAL "") + message( + WARNING + "TRT_VERSION defined by cmake and environment variable both, using the later one" + ) +endif() + +if(NOT $ENV{TRT_VERSION} STREQUAL "") + set(TRT_VERSION $ENV{TRT_VERSION}) +endif() + +string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION}) +set(TRT_MAJOR_VERSION "${_match}") +unset(_match) + +if(WIN32) + set(TensorRT_DIR "C:/Program Files/TensorRT-${TRT_VERSION}") + if(NOT EXISTS "${TensorRT_DIR}") + message( + FATAL_ERROR + "TensorRT_DIR=${TensorRT_DIR} does not exist!" + ) + endif() + + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) + message(DEBUG "Using ${_modules}") + else() + set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch + nvinfer_lean) + endif() + + set(TensorRT_LIBRARY_DIR "${TensorRT_DIR}/lib") + set(TensorRT_INCLUDE_DIR "${TensorRT_DIR}/include") +elseif(UNIX) + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _trt_arch) + set(_trt_include_candidates) + if(_trt_arch MATCHES "^(aarch64|arm64|arch64)$") + set(_trt_include_candidates "/usr/include/aarch64-linux-gnu" "/usr/include" + "/usr/local/cuda/targets/aarch64-linux/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/aarch64-linux-gnu/lib" + "/usr/lib/aarch64-linux-gnu" "/usr/lib/aarch64-linux-gnu/tegra" + "/usr/lib") + elseif(_trt_arch MATCHES "^(x86_64|amd64)$") + set(_trt_include_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/include" + "/usr/include/x86_64-linux-gnu" "/usr/include") + set(_trt_library_candidates + "/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" + "/usr/lib/x86_64-linux-gnu" "/usr/lib") + else() + message(FATAL_ERROR "Unknown architecture") + endif() + + set(_modules nvinfer nvinfer_plugin) + if(${TRT_MAJOR_VERSION} GREATER_EQUAL 8) + list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean) + endif() + + _guess_path(TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so" + ${_trt_library_candidates}) + message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}") + _guess_path(TensorRT_INCLUDE_DIR "NvInfer.h" ${_trt_include_candidates}) + message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}") +endif() + +foreach(lib IN LISTS _modules) + find_library( + TensorRT_${lib}_LIBRARY + NAMES ${lib} + HINTS ${TensorRT_LIBRARY_DIR}) + list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY}) +endforeach() + +target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES}) + +message(STATUS "Found TensorRT libs: ${TensorRT_LIBRARIES}") + +set_target_properties( + TensorRT + PROPERTIES C_STANDARD 17 + CXX_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + SKIP_BUILD_RPATH TRUE + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN" + INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}") + +unset(TRT_MAJOR_VERSION) +unset(_modules) +unset(_trt_lib_suffix) +unset(_trt_include_candidates) +unset(_trt_library_candidates) +unset(_trt_arch) diff --git a/vgg/README.md b/vgg/README.md index f29afd09..1debe217 100644 --- a/vgg/README.md +++ b/vgg/README.md @@ -1,33 +1,101 @@ -# vgg +# VGG -VGG 11-layer model (configuration "A") from - "Very Deep Convolutional Networks For Large-Scale Image Recognition" +## Introduction -For the Pytorch implementation, you can refer to [pytorchx/vgg](https://github.com/wang-xinyu/pytorchx/tree/master/vgg) +This is a TensorRT implementation of torchvision VGG variants from [Very Deep Convolutional Networks For Large-Scale Image Recognition](https://arxiv.org/pdf/1409.1556.pdf). It supports the common torchvision VGG11/VGG13/VGG16/VGG19 models and their batch-normalized variants. -VGG's architecture is simple, just some conv, relu, maxpool, and fc layers. +VGG's architecture is simple, just some conv, relu, maxpool, and fc layers, e.g., for VGG11: +```bash +VGG( + (features): Sequential( + (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (1): ReLU(inplace=True) + (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) + (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (4): ReLU(inplace=True) + (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) + (6): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (7): ReLU(inplace=True) + (8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (9): ReLU(inplace=True) + (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) + (11): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (12): ReLU(inplace=True) + (13): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (14): ReLU(inplace=True) + (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) + (16): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (17): ReLU(inplace=True) + (18): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + (19): ReLU(inplace=True) + (20): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) + ) + (avgpool): AdaptiveAvgPool2d(output_size=(7, 7)) + (classifier): Sequential( + (0): Linear(in_features=25088, out_features=4096, bias=True) + (1): ReLU(inplace=True) + (2): Dropout(p=0.5, inplace=False) + (3): Linear(in_features=4096, out_features=4096, bias=True) + (4): ReLU(inplace=True) + (5): Dropout(p=0.5, inplace=False) + (6): Linear(in_features=4096, out_features=1000, bias=True) + ) +) ``` -// 1. generate vgg.wts from [pytorchx/vgg](https://github.com/wang-xinyu/pytorchx/tree/master/vgg) -// 2. put vgg.wts into tensorrtx/vgg +## Usage -// 3. build and run +1. Set the cache locations used by PyTorch/Hugging Face. The export script only reads these paths from environment variables. -cd tensorrtx/vgg +```bash +export TORCH_HOME=/path/to/torch-cache +export HF_HOME=/path/to/hf-cache +``` + +2. Run `gen_wts.py` to generate `.wts` files. No inference is run during export. + +```bash +python gen_wts.py --model vgg11 +``` + +Omit `--model` to export all supported variants, or repeat it for multiple variants. -mkdir build +3. Supported model names: -cd build +| model name | torchvision config | batch norm | +| ---------- | ------------------ | ---------- | +| vgg11 | A | no | +| vgg11_bn | A | yes | +| vgg13 | B | no | +| vgg13_bn | B | yes | +| vgg16 | D | no | +| vgg16_bn | D | yes | +| vgg19 | E | no | +| vgg19_bn | E | yes | -cmake .. +4. build C++ code -make +```bash +pushd tensorrtx/vgg +cmake -S . -B build -G Ninja --fresh +cmake --build build +``` -sudo ./vgg -s // serialize model to plan file i.e. 'vgg.engine' -sudo ./vgg -d // deserialize plan file and run inference +5. serialize wts model to engine file -// 4. see if the output is same as pytorchx/vgg +```bash +./build/vgg -s ``` +Pass a model name to serialize a non-default variant: + +```bash +./build/vgg -s vgg16_bn +``` +6. run inference + +```bash +./build/vgg -d vgg11 +``` diff --git a/vgg/gen_wts.py b/vgg/gen_wts.py new file mode 100644 index 00000000..934fd7ba --- /dev/null +++ b/vgg/gen_wts.py @@ -0,0 +1,80 @@ +import argparse +import os +import struct +from pathlib import Path + +import torch +from torchvision import models + + +MODEL_NAMES = ( + "vgg11", + "vgg11_bn", + "vgg13", + "vgg13_bn", + "vgg16", + "vgg16_bn", + "vgg19", + "vgg19_bn", +) + + +def require_cache_env() -> None: + missing = [name for name in ("TORCH_HOME", "HF_HOME") if not os.environ.get(name)] + if missing: + raise RuntimeError( + f"Please set required cache environment variables: {', '.join(missing)}" + ) + print(f"Using TORCH_HOME={os.environ['TORCH_HOME']}") + print(f"Using HF_HOME={os.environ['HF_HOME']}") + + +def build_model(name: str) -> torch.nn.Module: + weights_cls = getattr(models, f"{name.upper()}_Weights") + model_fn = getattr(models, name) + model = model_fn(weights=weights_cls.DEFAULT) + model.eval() + return model + + +def write_wts(model: torch.nn.Module, output_path: Path) -> None: + state_dict = model.state_dict() + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w") as f: + f.write(f"{len(state_dict.keys())}\n") + for key, value in state_dict.items(): + print(f"key: {key}\tvalue: {value.shape}") + values = value.reshape(-1).cpu().numpy() + f.write(f"{key} {len(values)}") + for item in values: + f.write(" ") + f.write(struct.pack(">f", float(item)).hex()) + f.write("\n") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export torchvision VGG weights to tensorrtx .wts files." + ) + parser.add_argument( + "--model", + choices=MODEL_NAMES, + action="append", + help="VGG variant to export. Repeat the option for multiple variants. Defaults to all variants.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("../models"), + help="Directory for generated .wts files.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + require_cache_env() + model_names = args.model or MODEL_NAMES + for model_name in model_names: + print(f"writing {model_name}.wts") + write_wts(build_model(model_name), args.output_dir / f"{model_name}.wts") diff --git a/vgg/logging.h b/vgg/logging.h index 602b69fb..a303c7de 100644 --- a/vgg/logging.h +++ b/vgg/logging.h @@ -17,7 +17,6 @@ #ifndef TENSORRT_LOGGING_H #define TENSORRT_LOGGING_H -#include "NvInferRuntimeCommon.h" #include #include #include @@ -25,32 +24,25 @@ #include #include #include +#include +#include "NvInferRuntime.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) - { - } +class LogStreamConsumerBuffer : public std::stringbuf { + public: + LogStreamConsumerBuffer(std::ostream& stream, std::string prefix, bool shouldLog) + : mOutput(stream), mPrefix(std::move(prefix)), mShouldLog(shouldLog) {} - LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) - : mOutput(other.mOutput) - { - } + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) noexcept : mOutput(other.mOutput) {} - ~LogStreamConsumerBuffer() - { + ~LogStreamConsumerBuffer() override { // 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()) - { + if (pbase() != pptr()) { putOutput(); } } @@ -58,16 +50,13 @@ class LogStreamConsumerBuffer : public std::stringbuf // 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() - { + auto sync() -> int override { putOutput(); return 0; } - void putOutput() - { - if (mShouldLog) - { + void putOutput() { + if (mShouldLog) { // prepend timestamp std::time_t timestamp = std::time(nullptr); tm* tm_local = std::localtime(×tamp); @@ -88,12 +77,9 @@ class LogStreamConsumerBuffer : public std::stringbuf } } - void setShouldLog(bool shouldLog) - { - mShouldLog = shouldLog; - } + void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; } -private: + private: std::ostream& mOutput; std::string mPrefix; bool mShouldLog; @@ -103,15 +89,12 @@ class LogStreamConsumerBuffer : public std::stringbuf //! \class LogStreamConsumerBase //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer //! -class LogStreamConsumerBase -{ -public: +class LogStreamConsumerBase { + public: LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) - : mBuffer(stream, prefix, shouldLog) - { - } + : mBuffer(stream, prefix, shouldLog) {} -protected: + protected: LogStreamConsumerBuffer mBuffer; }; @@ -124,49 +107,49 @@ class LogStreamConsumerBase //! 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: +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) - { + : 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) noexcept + : 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) - { + private: + static auto severityOstream(Severity severity) -> std::ostream& { 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 ""; + static auto severityPrefix(Severity severity) -> std::string { + 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 ""; } } @@ -198,24 +181,19 @@ class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream //! 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) - { - } +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 + enum class TestResult { + kRUNNING, //!< The test is running + kPASSED, //!< The test passed + kFAILED, //!< The test failed + kWAIVED //!< The test was waived }; //! @@ -225,10 +203,7 @@ class Logger : public nvinfer1::ILogger //! 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; - } + auto getTRTLogger() -> nvinfer1::ILogger& { return *this; } //! //! \brief Implementation of the nvinfer1::ILogger::log() virtual method @@ -236,9 +211,8 @@ class Logger : public nvinfer1::ILogger //! 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; + void log(Severity severity, const char* msg) TRT_NOEXCEPT override { + LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << '\n'; } //! @@ -246,10 +220,7 @@ class Logger : public nvinfer1::ILogger //! //! \param severity The logger will only emit messages that have severity of this level or higher. //! - void setReportableSeverity(Severity severity) - { - mReportableSeverity = severity; - } + void setReportableSeverity(Severity severity) { mReportableSeverity = severity; } //! //! \brief Opaque handle that holds logging information for a particular test @@ -258,20 +229,15 @@ class Logger : public nvinfer1::ILogger //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used //! with Logger::reportTest{Start,End}(). //! - class TestAtom - { - public: + class TestAtom { + public: TestAtom(TestAtom&&) = default; - private: + private: friend class Logger; - TestAtom(bool started, const std::string& name, const std::string& cmdline) - : mStarted(started) - , mName(name) - , mCmdline(cmdline) - { - } + TestAtom(bool started, std::string name, std::string cmdline) + : mStarted(started), mName(std::move(name)), mCmdline(std::move(cmdline)) {} bool mStarted; std::string mName; @@ -289,9 +255,8 @@ class Logger : public nvinfer1::ILogger // //! \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); + static auto defineTest(const std::string& name, const std::string& cmdline) -> TestAtom { + return {false, name, cmdline}; } //! @@ -303,8 +268,7 @@ class Logger : public nvinfer1::ILogger //! \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) - { + static auto defineTest(const std::string& name, int argc, char const* const* argv) -> TestAtom { auto cmdline = genCmdlineString(argc, argv); return defineTest(name, cmdline); } @@ -316,8 +280,7 @@ class Logger : public nvinfer1::ILogger //! //! \param[in] testAtom The handle to the test that has started //! - static void reportTestStart(TestAtom& testAtom) - { + static void reportTestStart(TestAtom& testAtom) { reportTestResult(testAtom, TestResult::kRUNNING); assert(!testAtom.mStarted); testAtom.mStarted = true; @@ -332,98 +295,95 @@ class Logger : public nvinfer1::ILogger //! \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) - { + static void reportTestEnd(const TestAtom& testAtom, TestResult result) { assert(result != TestResult::kRUNNING); assert(testAtom.mStarted); reportTestResult(testAtom, result); } - static int reportPass(const TestAtom& testAtom) - { + static auto reportPass(const TestAtom& testAtom) -> int { reportTestEnd(testAtom, TestResult::kPASSED); return EXIT_SUCCESS; } - static int reportFail(const TestAtom& testAtom) - { + static auto reportFail(const TestAtom& testAtom) -> int { reportTestEnd(testAtom, TestResult::kFAILED); return EXIT_FAILURE; } - static int reportWaive(const TestAtom& testAtom) - { + static auto reportWaive(const TestAtom& testAtom) -> int { reportTestEnd(testAtom, TestResult::kWAIVED); return EXIT_SUCCESS; } - static int reportTest(const TestAtom& testAtom, bool pass) - { + static auto reportTest(const TestAtom& testAtom, bool pass) -> int { return pass ? reportPass(testAtom) : reportFail(testAtom); } - Severity getReportableSeverity() const - { - return mReportableSeverity; - } + [[nodiscard]] auto getReportableSeverity() const -> Severity { return mReportableSeverity; } -private: + 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 ""; + static auto severityPrefix(Severity severity) -> const char* { + 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 ""; + static auto testResultString(TestResult result) -> const char* { + 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) - { + static auto severityOstream(Severity severity) -> std::ostream& { return severity >= Severity::kINFO ? std::cout : std::cerr; } //! //! \brief method that implements logging test results //! - static void reportTestResult(const TestAtom& testAtom, TestResult result) - { + static void reportTestResult(const TestAtom& testAtom, TestResult result) { severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " - << testAtom.mCmdline << std::endl; + << testAtom.mCmdline << '\n'; } //! //! \brief generate a command line string from the given (argc, argv) values //! - static std::string genCmdlineString(int argc, char const* const* argv) - { + static auto genCmdlineString(int argc, char const* const* argv) -> std::string { std::stringstream ss; - for (int i = 0; i < argc; i++) - { + for (int i = 0; i < argc; i++) { if (i > 0) ss << " "; ss << argv[i]; @@ -434,8 +394,7 @@ class Logger : public nvinfer1::ILogger Severity mReportableSeverity; }; -namespace -{ +namespace { //! //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE @@ -444,9 +403,8 @@ namespace //! //! LOG_VERBOSE(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) -{ - return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); +inline auto LOG_VERBOSE(const Logger& logger) -> LogStreamConsumer { + return {logger.getReportableSeverity(), Severity::kVERBOSE}; } //! @@ -456,9 +414,8 @@ inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) //! //! LOG_INFO(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_INFO(const Logger& logger) -{ - return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); +inline auto LOG_INFO(const Logger& logger) -> LogStreamConsumer { + return {logger.getReportableSeverity(), Severity::kINFO}; } //! @@ -468,9 +425,8 @@ inline LogStreamConsumer LOG_INFO(const Logger& logger) //! //! LOG_WARN(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_WARN(const Logger& logger) -{ - return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); +inline auto LOG_WARN(const Logger& logger) -> LogStreamConsumer { + return {logger.getReportableSeverity(), Severity::kWARNING}; } //! @@ -480,9 +436,8 @@ inline LogStreamConsumer LOG_WARN(const Logger& logger) //! //! LOG_ERROR(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_ERROR(const Logger& logger) -{ - return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); +inline auto LOG_ERROR(const Logger& logger) -> LogStreamConsumer { + return {logger.getReportableSeverity(), Severity::kERROR}; } //! @@ -493,11 +448,10 @@ inline LogStreamConsumer LOG_ERROR(const Logger& logger) //! //! LOG_FATAL(logger) << "hello world" << std::endl; //! -inline LogStreamConsumer LOG_FATAL(const Logger& logger) -{ - return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); +inline auto LOG_FATAL(const Logger& logger) -> LogStreamConsumer { + return {logger.getReportableSeverity(), Severity::kINTERNAL_ERROR}; } -} // anonymous namespace +} // anonymous namespace -#endif // TENSORRT_LOGGING_H +#endif // TENSORRT_LOGGING_H diff --git a/vgg/macros.h b/vgg/macros.h new file mode 100644 index 00000000..6d851752 --- /dev/null +++ b/vgg/macros.h @@ -0,0 +1,37 @@ +#pragma once +#include + +#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 + +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) + +#if TRT_VERSION_LT(7, 2, 2) +#error "TensorRT >= 7.2.2 is required for this demo." +#endif + +#if TRT_VERSION_GE(8, 0, 0) +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif diff --git a/vgg/utils.h b/vgg/utils.h new file mode 100644 index 00000000..b9abe1eb --- /dev/null +++ b/vgg/utils.h @@ -0,0 +1,308 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "macros.h" + +using namespace nvinfer1; + +constexpr const std::size_t WORKSPACE_SIZE = 16 << 20; + +#define CHECK(status) \ + do { \ + auto ret = (status); \ + if (ret != cudaSuccess) { \ + std::cerr << "Cuda failure: " << ret << "\n"; \ + std::abort(); \ + } \ + } while (0) + +static void checkTrtEnv(int device = 0) { +#if TRT_VERSION_LT(8, 0, 0) + CHECK(cudaGetDevice(&device)); + cudaDeviceProp prop{}; + CHECK(cudaGetDeviceProperties(&prop, device)); + const int sm = prop.major * 10 + prop.minor; + if (sm > 86) { + std::cerr << "TensorRT < 8 does not support SM > 86 on this GPU."; + std::abort(); + } +#endif +} + +/** + * @brief TensorRT weight files have a simple space delimited format: + * [type] [size] + * + * @param file input weight file path + * @return std::map + */ +static std::map loadWeights(const std::string& file) { + std::cout << "Loading weights: " << file << "\n"; + std::map 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--) { + nvinfer1::Weights wt{nvinfer1::DataType::kFLOAT, nullptr, 0}; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> wt.count; + + // Load blob + auto* val = new float[wt.count]; + input >> std::hex; + for (auto x = 0ll; x < wt.count; ++x) { + uint32_t bits; + input >> bits; + std::memcpy(&val[x], &bits, sizeof(bits)); + } + wt.values = val; + weightMap[name] = wt; + } + + return weightMap; +} + +static void releaseWeights(std::map& weightMap) { + for (auto& weight : weightMap) { + delete[] static_cast(weight.second.values); + weight.second.values = nullptr; + } +} + +/** + * @brief a preprocess function aligning with ImageNet preprocess in torchvision, only support 3-channel image + * + * @param img opencv image with BGR layout + * @param bgr2rgb whether to convert BGR to RGB + * @param mean subtract mean + * @param std divide std + * @param n batch size + * @param h resize height + * @param w resize width + * @return std::vector contiguous flatten image data in float32 type + */ +static std::vector preprocess_img(cv::Mat& img, bool bgr2rgb, const std::array& mean, + const std::array& std, int n, int h, int w) { + const auto c = img.channels(); + const auto size = c * h * w; + if (c != 3) { + std::cerr << "this demo only supports 3 channel input image.\n"; + std::abort(); + } + if (bgr2rgb) { + cv::cvtColor(img, img, cv::COLOR_BGR2RGB); + } + cv::resize(img, img, cv::Size(w, h), 0, 0, cv::INTER_LINEAR); + img.convertTo(img, CV_32FC3, 1.f / 255); + img = (img - cv::Scalar(mean[0], mean[1], mean[2])) / cv::Scalar(std[0], std[1], std[2]); + std::vector chw(static_cast(n) * c * h * w, 0.f); + + // fill all batch with the same input image + for (int i = 0; i < n; ++i) { + for (int y = 0; y < h; ++y) { + for (int x = 0; x < w; ++x) { + const cv::Vec3f v = img.at(y, x); + chw[i * size + 0 * h * w + y * w + x] = v[0]; + chw[i * size + 1 * h * w + y * w + x] = v[1]; + chw[i * size + 2 * h * w + y * w + x] = v[2]; + } + } + } + return chw; +} + +static std::vector> topk(const std::vector& v, int k) { + if (k <= 0) + return {}; + auto stride = std::min(k, static_cast(v.size())); + + std::vector idx(v.size()); + std::iota(idx.begin(), idx.end(), 0); + + std::partial_sort(idx.begin(), idx.begin() + stride, idx.end(), [&](int a, int b) { return v[a] > v[b]; }); + + std::vector> out; + out.reserve(stride); + for (auto i = 0; i < stride; ++i) + out.emplace_back(idx[i], v[idx[i]]); + return out; +} + +static std::map loadImagenetLabelMap(const std::string& path) { + std::map labels; + std::ifstream in(path); + if (!in.is_open()) { + return labels; + } + std::string line; + while (std::getline(in, line)) { + auto colon = line.find(':'); + if (colon == std::string::npos) { + continue; + } + auto first_quote = line.find('\'', colon); + if (first_quote == std::string::npos) { + continue; + } + auto second_quote = line.find('\'', first_quote + 1); + if (second_quote == std::string::npos) { + continue; + } + int idx = std::stoi(line.substr(0, colon)); + labels[idx] = line.substr(first_quote + 1, second_quote - first_quote - 1); + } + return labels; +} + +static ILayer* addTransformLayer(INetworkDefinition* network, ITensor& input, bool bgr2rgb, + const std::array& mean, const std::array& std) { + struct ScaleParams { + std::array shift; + std::array scale; + }; + static std::vector> gScaleParams; + auto params = std::make_unique(); + params->shift = {-mean[0] / std[0], -mean[1] / std[1], -mean[2] / std[2]}; + params->scale = {1.f / (std[0] * 255.f), 1.f / (std[1] * 255.f), 1.f / (std[2] * 255.f)}; + + static const Weights empty{DataType::kFLOAT, nullptr, 0ll}; + const Weights shift{DataType::kFLOAT, params->shift.data(), 3ll}; + const Weights scale{DataType::kFLOAT, params->scale.data(), 3ll}; + + gScaleParams.emplace_back(std::move(params)); + + ITensor* in = &input; + if (input.getType() != DataType::kFLOAT) { +#if TRT_VERSION_GE(8, 0, 0) + auto* cast = network->addCast(input, DataType::kFLOAT); + assert(cast); + cast->setName("Cast to FP32"); + in = cast->getOutput(0); +#else + auto* identity = network->addIdentity(input); + assert(identity); + identity->setName("Convert to FP32"); + identity->setOutputType(0, DataType::kFLOAT); + in = identity->getOutput(0); +#endif + } + // Convert from NHWC to NCHW + auto* perm = network->addShuffle(*in); + assert(perm); + perm->setName("NHWC -> NCHW"); + perm->setFirstTranspose(Permutation{0, 3, 1, 2}); + + // Convert from BGR to RGB (optional) + ITensor* data{nullptr}; + if (bgr2rgb) { + auto add_slice = [&](int c, const char* name) -> ITensor* { + auto dims = perm->getOutput(0)->getDimensions(); + Dims4 start = {0, c, 0, 0}, stride = {1, 1, 1, 1}; + Dims4 size = {dims.d[0], 1, dims.d[2], dims.d[3]}; + auto* _slice = network->addSlice(*perm->getOutput(0), start, size, stride); + _slice->setName(name); + assert(_slice && _slice->getNbOutputs() == 1); + return _slice->getOutput(0); + }; + std::array channels = {add_slice(2, "R"), add_slice(1, "G"), add_slice(0, "B")}; + auto* cat = network->addConcatenation(channels.data(), 3); + assert(cat); + cat->setName("RGB"); + cat->setAxis(1); + data = cat->getOutput(0); + } else { + data = perm->getOutput(0); + } + + // Normalize + auto* trans = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, empty); + assert(trans); + trans->setName("mean & std"); +#if TRT_VERSION_GE(8, 0, 0) + trans->setChannelAxis(1); +#endif + return trans; +} + +static size_t getSize(DataType dt) { + switch (dt) { +#if TRT_VERSION_GE(8, 5, 1) + case DataType::kUINT8: +#endif + case DataType::kINT8: + return sizeof(int8_t); + case DataType::kFLOAT: + return sizeof(float); + case DataType::kHALF: + return sizeof(int16_t); + case DataType::kINT32: + return sizeof(int32_t); + default: { + std::cerr << "Unsupported data type\n"; + std::abort(); + } + } +} + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} diff --git a/vgg/vgg.cc b/vgg/vgg.cc new file mode 100644 index 00000000..1256275c --- /dev/null +++ b/vgg/vgg.cc @@ -0,0 +1,442 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include "logging.h" +#include "macros.h" +#include "utils.h" + +using WeightMap = std::map; +using M = nvinfer1::MatrixOperation; +using E = nvinfer1::ElementWiseOperation; +using LayerConfig = std::vector; +using LayerConfigMap = std::map; + +static Logger g_logger; + +struct VggVariant { + std::string name; + std::string cfg_name; + bool batch_norm; + std::string wts_path; + std::string engine_path; +}; + +static constexpr int N = 1; +static constexpr const int32_t INPUT_H = 224; +static constexpr const int32_t INPUT_W = 224; +static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; +static constexpr const std::array NAMES = {"data", "prob"}; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; +static constexpr const bool TRT_PREPROCESS = TRT_VERSION_GE(8, 5, 1); +static constexpr const std::array mean = {0.485f, 0.456f, 0.406f}; +static constexpr const std::array stdv = {0.229f, 0.224f, 0.225f}; + +const LayerConfigMap CFGS = { + {"A", {"64", "M", "128", "M", "256", "256", "M", "512", "512", "M", "512", "512", "M"}}, + {"B", {"64", "64", "M", "128", "128", "M", "256", "256", "M", "512", "512", "M", "512", "512", "M"}}, + {"D", + {"64", "64", "M", "128", "128", "M", "256", "256", "256", "M", "512", "512", "512", "M", "512", "512", "512", + "M"}}, + {"E", {"64", "64", "M", "128", "128", "M", "256", "256", "256", "256", "M", + "512", "512", "512", "512", "M", "512", "512", "512", "512", "M"}}}; + +static auto normalizeModelName(const std::string& raw_name) -> std::string { + std::string name; + name.reserve(raw_name.size()); + for (char ch : raw_name) { + const auto c = static_cast(ch); + name.push_back(ch == '-' ? '_' : static_cast(std::tolower(c))); + } + return name; +} + +static auto getVariantConfig(const std::string& raw_name) -> VggVariant { + const auto name = normalizeModelName(raw_name); + static const std::map> variants = { + {"vgg11", {"A", false}}, {"vgg11_bn", {"A", true}}, {"vgg13", {"B", false}}, {"vgg13_bn", {"B", true}}, + {"vgg16", {"D", false}}, {"vgg16_bn", {"D", true}}, {"vgg19", {"E", false}}, {"vgg19_bn", {"E", true}}, + }; + const auto iter = variants.find(name); + if (iter == variants.end()) { + std::cerr << "Unknown VGG variant: " << raw_name + << " (expected vgg11 | vgg11_bn | vgg13 | vgg13_bn | vgg16 | vgg16_bn | vgg19 | vgg19_bn)\n"; + std::abort(); + } + + return VggVariant{name, iter->second.first, iter->second.second, "models/" + name + ".wts", + "models/" + name + ".engine"}; +} + +static auto addBatchNorm2d(INetworkDefinition* network, WeightMap& w, ITensor& input, const std::string& lname, + float eps = 1e-5) -> IScaleLayer* { + const float* gamma = static_cast(w[lname + ".weight"].values); + const float* beta = static_cast(w[lname + ".bias"].values); + const float* mean = static_cast(w[lname + ".running_mean"].values); + const float* var = static_cast(w[lname + ".running_var"].values); + int64_t len = w[lname + ".running_var"].count; + + auto* scval = new float[len]; + for (int i = 0; i < len; i++) { + scval[i] = gamma[i] / sqrt(var[i] + eps); + } + Weights scale{DataType::kFLOAT, scval, len}; + + auto* shval = new 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}; + + auto* pval = new float[len]; + for (int i = 0; i < len; i++) { + pval[i] = 1.0f; + } + Weights power{DataType::kFLOAT, pval, len}; + + w[lname + ".scale"] = scale; + w[lname + ".shift"] = shift; + w[lname + ".power"] = power; + IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_1); + scale_1->setName(lname.c_str()); + return scale_1; +} + +static auto make_layers(INetworkDefinition* net, WeightMap& w, ITensor& input, const LayerConfig& cfg, + bool use_bn) -> ITensor* { + auto* tensor = &input; + int32_t idx = 0; + std::string name = "features."; + + for (const auto& _v : cfg) { + if (_v == "M") { + auto* pool = net->addPoolingNd(*tensor, PoolingType::kMAX, DimsHW{2, 2}); + pool->setStrideNd(DimsHW{2, 2}); + pool->setName((name + std::to_string(idx)).c_str()); + tensor = pool->getOutput(0); + ++idx; + continue; + } + const int32_t v = std::stoi(_v); + auto _name = "features." + std::to_string(idx); + auto* _conv = net->addConvolutionNd(*tensor, v, DimsHW{3, 3}, w.at(_name + ".weight"), w.at(_name + ".bias")); + _conv->setPaddingNd(DimsHW{1, 1}); + _conv->setName(_name.c_str()); + tensor = _conv->getOutput(0); + ++idx; + if (use_bn) { + auto _bn_name = "features." + std::to_string(idx); + auto* bn = addBatchNorm2d(net, w, *tensor, _bn_name); + tensor = bn->getOutput(0); + ++idx; + } + auto* relu = net->addActivation(*tensor, ActivationType::kRELU); + relu->setName(("features." + std::to_string(idx)).c_str()); + tensor = relu->getOutput(0); + ++idx; + } + + return tensor; +} + +auto create_engine(const VggVariant& variant, int32_t batch_size, IRuntime* runtime, IBuilder* builder, + IBuilderConfig* config, DataType dt) -> ICudaEngine* { + auto w = loadWeights(variant.wts_path); +#if TRT_VERSION_GE(10, 0, 0) + auto* net = builder->createNetworkV2(1U << static_cast(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)); +#else + auto* net = builder->createNetworkV2(1u << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)); +#endif + + ITensor* input{nullptr}; + if constexpr (TRT_PREPROCESS) { + // for simplicity, resize image on cpu side + dt = DataType::kUINT8; + input = net->addInput(NAMES[0], dt, Dims4{batch_size, INPUT_H, INPUT_W, 3}); + auto* trans = addTransformLayer(net, *input, true, mean, stdv); + input = trans->getOutput(0); + } else { + input = net->addInput(NAMES[0], dt, Dims4{batch_size, 3, INPUT_H, INPUT_W}); + } + assert(input); + + auto* features = make_layers(net, w, *input, CFGS.at(variant.cfg_name), variant.batch_norm); + + auto* _avg_pool = net->addPoolingNd(*features, PoolingType::kAVERAGE, Dims2{1, 1}); + auto* _flatten = net->addShuffle(*_avg_pool->getOutput(0)); + assert(_avg_pool && _flatten); + _flatten->setReshapeDimensions(Dims2{batch_size, -1}); + + auto* _fc1w = + net->addConstant(Dims2{4096, static_cast(512 * 7 * 7)}, w["classifier.0.weight"])->getOutput(0); + auto* _fc1b = net->addConstant(Dims2{1, 4096}, w["classifier.0.bias"])->getOutput(0); + auto* _fc2w = net->addConstant(Dims2{4096, 4096}, w["classifier.3.weight"])->getOutput(0); + auto* _fc2b = net->addConstant(Dims2{1, 4096}, w["classifier.3.bias"])->getOutput(0); + auto* _fc3w = net->addConstant(Dims2{1000, 4096}, w["classifier.6.weight"])->getOutput(0); + auto* _fc3b = net->addConstant(Dims2{1, 1000}, w["classifier.6.bias"])->getOutput(0); + assert(_fc1w && _fc1b && _fc2w && _fc2b && _fc3w && _fc3b); + + auto* _fc1_0 = net->addMatrixMultiply(*_flatten->getOutput(0), M::kNONE, *_fc1w, M::kTRANSPOSE); + auto* _fc1_1 = net->addElementWise(*_fc1_0->getOutput(0), *_fc1b, E::kSUM); + auto* _relu1 = net->addActivation(*_fc1_1->getOutput(0), ActivationType::kRELU); + + auto* _fc2_0 = net->addMatrixMultiply(*_relu1->getOutput(0), M::kNONE, *_fc2w, M::kTRANSPOSE); + auto* _fc2_1 = net->addElementWise(*_fc2_0->getOutput(0), *_fc2b, E::kSUM); + auto* _relu2 = net->addActivation(*_fc2_1->getOutput(0), ActivationType::kRELU); + + auto* _fc3_0 = net->addMatrixMultiply(*_relu2->getOutput(0), M::kNONE, *_fc3w, M::kTRANSPOSE); + auto* _fc3_1 = net->addElementWise(*_fc3_0->getOutput(0), *_fc3b, E::kSUM); + + _fc3_1->getOutput(0)->setName(NAMES[1]); + net->markOutput(*_fc3_1->getOutput(0)); + +#if TRT_VERSION_GE(8, 0, 0) + config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); + auto* _serialized = builder->buildSerializedNetwork(*net, *config); + auto* _engine = runtime->deserializeCudaEngine(_serialized->data(), _serialized->size()); + delete _serialized; + delete net; +#else + builder->setMaxBatchSize(N); + config->setMaxWorkspaceSize(WORKSPACE_SIZE); + auto* _engine = builder->buildEngineWithConfig(*net, *config); + net->destroy(); +#endif + std::cout << "build out" << '\n'; + + releaseWeights(w); + + return _engine; +} + +void APIToModel(const VggVariant& variant, int32_t batch_size, IRuntime* runtime, IHostMemory** model_stream) { + auto* builder = createInferBuilder(g_logger); + auto* config = builder->createBuilderConfig(); + + auto* engine = create_engine(variant, batch_size, runtime, builder, config, DataType::kFLOAT); + assert(engine != nullptr); + + (*model_stream) = engine->serialize(); + +#if TRT_VERSION_GE(8, 0, 0) + delete engine; + delete config; + delete builder; +#else + engine->destroy(); + config->destroy(); + builder->destroy(); +#endif +} + +auto doInference(IExecutionContext& context, void* input, std::size_t batch_size) -> std::vector> { + const ICudaEngine& engine = context.getEngine(); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + std::vector buffers; + +#if TRT_VERSION_GE(8, 0, 0) + const int32_t nIO = engine.getNbIOTensors(); +#else + const int32_t nIO = engine.getNbBindings(); +#endif + + buffers.resize(nIO); + for (auto i = 0; i < nIO; ++i) { + std::size_t size = 0; +#if TRT_VERSION_GE(8, 0, 0) + auto* tensor_name = engine.getIOTensorName(i); + const std::string name = tensor_name; + auto s = getSize(engine.getTensorDataType(tensor_name)); + size = s * batch_size * (name == NAMES[0] ? SIZES[0] : SIZES[1]); + CHECK(cudaMalloc(&buffers[i], size)); + if (name == NAMES[0]) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } +#else + const int32_t idx = engine.getBindingIndex(NAMES[i]); + auto s = getSize(engine.getBindingDataType(idx)); + assert(idx == i); + size = s * batch_size * SIZES[i]; + CHECK(cudaMalloc(&buffers[i], size)); + if (i == 0) { + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); + } +#endif + } + +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } +#else + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } +#endif + + std::vector> prob; + for (int i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) + const std::string name = engine.getIOTensorName(i); + if (name == NAMES[0]) { + continue; + } + constexpr auto output_size = SIZES[1]; +#else + if (i == 0) { + continue; + } + const auto output_size = SIZES[i]; +#endif + std::vector tmp(batch_size * output_size, std::nanf("")); + std::size_t size = batch_size * output_size * sizeof(float); + CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); + prob.emplace_back(tmp); + } + CHECK(cudaStreamSynchronize(stream)); + + for (auto i = 0; i < nIO; ++i) { + CHECK(cudaFree(buffers[i])); + } + CHECK(cudaStreamDestroy(stream)); + return prob; +} + +auto main(int argc, char** argv) -> int { + checkTrtEnv(); + if (argc < 2 || argc > 3) { + std::cerr << "arguments not right!" << '\n'; + std::cerr << "./vgg -s [model] // serialize model to plan file" << '\n'; + std::cerr << "./vgg -d [model] // deserialize plan file and run inference" << '\n'; + std::cerr << "model choices: vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn" << '\n'; + return -1; + } + const auto variant = getVariantConfig(argc == 3 ? argv[2] : "vgg11"); + std::cout << "Using VGG variant: " << variant.name << '\n'; + + auto* runtime = createInferRuntime(g_logger); + assert(runtime != nullptr); + char* trt_model_stream{nullptr}; + std::streamsize size{0}; + + if (std::string(argv[1]) == "-s") { + IHostMemory* mem{nullptr}; + APIToModel(variant, 1, runtime, &mem); + assert(mem != nullptr); + + std::ofstream _plan(variant.engine_path, std::ios::binary | std::ios::trunc); + if (!_plan) { + std::cerr << "could not open plan output file" << '\n'; + return -1; + } + if (mem->size() > static_cast(std::numeric_limits::max())) { + std::cerr << "this model is too large to serialize\n"; + return -1; + } + const auto* data_ptr = reinterpret_cast(mem->data()); + auto data_size = static_cast(mem->size()); + _plan.write(data_ptr, data_size); +#if TRT_VERSION_GE(8, 0, 0) + delete mem; + delete runtime; +#else + mem->destroy(); + runtime->destroy(); +#endif + return 0; + } else if (std::string(argv[1]) == "-d") { + std::ifstream _file(variant.engine_path, std::ios::binary); + if (_file.good()) { + _file.seekg(0, _file.end); + size = _file.tellg(); + _file.seekg(0, _file.beg); + trt_model_stream = new char[size]; + assert(trt_model_stream); + _file.read(trt_model_stream, size); + _file.close(); + } else { + std::cerr << "could not open engine file" << '\n'; + return -1; + } + } else { + return 1; + } + +#if TRT_VERSION_GE(8, 0, 0) + auto* engine = runtime->deserializeCudaEngine(trt_model_stream, size); +#else + auto* engine = runtime->deserializeCudaEngine(trt_model_stream, size, nullptr); +#endif + assert(engine != nullptr); + auto* context = engine->createExecutionContext(); + assert(context != nullptr); + + const std::string img_path = "assets/cats.jpg"; + void* input = nullptr; + std::vector flat_img; + cv::Mat img; + if constexpr (TRT_PREPROCESS) { + // for simplicity, resize image on cpu side + img = cv::imread(img_path, cv::IMREAD_COLOR); + cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR); + input = static_cast(img.data); + } else { + img = cv::imread(img_path, cv::IMREAD_COLOR); + flat_img = preprocess_img(img, true, mean, stdv, N, INPUT_H, INPUT_W); + input = flat_img.data(); + } + assert(input); + + auto firstProb = doInference(*context, input, 1); + printFirstOutputs("vgg", firstProb[0].data(), firstProb[0].size()); + std::cout << "prediction result:\n"; + auto labels = loadImagenetLabelMap(LABELS_PATH); + int _top = 0; + for (auto& [idx, logits] : topk(firstProb[0], 3)) { + std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits << ", label: " << labels[idx] + << '\n'; + } + + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int32_t i = 0; i < kBenchmarkRuns; ++i) { + auto _start = std::chrono::steady_clock::now(); + (void)doInference(*context, input, 1); + auto _end = std::chrono::steady_clock::now(); + auto _time = std::chrono::duration_cast(_end - _start).count(); + latencies.push_back(static_cast(_time) / 1000.0); + } + printBenchmark("vgg", latencies); + + delete[] trt_model_stream; +#if TRT_VERSION_GE(8, 0, 0) + delete context; + delete engine; + delete runtime; +#else + context->destroy(); + engine->destroy(); + runtime->destroy(); +#endif + return 0; +} diff --git a/vgg/vgg11.cpp b/vgg/vgg11.cpp deleted file mode 100644 index 45ac7514..00000000 --- a/vgg/vgg11.cpp +++ /dev/null @@ -1,289 +0,0 @@ -#include "NvInfer.h" -#include "cuda_runtime_api.h" -#include -#include -#include -#include -#include -#include -#include "logging.h" - -#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] -std::map loadWeights(const std::string file) -{ - std::cout << "Loading weights: " << file << std::endl; - std::map 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(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 { 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 weightMap = loadWeights("../vgg.wts"); - Weights emptywts{DataType::kFLOAT, nullptr, 0}; - - IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{3, 3}, weightMap["features.0.weight"], weightMap["features.0.bias"]); - assert(conv1); - conv1->setPaddingNd(DimsHW{1, 1}); - IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - assert(relu1); - IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); - assert(pool1); - pool1->setStrideNd(DimsHW{2, 2}); - - conv1 = network->addConvolutionNd(*pool1->getOutput(0), 128, DimsHW{3, 3}, weightMap["features.3.weight"], weightMap["features.3.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); - pool1->setStrideNd(DimsHW{2, 2}); - - conv1 = network->addConvolutionNd(*pool1->getOutput(0), 256, DimsHW{3, 3}, weightMap["features.6.weight"], weightMap["features.6.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - conv1 = network->addConvolutionNd(*relu1->getOutput(0), 256, DimsHW{3, 3}, weightMap["features.8.weight"], weightMap["features.8.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); - pool1->setStrideNd(DimsHW{2, 2}); - - conv1 = network->addConvolutionNd(*pool1->getOutput(0), 512, DimsHW{3, 3}, weightMap["features.11.weight"], weightMap["features.11.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - conv1 = network->addConvolutionNd(*relu1->getOutput(0), 512, DimsHW{3, 3}, weightMap["features.13.weight"], weightMap["features.13.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); - pool1->setStrideNd(DimsHW{2, 2}); - - conv1 = network->addConvolutionNd(*pool1->getOutput(0), 512, DimsHW{3, 3}, weightMap["features.16.weight"], weightMap["features.16.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - conv1 = network->addConvolutionNd(*relu1->getOutput(0), 512, DimsHW{3, 3}, weightMap["features.18.weight"], weightMap["features.18.bias"]); - conv1->setPaddingNd(DimsHW{1, 1}); - relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); - pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); - pool1->setStrideNd(DimsHW{2, 2}); - - IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool1->getOutput(0), 4096, weightMap["classifier.0.weight"], weightMap["classifier.0.bias"]); - assert(fc1); - relu1 = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU); - fc1 = network->addFullyConnected(*relu1->getOutput(0), 4096, weightMap["classifier.3.weight"], weightMap["classifier.3.bias"]); - relu1 = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU); - fc1 = network->addFullyConnected(*relu1->getOutput(0), 1000, weightMap["classifier.6.weight"], weightMap["classifier.6.bias"]); - - fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME); - std::cout << "set name out" << std::endl; - network->markOutput(*fc1->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(); - config->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 << "./vgg -s // serialize model to plan file" << std::endl; - std::cerr << "./vgg -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("vgg.engine", std::ios::binary); - if (!p) { - std::cerr << "could not open plan output file" << std::endl; - return -1; - } - p.write(reinterpret_cast(modelStream->data()), modelStream->size()); - modelStream->destroy(); - return 1; - } else if (std::string(argv[1]) == "-d") { - std::ifstream file("vgg.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; - } - - static 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); - assert(engine != nullptr); - IExecutionContext* context = engine->createExecutionContext(); - assert(context != nullptr); - delete[] trtModelStream; - - // Run inference - static float prob[OUTPUT_SIZE]; - for (int i = 0; i < 10; 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(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; -} diff --git a/vit/CMakeLists.txt b/vit/CMakeLists.txt index 2a3676c8..b0ed0da3 100644 --- a/vit/CMakeLists.txt +++ b/vit/CMakeLists.txt @@ -11,7 +11,7 @@ endif() set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD 20) set(CMAKE_CUDA_STANDARD_REQUIRED ON) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) set(CMAKE_INCLUDE_CURRENT_DIR TRUE) diff --git a/vit/FindTensorRT.cmake b/vit/FindTensorRT.cmake index f529daa2..869ad16e 100644 --- a/vit/FindTensorRT.cmake +++ b/vit/FindTensorRT.cmake @@ -70,8 +70,11 @@ if(WIN32) endif() if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10) - set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10 - nvinfer_dispatch_10 nvinfer_lean_10) + set(_trt_lib_suffix "_${TRT_MAJOR_VERSION}") + set(_modules nvinfer${_trt_lib_suffix} nvinfer_plugin${_trt_lib_suffix} + nvinfer_vc_plugin${_trt_lib_suffix} + nvinfer_dispatch${_trt_lib_suffix} + nvinfer_lean${_trt_lib_suffix}) message(DEBUG "Using ${_modules}") else() set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch @@ -137,6 +140,7 @@ set_target_properties( unset(TRT_MAJOR_VERSION) unset(_modules) +unset(_trt_lib_suffix) unset(_trt_include_candidates) unset(_trt_library_candidates) unset(_trt_arch) diff --git a/vit/README.md b/vit/README.md index d19217df..10e639ae 100644 --- a/vit/README.md +++ b/vit/README.md @@ -19,21 +19,33 @@ This is a handwritten TensorRT implementation of the Vision Transformers[arxiv.o - Support a dummy profiler by default - Support a dummy output allocator by default - Use optimization profile by default +- Support ViT-B/16, ViT-B/32, ViT-L/16, ViT-L/32 and ViT-H/14 +- Use an optimization profile by default (`min=1 / opt=1 / max=2`) ### 2.2 Current limitations - cannot use `IAttenion` with TensorRT SDK 10.14 ~ 10.15 because of the bugs in TensorRT - TensorRT < 8 is not supported because some ops are not inplemented in cuDNN - SM < 86, TensorRT < 10, CUDA < 12 cases are _NOT_ fully tested yet +- ViT-H/14 uses the HuggingFace ImageNet-21k checkpoint + (`google/vit-huge-patch14-224-in21k`, 21843 classes); no public 1k + fine-tuned checkpoint is available. ### 2.3 Usage 1. use `gen_wts.py` to generate `.wts` file. ```bash -python gen_wts.py +# Choose one of: ViT-B/16, ViT-B/32, ViT-L/16, ViT-L/32, ViT-H/14 +# Aliases b16, b32, l16, l32 and h14 are also accepted. +python gen_wts.py ViT-B/16 +# -> writes models/ViT-B-16.wts ``` +`gen_wts.py` sets model download caches before importing `torch` and +`transformers`: `TORCH_HOME=/mnt/data/storage/torch` and +`HF_HOME=/mnt/data/storage/huggingface`. + 2. build C++ code ```bash @@ -45,15 +57,22 @@ cmake --build build 3. serialize `.wts` model to engine file. ```bash -./build/vit -s +./build/vit -s +# Example: +./build/vit -s models/ViT-B-16.wts models/ViT-B-16.engine ViT-B/16 ``` 4. run inference ```bash -./build/vit -d +./build/vit -d +# Example: +./build/vit -d models/ViT-B-16.engine assets/cats.jpg ``` +The engine uses a dynamic batch profile (`min=1 / opt=1 / max=2`). Passing a +directory with two images runs both samples in one batch. + On **RTX 4080, TensorRT 10.15.1 SDK**, the output looks like: ```bash @@ -110,7 +129,7 @@ Where: - (N): batch size (represented by `N` in your code) - (L): sequence length (number of tokens; dynamic in code via `-1`) -- (D): hidden size, fixed at 768 in this implementation +- (D): hidden size, configured by the selected variant The attention head configuration: @@ -141,7 +160,7 @@ For a standard Transformer block: $$ \mathbf{W}_2 \in \mathbb{R}^{4D \times D}, \ \mathbf{b}_2 \in \mathbb{R}^{D} $$ - Here ($4 D = 3072$). + Here the FFN dimension is configured by the selected variant. ### 3.3 High-Level Block Structure diff --git a/vit/cuda_allocator.cc b/vit/cuda_allocator.cc index 2d46469a..3f06aa08 100644 --- a/vit/cuda_allocator.cc +++ b/vit/cuda_allocator.cc @@ -16,6 +16,7 @@ struct CudaOutputAllocator::Allocation { void* ptr{nullptr}; std::size_t size{0}; OutputAllocKind kind{OutputAllocKind::kCudaMallocManaged}; + cudaStream_t stream{}; CUmemGenericAllocationHandle handle{}; CUdeviceptr addr{}; std::size_t mapped_size{0}; @@ -61,7 +62,7 @@ CudaOutputAllocator::~CudaOutputAllocator() { } } -#if TRT_VERSION < 10000 +#if TRT_VERSION_LT(10, 0, 0) // NOLINTNEXTLINE(bugprone-easily-swappable-parameters) void* CudaOutputAllocator::reallocateOutput(const char* tensorName, void* currentMemory, uint64_t size, uint64_t alignment) TRT_NOEXCEPT { @@ -80,7 +81,7 @@ void* CudaOutputAllocator::reallocateOutput(const char* tensorName, void* curren return currentMemory; } - Allocation fresh = allocate(static_cast(size)); + Allocation fresh = allocate(static_cast(size), stream_); if (!fresh.ptr) { return nullptr; } @@ -98,10 +99,12 @@ void* CudaOutputAllocator::reallocateOutputAsync(const char* tensorName, void* c if (stream == nullptr) { stream = stream_; } - stream_ = stream; std::lock_guard lock(mutex_); auto& alloc = allocations_[tensorName]; if (alloc.ptr && size <= alloc.size) { + if (alloc.kind == OutputAllocKind::kCudaMallocAsync) { + alloc.stream = stream; + } return alloc.ptr; } if (alloc.ptr) { @@ -110,7 +113,7 @@ void* CudaOutputAllocator::reallocateOutputAsync(const char* tensorName, void* c return currentMemory; } - Allocation fresh = allocate(static_cast(size)); + Allocation fresh = allocate(static_cast(size), stream); if (!fresh.ptr) { return nullptr; } @@ -121,19 +124,26 @@ void* CudaOutputAllocator::reallocateOutputAsync(const char* tensorName, void* c void CudaOutputAllocator::notifyShape(const char* /*tensorName*/, nvinfer1::Dims const& /*dims*/) TRT_NOEXCEPT {} -CudaOutputAllocator::Allocation CudaOutputAllocator::allocate(std::size_t size) { +CudaOutputAllocator::Allocation CudaOutputAllocator::allocate(std::size_t size, cudaStream_t stream) { Allocation alloc{}; if (size == 0) { return alloc; } + if (stream == nullptr) { + stream = stream_; + } + if (cudaSetDevice(device_) != cudaSuccess) { + return alloc; + } if (kind_ == OutputAllocKind::kCudaMallocAsync) { void* ptr = nullptr; - if (cudaMallocAsync(&ptr, size, stream_) != cudaSuccess) { + if (cudaMallocAsync(&ptr, size, stream) != cudaSuccess) { return alloc; } alloc.ptr = ptr; alloc.size = size; alloc.kind = OutputAllocKind::kCudaMallocAsync; + alloc.stream = stream; return alloc; } if (kind_ == OutputAllocKind::kCudaMallocManaged) { @@ -147,9 +157,6 @@ CudaOutputAllocator::Allocation CudaOutputAllocator::allocate(std::size_t size) return alloc; } - if (cudaSetDevice(device_) != cudaSuccess) { - return alloc; - } if (cuInit(0) != CUDA_SUCCESS) { return alloc; } @@ -205,8 +212,9 @@ void CudaOutputAllocator::release(const std::string& /*tensorName*/, Allocation& if (!alloc.ptr) { return; } + CHECK(cudaSetDevice(device_)); if (alloc.kind == OutputAllocKind::kCudaMallocAsync) { - cudaFreeAsync(alloc.ptr, stream_); + cudaFreeAsync(alloc.ptr, alloc.stream); } else if (alloc.kind == OutputAllocKind::kCudaMallocManaged) { cudaFree(alloc.ptr); } else if (alloc.kind == OutputAllocKind::kCuMem) { diff --git a/vit/cuda_allocator.h b/vit/cuda_allocator.h index dcd6c09e..c5e4bc2e 100644 --- a/vit/cuda_allocator.h +++ b/vit/cuda_allocator.h @@ -17,7 +17,7 @@ class CudaOutputAllocator final : public nvinfer1::IOutputAllocator { explicit CudaOutputAllocator(cudaStream_t stream, OutputAllocKind kind, int device = 0); ~CudaOutputAllocator() override; -#if TRT_VERSION < 10000 +#if TRT_VERSION_LT(10, 0, 0) void* reallocateOutput(const char* tensorName, void* currentMemory, uint64_t size, uint64_t alignment) TRT_NOEXCEPT override; #else @@ -32,7 +32,7 @@ class CudaOutputAllocator final : public nvinfer1::IOutputAllocator { private: struct Allocation; - Allocation allocate(std::size_t size); + Allocation allocate(std::size_t size, cudaStream_t stream); void release(const std::string& tensorName, Allocation& alloc); cudaStream_t stream_{}; diff --git a/vit/gen_wts.py b/vit/gen_wts.py index 75086eb6..838c5b43 100644 --- a/vit/gen_wts.py +++ b/vit/gen_wts.py @@ -1,64 +1,159 @@ +"""Export ViT weights from HuggingFace to a .wts file for TensorRT. + +Usage: + python gen_wts.py [] + +Examples: + python gen_wts.py ViT-B/16 + python gen_wts.py ViT-L/16 models/vit_l16.wts + python gen_wts.py b32 # alias accepted + +Model types supported (must match the table in vit.cc::getVariantConfig): + ViT-B/16 -> google/vit-base-patch16-224 (img=224, classes=1000) + ViT-B/32 -> google/vit-base-patch32-384 (img=384, classes=1000) + ViT-L/16 -> google/vit-large-patch16-224 (img=224, classes=1000) + ViT-L/32 -> google/vit-large-patch32-384 (img=384, classes=1000) + ViT-H/14 -> google/vit-huge-patch14-224-in21k (img=224, classes=21843, + ImageNet-21k pretrain only; + no public 1k fine-tuned ckpt) + +Default output path (when omitted) is models/.wts where slashes +are replaced by hyphens, e.g. ViT-B/16 -> models/ViT-B-16.wts. This matches +the recommended layout used by exp/run_experiment.py. +""" + +from __future__ import annotations + +import argparse +import os import struct +from pathlib import Path -import cv2 -import numpy as np import torch -from transformers import AutoConfig, AutoImageProcessor, AutoModelForImageClassification +from transformers import AutoConfig, AutoModelForImageClassification +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +MODELS_DIR = REPO_ROOT / "models" -def read_imagenet_labels() -> dict[int, str]: - """ - read ImageNet 1000 labels +# Canonical name -> (HF hub id, image size, num classes) +VARIANTS: dict[str, tuple[str, int, int]] = { + "ViT-B/16": ("google/vit-base-patch16-224", 224, 1000), + "ViT-B/32": ("google/vit-base-patch32-384", 384, 1000), + "ViT-L/16": ("google/vit-large-patch16-224", 224, 1000), + "ViT-L/32": ("google/vit-large-patch32-384", 384, 1000), + "ViT-H/14": ("google/vit-huge-patch14-224-in21k", 224, 21843), +} - Returns: - dict[int, str]: labels dict - """ - clsid2label = {} - with open("../assets/imagenet1000_clsidx_to_labels.txt", "r") as f: - for i in f.readlines(): - k, v = i.split(": ") - clsid2label.setdefault(int(k), v[1:-3]) - return clsid2label +ALIASES: dict[str, str] = { + "b16": "ViT-B/16", + "B16": "ViT-B/16", + "B/16": "ViT-B/16", + "b32": "ViT-B/32", + "B32": "ViT-B/32", + "B/32": "ViT-B/32", + "l16": "ViT-L/16", + "L16": "ViT-L/16", + "L/16": "ViT-L/16", + "l32": "ViT-L/32", + "L32": "ViT-L/32", + "L/32": "ViT-L/32", + "h14": "ViT-H/14", + "H14": "ViT-H/14", + "H/14": "ViT-H/14", +} -USE_HF_PREPROCESS = False +def normalize_model_type(name: str) -> str: + if name in VARIANTS: + return name + if name in ALIASES: + return ALIASES[name] + upper = name.upper().replace(" ", "") + if upper.startswith("VIT-"): + cand = "ViT-" + upper[4:] + if cand in VARIANTS: + return cand + raise SystemExit( + f"Unknown model_type: {name!r}. Choose from: {', '.join(VARIANTS.keys())}" + ) -if __name__ == "__main__": - hub_model_id = "google/vit-base-patch16-224" - config = AutoConfig.from_pretrained(hub_model_id) + +def safe_filename(model_type: str) -> str: + """ViT-B/16 -> ViT-B-16 (filesystem-safe).""" + return model_type.replace("/", "-") + + +def require_cache_env() -> None: + missing = [name for name in ("TORCH_HOME", "HF_HOME") if not os.environ.get(name)] + if missing: + raise RuntimeError( + f"Please set required cache environment variables: {', '.join(missing)}" + ) + print(f"Using TORCH_HOME={os.environ['TORCH_HOME']}") + print(f"Using HF_HOME={os.environ['HF_HOME']}") + + +def export_wts(model: torch.nn.Module, out_path: Path) -> None: + out_path.parent.mkdir(parents=True, exist_ok=True) + sd = model.state_dict() + with open(out_path, "w") as f: + f.write(f"{len(sd)}\n") + for k, v in sd.items(): + vr = v.detach().reshape(-1).cpu().numpy() + f.write(f"{k} {vr.size}") + for vv in vr: + f.write(" ") + f.write(struct.pack(">f", float(vv)).hex()) + f.write("\n") + size_mib = out_path.stat().st_size / 1024 / 1024 + print(f"[ok] wrote {out_path} ({size_mib:.1f} MiB, {len(sd)} tensors)") + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawTextHelpFormatter, + ) + parser.add_argument("model_type", help="e.g. ViT-B/16, ViT-L/32, b16, h14") + parser.add_argument( + "output", + nargs="?", + default=None, + help="output .wts path (default: models/.wts)", + ) + args = parser.parse_args() + + require_cache_env() + + model_type = normalize_model_type(args.model_type) + hub_id, img_size, num_classes = VARIANTS[model_type] + out_path = ( + Path(args.output) + if args.output + else MODELS_DIR / f"{safe_filename(model_type)}.wts" + ) + + print( + f"[load] model_type={model_type} hub_id={hub_id} img={img_size} classes={num_classes}" + ) + config = AutoConfig.from_pretrained(hub_id) config._attn_implementation = "eager" + # Force the classifier head size: in21k checkpoints (e.g. ViT-H/14) ship + # without a classifier and HF would otherwise default num_labels=2, + # producing a tiny random head incompatible with our engine table. + config.num_labels = num_classes + config.id2label = {i: str(i) for i in range(num_classes)} + config.label2id = {str(i): i for i in range(num_classes)} model = AutoModelForImageClassification.from_pretrained( - hub_model_id, - ignore_mismatched_sizes=False, + hub_id, + ignore_mismatched_sizes=True, config=config, ) - model.eval() - img = cv2.imread("../assets/cats.jpg", cv2.IMREAD_COLOR) - - if USE_HF_PREPROCESS: - image_processor = AutoImageProcessor.from_pretrained(hub_model_id) - img = image_processor(img, return_tensors="pt") - img = img["pixel_values"] - else: - img: np.array = cv2.resize(img, (224, 224), cv2.INTER_LINEAR) - img = (img.astype(np.float32) / 255.0 - np.array([0.5, 0.5, 0.5])) / np.array([0.5, 0.5, 0.5]) - img = torch.from_numpy(np.transpose(img, (2, 0, 1))[None, ...]) - - output = model(img) - labels = read_imagenet_labels() - for i, j in enumerate(torch.topk(output.logits[0], k=3).indices): - print(f"Top: {i} is {labels[int(j)]}") - - f = open("../models/vit.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") + export_wts(model, out_path) + + +if __name__ == "__main__": + main() diff --git a/vit/macros.h b/vit/macros.h index daae7bfb..6d851752 100644 --- a/vit/macros.h +++ b/vit/macros.h @@ -16,14 +16,19 @@ #endif #endif // API_EXPORTS -#define TRT_VERSION \ - ((NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + (NV_TENSORRT_PATCH * 10) + NV_TENSORRT_BUILD) +#define TRT_VERSION_ENCODE(major, minor, patch, build) \ + (((major) * 1000000) + ((minor) * 10000) + ((patch) * 100) + (build)) +#define TRT_VERSION TRT_VERSION_ENCODE(NV_TENSORRT_MAJOR, NV_TENSORRT_MINOR, NV_TENSORRT_PATCH, NV_TENSORRT_BUILD) +#define TRT_VERSION_GE(major, minor, patch) \ + ((NV_TENSORRT_MAJOR > (major)) || (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR > (minor)) || \ + (NV_TENSORRT_MAJOR == (major) && NV_TENSORRT_MINOR == (minor) && NV_TENSORRT_PATCH >= (patch))) +#define TRT_VERSION_LT(major, minor, patch) (!TRT_VERSION_GE((major), (minor), (patch))) -#if TRT_VERSION < 7220 +#if TRT_VERSION_LT(7, 2, 2) #error "TensorRT >= 7.2.2 is required for this demo." #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) #define TRT_NOEXCEPT noexcept #define TRT_CONST_ENQUEUE const #else diff --git a/vit/utils.h b/vit/utils.h index 6a86347c..92ba2a87 100644 --- a/vit/utils.h +++ b/vit/utils.h @@ -4,7 +4,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -25,7 +27,7 @@ namespace { } while (0) static void checkTrtEnv(int device = 0) { -#if TRT_VERSION < 8000 +#if TRT_VERSION_LT(8, 0, 0) CHECK(cudaGetDevice(&device)); cudaDeviceProp prop{}; CHECK(cudaGetDeviceProperties(&prop, device)); @@ -162,4 +164,43 @@ static auto loadImagenetLabelMap(const std::string& path) { } return labels; } + +static constexpr int32_t kBenchmarkRuns = 200; +static constexpr std::size_t kMaxFirstOutputs = 10; + +inline auto percentile(const std::vector& sorted, double percent) -> double { + assert(!sorted.empty()); + const double rank = percent / 100.0 * static_cast(sorted.size() - 1); + const auto lower = static_cast(rank); + const auto upper = std::min(lower + 1, sorted.size() - 1); + if (lower == upper) { + return sorted[lower]; + } + const double weight = rank - static_cast(lower); + return sorted[lower] * (1.0 - weight) + sorted[upper] * weight; +} + +inline void printBenchmark(const std::string& tag, const std::vector& latenciesMs, int64_t batchSize = 1) { + assert(!latenciesMs.empty()); + auto sorted = latenciesMs; + std::sort(sorted.begin(), sorted.end()); + const double avg = + std::accumulate(latenciesMs.begin(), latenciesMs.end(), 0.0) / static_cast(latenciesMs.size()); + std::cout << "[" << tag << "] benchmark_runs=" << latenciesMs.size() << " batch=" << batchSize << " AVG=" << avg + << "ms P50=" << percentile(sorted, 50.0) << "ms P90=" << percentile(sorted, 90.0) + << "ms P95=" << percentile(sorted, 95.0) << "ms P99=" << percentile(sorted, 99.0) << "ms\n"; +} + +inline void printFirstOutputs(const std::string& tag, const float* values, std::size_t count) { + const auto limit = std::min(count, kMaxFirstOutputs); + std::cout << "[" << tag << "] first_outputs="; + for (std::size_t i = 0; i < limit; ++i) { + if (i > 0) { + std::cout << ", "; + } + std::cout << std::setprecision(4) << values[i]; + } + std::cout << '\n'; +} + } // namespace diff --git a/vit/vit.cc b/vit/vit.cc index 0a95b60f..808bcb61 100644 --- a/vit/vit.cc +++ b/vit/vit.cc @@ -1,9 +1,24 @@ #include #include +#include +#include #include +#include +#include +#include +#include +#include +#include #include +#include #include #include +#include +#include +#include +#include +#include + #include "cuda_allocator.h" #include "logging.h" #include "macros.h" @@ -16,17 +31,124 @@ using M = nvinfer1::MatrixOperation; using E = nvinfer1::ElementWiseOperation; using NDCF = nvinfer1::NetworkDefinitionCreationFlag; -static constexpr const int64_t N = 1; -static constexpr const int64_t INPUT_H = 224; -static constexpr const int64_t INPUT_W = 224; +// DataType::kHALF -> FP16 +// DataType::kFLOAT -> FP32 +static constexpr DataType BUILD_PRECISION = DataType::kHALF; +static constexpr int64_t BUILD_MIN_BATCH = 1; +static constexpr int64_t BUILD_OPT_BATCH = 1; +static constexpr int64_t BUILD_MAX_BATCH = 2; + +template +requires(sizeof(T) > sizeof(int32_t)) static auto toI32(T value) -> int32_t { + assert(std::in_range(value)); + return static_cast(value); +} + +// ViT model variant table. +// All variants share the same architecture; only sizes differ. +// `model_type` is the canonical name (e.g. "ViT-B/16"); aliases without slash +// (e.g. "b16") are accepted via normalizeModelType(). +struct ViTConfig { + int64_t hidden = 768; + int64_t num_layers = 12; + int64_t num_heads = 12; + int64_t ff_dim = 3072; + int64_t num_classes = 1000; + int64_t patch = 16; + int64_t img_size = 224; + // Optimization profile used by the generated engine. + int64_t min_batch = BUILD_MIN_BATCH; + int64_t opt_batch = BUILD_OPT_BATCH; + int64_t max_batch = BUILD_MAX_BATCH; + float lnorm_eps = 1e-12f; + std::string model_type = "ViT-B/16"; + std::string wts_path; + std::string engine_path; + [[nodiscard]] int64_t num_patches() const { return (img_size / patch) * (img_size / patch); } + [[nodiscard]] int64_t seq_len() const { return num_patches() + 1; } +}; + +static std::string normalizeModelType(const std::string& s) { + // Accept "ViT-B/16", "vit-b/16", "b16", "B/16" -> "ViT-B/16". + std::string t; + for (char c : s) { + if (c != '-' && c != '/' && c != ' ') { + t += static_cast(std::toupper(static_cast(c))); + } + } + // Strip leading "VIT" + if (t.starts_with("VIT")) + t = t.substr(3); + if (t == "B16") + return "ViT-B/16"; + if (t == "B32") + return "ViT-B/32"; + if (t == "L16") + return "ViT-L/16"; + if (t == "L32") + return "ViT-L/32"; + if (t == "H14") + return "ViT-H/14"; + return s; // unknown, return as-is for error reporting +} + +static ViTConfig getVariantConfig(const std::string& raw_name) { + const std::string name = normalizeModelType(raw_name); + ViTConfig c; + c.model_type = name; + if (name == "ViT-B/16") { + c.hidden = 768; + c.num_layers = 12; + c.num_heads = 12; + c.ff_dim = 3072; + c.patch = 16; + c.img_size = 224; + c.num_classes = 1000; + } else if (name == "ViT-B/32") { + c.hidden = 768; + c.num_layers = 12; + c.num_heads = 12; + c.ff_dim = 3072; + c.patch = 32; + c.img_size = 384; + c.num_classes = 1000; + } else if (name == "ViT-L/16") { + c.hidden = 1024; + c.num_layers = 24; + c.num_heads = 16; + c.ff_dim = 4096; + c.patch = 16; + c.img_size = 224; + c.num_classes = 1000; + } else if (name == "ViT-L/32") { + c.hidden = 1024; + c.num_layers = 24; + c.num_heads = 16; + c.ff_dim = 4096; + c.patch = 32; + c.img_size = 384; + c.num_classes = 1000; + } else if (name == "ViT-H/14") { + // HF only ships an ImageNet-21k checkpoint for huge: 21843 classes. + c.hidden = 1280; + c.num_layers = 32; + c.num_heads = 16; + c.ff_dim = 5120; + c.patch = 14; + c.img_size = 224; + c.num_classes = 21843; + } else { + std::cerr << "Unknown model_type: " << raw_name + << " (expected ViT-B/16 | ViT-B/32 | ViT-L/16 | ViT-L/32 | ViT-H/14)\n"; + std::abort(); + } + return c; +} -static constexpr const char* WTS_PATH = "../models/vit.wts"; -static constexpr const char* ENGINE_PATH = "../models/vit.engine"; -static constexpr const char* LABELS_PATH = "../assets/imagenet1000_clsidx_to_labels.txt"; static constexpr const std::array NAMES = {"input", "logits"}; -static constexpr const std::array SIZES = {3 * INPUT_H * INPUT_W, 1000}; static constexpr const std::array mean = {0.5f, 0.5f, 0.5f}; static constexpr const std::array stdv = {0.5f, 0.5f, 0.5f}; +static constexpr const char* LABELS_PATH = "assets/imagenet1000_clsidx_to_labels.txt"; static Logger gLogger; @@ -36,12 +158,14 @@ static auto bytesPerElement(DataType t) -> std::size_t { return 4; case DataType::kHALF: return 2; + case DataType::kINT64: + return 8; case DataType::kINT32: return 4; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) case DataType::kBOOL: #endif -#if TRT_VERSION >= 8500 +#if TRT_VERSION_GE(8, 5, 0) case DataType::kUINT8: #endif case DataType::kINT8: @@ -76,11 +200,90 @@ static void convertWeightMapToHalf(WeightMap& w) { struct ViTParam { uint32_t index; uint32_t head_num; + int64_t hidden; + int64_t ff_dim; float lnorm_eps = 1e-12f; }; +struct RepeatedWeightsStorage { + std::vector> f32; + std::vector> f16; + std::vector> i64; +}; + +static auto makeRepeatedBatchWeights(const Weights& src, int64_t repeat, RepeatedWeightsStorage& storage) -> Weights { + if (repeat <= 0 || src.values == nullptr || src.count <= 0) { + std::cerr << "invalid repeated weight\n"; + std::abort(); + } + + if (src.type == DataType::kHALF) { + const auto* values = reinterpret_cast(src.values); + storage.f16.emplace_back(static_cast(src.count * repeat)); + auto& dst = storage.f16.back(); + for (int64_t i = 0; i < repeat; ++i) { + std::copy(values, values + src.count, dst.begin() + i * src.count); + } + return Weights{src.type, dst.data(), static_cast(dst.size())}; + } + + if (src.type == DataType::kFLOAT) { + const auto* values = reinterpret_cast(src.values); + storage.f32.emplace_back(static_cast(src.count * repeat)); + auto& dst = storage.f32.back(); + for (int64_t i = 0; i < repeat; ++i) { + std::copy(values, values + src.count, dst.begin() + i * src.count); + } + return Weights{src.type, dst.data(), static_cast(dst.size())}; + } + + std::cerr << "unsupported repeated weight dtype\n"; + std::abort(); +} + +static auto addShapeConstant(INetworkDefinition* net, std::initializer_list values, + RepeatedWeightsStorage& storage) -> ITensor* { + storage.i64.emplace_back(values); + auto& data = storage.i64.back(); + Dims dims{}; + dims.nbDims = 1; + dims.d[0] = static_cast(data.size()); + auto* c = net->addConstant( + dims, Weights{.type = DataType::kINT64, .values = data.data(), .count = static_cast(data.size())}); + return c->getOutput(0); +} + +static auto addBatchedConstantSlice(INetworkDefinition* net, ITensor& input, ITensor& full, int64_t dim1, int64_t dim2, + RepeatedWeightsStorage& storage) -> ITensor* { + auto* input_shape = net->addShape(input); + Dims batch_start{}; + batch_start.nbDims = 1; + batch_start.d[0] = 0; + Dims batch_size{}; + batch_size.nbDims = 1; + batch_size.d[0] = 1; + Dims batch_stride{}; + batch_stride.nbDims = 1; + batch_stride.d[0] = 1; + auto* batch = net->addSlice(*input_shape->getOutput(0), batch_start, batch_size, batch_stride); + auto* size_tail = addShapeConstant(net, {dim1, dim2}, storage); + const std::array size_inputs = {batch->getOutput(0), size_tail}; + auto* size = net->addConcatenation(size_inputs.data(), toI32(size_inputs.size())); + size->setAxis(0); + + auto* start = addShapeConstant(net, {0, 0, 0}, storage); + auto* stride = addShapeConstant(net, {1, 1, 1}, storage); + const auto full_dims = full.getDimensions(); + auto* slice = + net->addSlice(full, Dims3{0, 0, 0}, Dims3{full_dims.d[0], full_dims.d[1], full_dims.d[2]}, Dims3{1, 1, 1}); + slice->setInput(1, *start); + slice->setInput(2, *size->getOutput(0)); + slice->setInput(3, *stride); + return slice->getOutput(0); +} + static auto addGeLU(INetworkDefinition* net, ITensor& input) -> ILayer* { -#if TRT_VERSION < 10000 +#if TRT_VERSION_LT(10, 0, 0) // tanh approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) const auto inputDims = input.getDimensions(); @@ -94,10 +297,11 @@ static auto addGeLU(INetworkDefinition* net, ITensor& input) -> ILayer* { static float _one = 1.0f; static float _sqrt_2_div_pi = std::sqrt(2.0f / M_PI); static float _coeff = 0.044715f; - auto* _w_half = net->addConstant(scalarDims, Weights{DataType::kFLOAT, &_half, 1}); - auto* _w_one = net->addConstant(scalarDims, Weights{DataType::kFLOAT, &_one, 1}); - auto* _w_sqrt_2_div_pi = net->addConstant(scalarDims, Weights{DataType::kFLOAT, &_sqrt_2_div_pi, 1}); - auto* _w_coeff = net->addConstant(scalarDims, Weights{DataType::kFLOAT, &_coeff, 1}); + auto* _w_half = net->addConstant(scalarDims, Weights{.type = DataType::kFLOAT, .values = &_half, .count = 1}); + auto* _w_one = net->addConstant(scalarDims, Weights{.type = DataType::kFLOAT, .values = &_one, .count = 1}); + auto* _w_sqrt_2_div_pi = + net->addConstant(scalarDims, Weights{.type = DataType::kFLOAT, .values = &_sqrt_2_div_pi, .count = 1}); + auto* _w_coeff = net->addConstant(scalarDims, Weights{.type = DataType::kFLOAT, .values = &_coeff, .count = 1}); auto* _x2 = net->addElementWise(input, input, E::kPROD); auto* x3_0 = net->addElementWise(*_x2->getOutput(0), input, E::kPROD); @@ -117,7 +321,7 @@ static auto addGeLU(INetworkDefinition* net, ITensor& input) -> ILayer* { static auto addLinearNorm(INetworkDefinition* net, ITensor& input, ITensor& scale, ITensor& bias, uint32_t axesMask) noexcept -> ILayer* { -#if TRT_VERSION >= 11500 +#if TRT_VERSION_GE(10, 15, 0) auto* ln = net->addNormalizationV2(input, scale, bias, axesMask); #else auto* ln = net->addNormalization(input, scale, bias, axesMask); @@ -129,11 +333,13 @@ static auto addLinearNorm(INetworkDefinition* net, ITensor& input, ITensor& scal auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTParam& param) -> ITensor* { std::string name = "vit.encoder.layer." + std::to_string(param.index); auto attn_name = name + ".attention"; - int64_t attn_head_size = 768LL / param.head_num; + int64_t H = param.hidden; + int64_t F = param.ff_dim; + int64_t attn_head_size = H / param.head_num; - auto* qw = net->addConstant(Dims3{1, 768, 768}, w.at(attn_name + ".attention.query.weight")); - auto* kw = net->addConstant(Dims3{1, 768, 768}, w.at(attn_name + ".attention.key.weight")); - auto* vw = net->addConstant(Dims3{1, 768, 768}, w.at(attn_name + ".attention.value.weight")); + auto* qw = net->addConstant(Dims3{1, H, H}, w.at(attn_name + ".attention.query.weight")); + auto* kw = net->addConstant(Dims3{1, H, H}, w.at(attn_name + ".attention.key.weight")); + auto* vw = net->addConstant(Dims3{1, H, H}, w.at(attn_name + ".attention.value.weight")); /* 1. layer norm before attention */ auto pre_ln_name = name + ".layernorm_before"; auto dims = input.getDimensions(); @@ -143,9 +349,9 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa auto* pre_lnorm = addLinearNorm(net, input, *ln_scale->getOutput(0), *ln_bias->getOutput(0), axes); /** 2. multi-head self-attention */ - auto* qb = net->addConstant(Dims3{1, 1, 768}, w.at(attn_name + ".attention.query.bias")); - auto* kb = net->addConstant(Dims3{1, 1, 768}, w.at(attn_name + ".attention.key.bias")); - auto* vb = net->addConstant(Dims3{1, 1, 768}, w.at(attn_name + ".attention.value.bias")); + auto* qb = net->addConstant(Dims3{1, 1, H}, w.at(attn_name + ".attention.query.bias")); + auto* kb = net->addConstant(Dims3{1, 1, H}, w.at(attn_name + ".attention.key.bias")); + auto* vb = net->addConstant(Dims3{1, 1, H}, w.at(attn_name + ".attention.value.bias")); auto* _lno = pre_lnorm->getOutput(0); // 2.1 Q, K attention matmul auto* _q_attn = net->addMatrixMultiply(*_lno, M::kNONE, *qw->getOutput(0), M::kTRANSPOSE); @@ -181,18 +387,20 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa auto* qk_scale_w = net->addConstant(Dims4{1, 1, 1, 1}, w.at(attn_name + ".scale")); // 2.3 QKV attention output and reshape -#if TRT_VERSION >= 11400 && TRT_VERSION < 11500 +#if TRT_VERSION_GE(10, 14, 0) && TRT_VERSION_LT(10, 15, 0) gLogger.log(Severity::kWARNING, "IAttention is available in TensorRT 10.14.1 SDK but have bugs, use 10.15.1+ to enable native fused " "kernel"); #endif -#if TRT_VERSION >= 11500 +#if TRT_VERSION_GE(10, 15, 0) using ANO = AttentionNormalizationOp; auto* q_scaled = net->addElementWise(*q_s->getOutput(0), *qk_scale_w->getOutput(0), E::kPROD)->getOutput(0); auto* attn = net->addAttention(*q_scaled, *k_s->getOutput(0), *v_s->getOutput(0), ANO::kSOFTMAX, false); assert(attn != nullptr); - auto status = attn->setDecomposable(false); - assert(status); + if (!attn->setDecomposable(false)) { + std::cerr << "setDecomposable failed\n"; + std::abort(); + } auto* attn_out = net->addShuffle(*attn->getOutput(0)); #else auto* qk = net->addMatrixMultiply(*q_s->getOutput(0), M::kNONE, *k_s->getOutput(0), M::kTRANSPOSE); @@ -204,10 +412,10 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa auto* attn_out = net->addShuffle(*attn_qkv->getOutput(0)); #endif attn_out->setFirstTranspose({0, 2, 1, 3}); - attn_out->setReshapeDimensions(Dims3{0, 0, 768}); + attn_out->setReshapeDimensions(Dims3{0, 0, H}); // 2.4 attention output projection - auto* out_proj_w = net->addConstant(Dims3{1, 768, 768}, w.at(attn_name + ".output.dense.weight"))->getOutput(0); - auto* out_proj_b = net->addConstant(Dims3{1, 1, 768}, w.at(attn_name + ".output.dense.bias"))->getOutput(0); + auto* out_proj_w = net->addConstant(Dims3{1, H, H}, w.at(attn_name + ".output.dense.weight"))->getOutput(0); + auto* out_proj_b = net->addConstant(Dims3{1, 1, H}, w.at(attn_name + ".output.dense.bias"))->getOutput(0); auto* attn_fcw = net->addMatrixMultiply(*attn_out->getOutput(0), M::kNONE, *out_proj_w, M::kTRANSPOSE); auto* attn_fcb = net->addElementWise(*attn_fcw->getOutput(0), *out_proj_b, E::kSUM); attn_fcb->setName((attn_name + ".out_proj").c_str()); @@ -227,8 +435,8 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa /** 6. intermediate (feed-forward) layer and activation */ auto intermediate_name = name + ".intermediate.dense"; std::cout << "Building: " << intermediate_name << "\n"; - auto* iw = net->addConstant(Dims3{1, 3072, 768}, w[intermediate_name + ".weight"]); - auto* ib = net->addConstant(Dims3{1, 1, 3072}, w[intermediate_name + ".bias"]); + auto* iw = net->addConstant(Dims3{1, F, H}, w[intermediate_name + ".weight"]); + auto* ib = net->addConstant(Dims3{1, 1, F}, w[intermediate_name + ".bias"]); ib->setName((intermediate_name + ".bias").c_str()); auto* inter0 = net->addMatrixMultiply(*post_lnorm->getOutput(0), M::kNONE, *iw->getOutput(0), M::kTRANSPOSE); auto* inter1 = net->addElementWise(*inter0->getOutput(0), *ib->getOutput(0), E::kSUM); @@ -237,8 +445,8 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa /** 7. output projection */ auto output_name = name + ".output.dense"; std::cout << "Building: " << output_name << "\n"; - auto* ow = net->addConstant(Dims3{1, 768, 3072}, w[output_name + ".weight"]); - auto* ob = net->addConstant(Dims3{1, 1, 768}, w[output_name + ".bias"]); + auto* ow = net->addConstant(Dims3{1, H, F}, w[output_name + ".weight"]); + auto* ob = net->addConstant(Dims3{1, 1, H}, w[output_name + ".bias"]); ob->setName((output_name + ".bias").c_str()); auto* out0 = net->addMatrixMultiply(*inter_act->getOutput(0), M::kNONE, *ow->getOutput(0), M::kTRANSPOSE); auto* out1 = net->addElementWise(*out0->getOutput(0), *ob->getOutput(0), E::kSUM); @@ -250,65 +458,99 @@ auto ViTLayer(INetworkDefinition* net, WeightMap& w, ITensor& input, const ViTPa } // Creat the engine using only the API without any parser. -auto createEngine(int64_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, +auto createEngine(const ViTConfig& cfg, IRuntime* runtime, IBuilder* builder, IBuilderConfig* config, DataType dt) -> ICudaEngine* { - WeightMap w = loadWeights(WTS_PATH); + WeightMap w = loadWeights(cfg.wts_path); if (dt == DataType::kHALF) { convertWeightMapToHalf(w); } -#if TRT_VERSION >= 10000 + RepeatedWeightsStorage repeated_storage; + +#if TRT_VERSION_GE(10, 0, 0) auto* net = builder->createNetworkV2(1U << static_cast(NDCF::kSTRONGLY_TYPED)); #else auto* net = builder->createNetworkV2(1U << static_cast(NDCF::kEXPLICIT_BATCH)); #endif // 1. patch embedding - ITensor* data = net->addInput(NAMES[0], dt, Dims4{-1, 3, INPUT_H, INPUT_W}); + Dims input_dims{.nbDims = 4, .d = {-1, 3, cfg.img_size, cfg.img_size}}; + ITensor* data = net->addInput(NAMES[0], dt, input_dims); std::string name = "vit.embeddings.patch_embeddings.projection."; - auto* embed = net->addConvolutionNd(*data, 768, DimsHW{16, 16}, w[name + "weight"], w[name + "bias"]); + auto* embed = net->addConvolutionNd(*data, cfg.hidden, DimsHW{cfg.patch, cfg.patch}, w[name + "weight"], + w[name + "bias"]); embed->setName("patch embedding"); - embed->setStrideNd(DimsHW{16, 16}); + embed->setStrideNd(DimsHW{cfg.patch, cfg.patch}); auto* s = net->addShuffle(*embed->getOutput(0)); - s->setReshapeDimensions(Dims3{0, 768, 14LL * 14}); + s->setReshapeDimensions(Dims3{0, cfg.hidden, cfg.num_patches()}); s->setSecondTranspose({0, 2, 1}); // 2. add cls token and position embedding - auto* cls_token = net->addConstant(Dims3{1, 1, 768}, w["vit.embeddings.cls_token"]); - auto* pos_embed = net->addConstant(Dims3{1, 197, 768}, w["vit.embeddings.position_embeddings"]); - const std::array _cat = {cls_token->getOutput(0), s->getOutput(0)}; + auto cls_weights = makeRepeatedBatchWeights(w["vit.embeddings.cls_token"], cfg.max_batch, repeated_storage); + auto pos_weights = + makeRepeatedBatchWeights(w["vit.embeddings.position_embeddings"], cfg.max_batch, repeated_storage); + auto* cls_full = net->addConstant(Dims3{cfg.max_batch, 1, cfg.hidden}, cls_weights); + auto* pos_full = net->addConstant(Dims3{cfg.max_batch, cfg.seq_len(), cfg.hidden}, pos_weights); + auto* cls_token = addBatchedConstantSlice(net, *data, *cls_full->getOutput(0), 1, cfg.hidden, repeated_storage); + auto* pos_embed = + addBatchedConstantSlice(net, *data, *pos_full->getOutput(0), cfg.seq_len(), cfg.hidden, repeated_storage); + const std::array _cat = {cls_token, s->getOutput(0)}; auto* cat = net->addConcatenation(_cat.data(), 2); cat->setAxis(1); cat->setName("cat_clstoken_embed"); - auto* pos_added = net->addElementWise(*cat->getOutput(0), *pos_embed->getOutput(0), ElementWiseOperation::kSUM); + auto* pos_added = net->addElementWise(*cat->getOutput(0), *pos_embed, ElementWiseOperation::kSUM); pos_added->setName("position_embed"); // 3. transformer encoder layers ITensor* input = pos_added->getOutput(0); - for (auto i = 0u; i < 12; i++) { - auto* vit = ViTLayer(net, w, *input, {.index = i, .head_num = 12, .lnorm_eps = 1e-12f}); + for (int64_t i = 0; i < cfg.num_layers; i++) { + auto* vit = ViTLayer(net, w, *input, + {.index = static_cast(i), + .head_num = static_cast(cfg.num_heads), + .hidden = cfg.hidden, + .ff_dim = cfg.ff_dim, + .lnorm_eps = cfg.lnorm_eps}); input = vit; } // 4. layer norm after transformer encoder - auto* ln_scale = net->addConstant(Dims3{1, 1, 768}, w["vit.layernorm.weight"]); - auto* ln_bias = net->addConstant(Dims3{1, 1, 768}, w["vit.layernorm.bias"]); + auto* ln_scale = net->addConstant(Dims3{1, 1, cfg.hidden}, w["vit.layernorm.weight"]); + auto* ln_bias = net->addConstant(Dims3{1, 1, cfg.hidden}, w["vit.layernorm.bias"]); uint32_t axes = 1U << static_cast(input->getDimensions().nbDims - 1); auto* post_lnorm = addLinearNorm(net, *input, *ln_scale->getOutput(0), *ln_bias->getOutput(0), axes); - // 6. classifier head - auto* slice = net->addSlice(*post_lnorm->getOutput(0), Dims3{0, 0, 0}, Dims3{N, 1, 768}, Dims3{1, 1, 1}); - auto* shuffle = net->addShuffle(*slice->getOutput(0)); - shuffle->setReshapeDimensions(Dims2{N, 768}); - auto* cls_w = net->addConstant(DimsHW{1000, 768}, w["classifier.weight"]); - auto* cls_b = net->addConstant(DimsHW{1, 1000}, w["classifier.bias"]); + // 6. classifier head -- take CLS token (index 0 along seq axis) via Gather + static int32_t cls_idx_data = 0; + Weights cls_idx_w{DataType::kINT32, &cls_idx_data, 1}; + Dims idx_dims; + idx_dims.nbDims = 1; + idx_dims.d[0] = 1; + auto* idx_const = net->addConstant(idx_dims, cls_idx_w)->getOutput(0); + auto* gather = net->addGather(*post_lnorm->getOutput(0), *idx_const, 1); + auto* shuffle = net->addShuffle(*gather->getOutput(0)); + shuffle->setReshapeDimensions(Dims2{-1, cfg.hidden}); + auto* cls_w = net->addConstant(DimsHW{cfg.num_classes, cfg.hidden}, w["classifier.weight"]); + auto* cls_b = net->addConstant(DimsHW{1, cfg.num_classes}, w["classifier.bias"]); auto* cls_0 = net->addMatrixMultiply(*shuffle->getOutput(0), M::kNONE, *cls_w->getOutput(0), M::kTRANSPOSE); auto* cls_1 = net->addElementWise(*cls_b->getOutput(0), *cls_0->getOutput(0), E::kSUM); + cls_1->getOutput(0)->setName(NAMES[1]); net->markOutput(*cls_1->getOutput(0)); - Dims4 _min{1, 3, INPUT_H, INPUT_W}, _opt{N, 3, INPUT_H, INPUT_W}, _max{2 * N, 3, INPUT_H, INPUT_W}; -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, WORKSPACE_SIZE); - config->setHardwareCompatibilityLevel(HardwareCompatibilityLevel::kAMPERE_PLUS); + config->setBuilderOptimizationLevel(5); +#if TRT_VERSION_LT(10, 0, 0) + // Strongly-typed networks (TRT 10+) take their precision from the + // tensor types declared on inputs/weights and reject BuilderFlag::kFP16. + // Pre-TRT-10 networks are weakly typed, so explicitly request FP16 + // tactics when BUILD_PRECISION == kHALF; otherwise TRT may fall back + // to FP32 kernels even though the inputs/weights are half. + if (BUILD_PRECISION == DataType::kHALF) { + config->setFlag(BuilderFlag::kFP16); + } +#endif + Dims _min{.nbDims = 4, .d = {cfg.min_batch, 3, cfg.img_size, cfg.img_size}}; + Dims _opt{.nbDims = 4, .d = {cfg.opt_batch, 3, cfg.img_size, cfg.img_size}}; + Dims _max{.nbDims = 4, .d = {cfg.max_batch, 3, cfg.img_size, cfg.img_size}}; auto* profile = builder->createOptimizationProfile(); profile->setDimensions(NAMES[0], OptProfileSelector::kMIN, _min); profile->setDimensions(NAMES[0], OptProfileSelector::kOPT, _opt); @@ -318,7 +560,7 @@ auto createEngine(int64_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfi ICudaEngine* engine = runtime->deserializeCudaEngine(mem->data(), mem->size()); delete net; #else - builder->setMaxBatchSize(N); + builder->setMaxBatchSize(cfg.max_batch); config->setMaxWorkspaceSize(WORKSPACE_SIZE); ICudaEngine* engine = builder->buildEngineWithConfig(*net, *config); net->destroy(); @@ -341,49 +583,71 @@ auto createEngine(int64_t N, IRuntime* runtime, IBuilder* builder, IBuilderConfi return engine; } -std::vector> doInference(IExecutionContext& context, __half* input, std::size_t batchSize) { +std::vector> doInference(IExecutionContext& context, const void* input, std::size_t batchSize, + const ViTConfig& cfg) { const ICudaEngine& engine = context.getEngine(); cudaStream_t stream; CHECK(cudaStreamCreate(&stream)); std::vector buffers; -#if TRT_VERSION >= 10000 +#if TRT_VERSION_GE(10, 0, 0) auto allocator = CudaOutputAllocator::Create(stream); #endif -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) const int32_t nIO = engine.getNbIOTensors(); #else const int32_t nIO = engine.getNbBindings(); #endif + // SIZES per IO: input is C*H*W per sample; output is num_classes per sample. + const int64_t in_per_sample = 3LL * cfg.img_size * cfg.img_size; + const int64_t out_per_sample = cfg.num_classes; + auto sizeOf = [&](int i) -> int64_t { + return i == 0 ? in_per_sample : out_per_sample; + }; + +#if TRT_VERSION_GE(8, 0, 0) + if (engine.getTensorShape(NAMES[0]).d[0] == -1) { + Dims in_dims{.nbDims = 4, .d = {static_cast(batchSize), 3, cfg.img_size, cfg.img_size}}; + if (!context.setInputShape(NAMES[0], in_dims)) { + std::cerr << "setInputShape failed batch=" << batchSize << "\n"; + std::abort(); + } + } +#endif + buffers.resize(nIO, nullptr); for (auto i = 0; i < nIO; ++i) { -#if TRT_VERSION >= 8000 - // TensorRT 8+ use name based SDK +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); const auto dtype = engine.getTensorDataType(tensor_name); - std::size_t size = batchSize * SIZES[i] * bytesPerElement(dtype); -#if TRT_VERSION >= 10000 - // TensorRT 10+ use outuput allocator - if (i == 0) { + std::size_t size = batchSize * sizeOf(i) * bytesPerElement(dtype); +#if TRT_VERSION_GE(10, 0, 0) + if (engine.getTensorIOMode(tensor_name) == TensorIOMode::kINPUT) { CHECK(cudaMalloc(&buffers[i], size)); CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } } else { context.setOutputAllocator(tensor_name, allocator.get()); } #else - if (i != 0) { + if (engine.getTensorIOMode(tensor_name) == TensorIOMode::kINPUT) { CHECK(cudaMalloc(&buffers[i], size)); + CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } else { CHECK(cudaMalloc(&buffers[i], size)); - CHECK(cudaMemcpyAsync(buffers[i], input, size, cudaMemcpyHostToDevice, stream)); } - context.setTensorAddress(tensor_name, buffers[i]); + if (!context.setTensorAddress(tensor_name, buffers[i])) { + std::cerr << "setTensorAddress failed\n"; + std::abort(); + } #endif #else - std::size_t size = batchSize * SIZES[i] * sizeof(float); + std::size_t size = batchSize * sizeOf(i) * sizeof(float); const int32_t idx = engine.getBindingIndex(NAMES[i]); assert(idx == i); CHECK(cudaMalloc(&buffers[i], size)); @@ -393,37 +657,50 @@ std::vector> doInference(IExecutionContext& context, __half* #endif } -#if TRT_VERSION >= 8000 - assert(context.enqueueV3(stream)); +#if TRT_VERSION_GE(8, 0, 0) + if (!context.enqueueV3(stream)) { + std::cerr << "enqueueV3 failed\n"; + std::abort(); + } #else - assert(context.enqueueV2(buffers.data(), stream, nullptr)); + if (!context.enqueueV2(buffers.data(), stream, nullptr)) { + std::cerr << "enqueueV2 failed\n"; + std::abort(); + } #endif std::vector> prob; - for (int i = 1; i < nIO; ++i) { -#if TRT_VERSION >= 10000 + for (int i = 0; i < nIO; ++i) { +#if TRT_VERSION_GE(8, 0, 0) auto* tensor_name = engine.getIOTensorName(i); + if (engine.getTensorIOMode(tensor_name) != TensorIOMode::kOUTPUT) + continue; const auto dtype = engine.getTensorDataType(tensor_name); - std::size_t size = batchSize * SIZES[i] * bytesPerElement(dtype); + std::size_t count = batchSize * out_per_sample; + std::size_t size = count * bytesPerElement(dtype); +#if TRT_VERSION_GE(10, 0, 0) void* out_ptr = allocator->getBuffer(tensor_name); - // D2H data transfer +#else + void* out_ptr = buffers[i]; +#endif if (dtype == DataType::kHALF) { - std::vector<__half> tmp_h(batchSize * SIZES[i]); + std::vector<__half> tmp_h(count); CHECK(cudaMemcpyAsync(tmp_h.data(), out_ptr, size, cudaMemcpyDeviceToHost, stream)); CHECK(cudaStreamSynchronize(stream)); - std::vector tmp(batchSize * SIZES[i]); - for (std::size_t j = 0; j < tmp.size(); ++j) { + std::vector tmp(count); + for (std::size_t j = 0; j < tmp.size(); ++j) tmp[j] = __half2float(tmp_h[j]); - } prob.emplace_back(std::move(tmp)); } else { - std::vector tmp(batchSize * SIZES[i], std::nanf("")); + std::vector tmp(count, std::nanf("")); CHECK(cudaMemcpyAsync(tmp.data(), out_ptr, size, cudaMemcpyDeviceToHost, stream)); prob.emplace_back(std::move(tmp)); } #else - std::vector tmp(batchSize * SIZES[i], std::nanf("")); - std::size_t size = batchSize * SIZES[i] * sizeof(float); + if (i == 0) + continue; + std::vector tmp(batchSize * sizeOf(i), std::nanf("")); + std::size_t size = batchSize * sizeOf(i) * sizeof(float); CHECK(cudaMemcpyAsync(tmp.data(), buffers[i], size, cudaMemcpyDeviceToHost, stream)); prob.emplace_back(std::move(tmp)); #endif @@ -435,23 +712,23 @@ std::vector> doInference(IExecutionContext& context, __half* CHECK(cudaFree(buffer)); } } -#if TRT_VERSION >= 10000 +#if TRT_VERSION_GE(10, 0, 0) allocator.reset(); #endif CHECK(cudaStreamDestroy(stream)); return prob; } -void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { +void APIToModel(const ViTConfig& cfg, IRuntime* runtime, IHostMemory** modelStream) { IBuilder* builder = createInferBuilder(gLogger); IBuilderConfig* config = builder->createBuilderConfig(); - ICudaEngine* engine = createEngine(N, runtime, builder, config, DataType::kHALF); + ICudaEngine* engine = createEngine(cfg, runtime, builder, config, BUILD_PRECISION); assert(engine != nullptr); (*modelStream) = engine->serialize(); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete engine; delete config; delete builder; @@ -462,110 +739,241 @@ void APIToModel(int32_t N, IRuntime* runtime, IHostMemory** modelStream) { #endif } +// Recover ViTConfig fields needed by inference (img_size, num_classes) from the +// deserialized engine. The other fields (hidden/layers/heads/ff/patch) are baked +// into the engine and not needed at runtime. +static ViTConfig configFromEngine(const ICudaEngine& engine) { + ViTConfig cfg; + auto in_dims = engine.getTensorShape(NAMES[0]); // [N or -1, 3, H, W] + auto out_dims = engine.getTensorShape(NAMES[1]); // [N or -1, num_classes] + if (in_dims.nbDims != 4 || out_dims.nbDims != 2) { + std::cerr << "Unexpected engine IO shapes\n"; + std::abort(); + } + cfg.img_size = in_dims.d[2]; + cfg.num_classes = out_dims.d[1]; + if (in_dims.d[0] == -1) { + auto pmin = engine.getProfileShape(NAMES[0], 0, OptProfileSelector::kMIN); + auto popt = engine.getProfileShape(NAMES[0], 0, OptProfileSelector::kOPT); + auto pmax = engine.getProfileShape(NAMES[0], 0, OptProfileSelector::kMAX); + cfg.min_batch = pmin.d[0]; + cfg.opt_batch = popt.d[0]; + cfg.max_batch = pmax.d[0]; + } else { + cfg.min_batch = in_dims.d[0]; + cfg.opt_batch = in_dims.d[0]; + cfg.max_batch = in_dims.d[0]; + } + return cfg; +} + +static auto loadLabels() -> std::map { + auto labels = loadImagenetLabelMap(LABELS_PATH); + if (labels.empty()) { + std::cerr << "failed to load labels: " << LABELS_PATH << "\n"; + std::abort(); + } + return labels; +} + +// Collect image files from a path: if a directory, return all *.jpg/*.jpeg/*.png +// files; if a single file, return just that file. +static std::vector collectImages(const std::string& path) { + namespace fs = std::filesystem; + std::vector out; + fs::path p(path); + if (!fs::exists(p)) { + std::cerr << "image path does not exist: " << path << "\n"; + std::abort(); + } + if (fs::is_regular_file(p)) { + out.push_back(path); + } else if (fs::is_directory(p)) { + for (auto& e : fs::directory_iterator(p)) { + if (!e.is_regular_file()) + continue; + auto ext = e.path().extension().string(); + for (auto& c : ext) { + c = static_cast(std::tolower(static_cast(c))); + } + if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp") { + out.push_back(e.path().string()); + } + } + std::sort(out.begin(), out.end()); + } + return out; +} + auto main(int argc, char** argv) -> int { std::cout << "TensorRT version: " << TRT_VERSION << "\n"; - if (argc != 2) { - std::cerr << "arguments not right!\n"; - std::cerr << "./vit -s // serialize model to plan file\n"; - std::cerr << "./vit -d // deserialize plan file and run inference\n"; - + if (argc < 2 || (std::string(argv[1]) == "-s" && argc != 5) || + (std::string(argv[1]) == "-d" && argc != 4 && argc != 5)) { + std::cerr << "usage:\n" + << " ./vit -s \n" + << " ./vit -d [--profile]\n" + << " model_type: ViT-B/16 | ViT-B/32 | ViT-L/16 | ViT-L/32 | ViT-H/14\n" + << " (aliases: b16, b32, l16, l32, h14 also accepted)\n" + << " build precision is configured in code via BUILD_PRECISION (see top of vit.cc).\n"; return 1; } + std::string mode = argv[1]; + #ifndef NDEBUG - gLogger.setReportableSeverity(nvinfer1::ILogger::Severity::kVERBOSE); + gLogger.setReportableSeverity(nvinfer1::ILogger::Severity::kINFO); #endif IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); - char* trtModelStream{nullptr}; - std::streamsize size{0}; - if (std::string(argv[1]) == "-s") { + if (mode == "-s") { + std::string wts_path = argv[2]; + std::string engine_path = argv[3]; + std::string model_type = argv[4]; + ViTConfig cfg = getVariantConfig(model_type); + cfg.wts_path = wts_path; + cfg.engine_path = engine_path; + const char* prec_str = (BUILD_PRECISION == DataType::kHALF) ? "fp16" : "fp32"; + std::cout << "[cfg] " << cfg.model_type << " hidden=" << cfg.hidden << " layers=" << cfg.num_layers + << " heads=" << cfg.num_heads << " ff=" << cfg.ff_dim << " patch=" << cfg.patch + << " img=" << cfg.img_size << " classes=" << cfg.num_classes + << " batch[min/opt/max]=" << cfg.min_batch << "/" << cfg.opt_batch << "/" << cfg.max_batch + << " precision=" << prec_str << "\n[wts] " << wts_path << "\n[engine] " << engine_path << "\n"; + IHostMemory* modelStream{nullptr}; - APIToModel(N, runtime, &modelStream); + APIToModel(cfg, runtime, &modelStream); assert(modelStream != nullptr); - std::ofstream p(ENGINE_PATH, std::ios::binary | std::ios::trunc); + std::ofstream p(engine_path, std::ios::binary | std::ios::trunc); if (!p) { - std::cerr << "could not open plan output file\n"; - return -1; - } - if (modelStream->size() > static_cast(std::numeric_limits::max())) { - std::cerr << "this model is too large to serialize\n"; + std::cerr << "could not open plan output file: " << engine_path << "\n"; return -1; } const auto* data_ptr = reinterpret_cast(modelStream->data()); auto data_size = static_cast(modelStream->size()); p.write(data_ptr, data_size); -#if TRT_VERSION >= 8000 +#if TRT_VERSION_GE(8, 0, 0) delete modelStream; #else modelStream->destroy(); #endif return 0; - } else if (std::string(argv[1]) == "-d") { - std::ifstream file(ENGINE_PATH, 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 << "read engine file error!\n"; + } else if (mode == "-d") { + std::string engine_path = argv[2]; + std::string image_dir = argv[3]; + bool enable_profile = false; + if (argc == 5) { + if (std::string(argv[4]) != "--profile") { + std::cerr << "unknown option: " << argv[4] << "\n"; + return 1; + } + enable_profile = true; + } + + std::ifstream file(engine_path, std::ios::binary); + if (!file.good()) { + std::cerr << "read engine file error: " << engine_path << "\n"; + return -1; + } + file.seekg(0, file.end); + std::streamsize size = file.tellg(); + file.seekg(0, file.beg); + if (size <= 0) { + std::cerr << "empty engine file: " << engine_path << "\n"; return -1; } + std::vector trtModelStream(static_cast(size)); + file.read(trtModelStream.data(), size); + file.close(); -#if TRT_VERSION >= 8000 - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); +#if TRT_VERSION_GE(8, 0, 0) + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream.data(), size); #else - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream.data(), size, nullptr); #endif assert(engine != nullptr); + trtModelStream.clear(); + trtModelStream.shrink_to_fit(); auto* context = engine->createExecutionContext(); assert(context != nullptr); + Profiler profiler("VisionTransformerProfiler"); + if (enable_profile) { + context->setProfiler(&profiler); + } - // VIT use default BGR order - auto img = cv::imread("../assets/cats.jpg", cv::IMREAD_COLOR); - auto input = preprocess_img(img, false, mean, stdv, N, INPUT_H, INPUT_W); + ViTConfig cfg = configFromEngine(*engine); + std::cout << "[engine] img=" << cfg.img_size << " classes=" << cfg.num_classes + << " batch[min/opt/max]=" << cfg.min_batch << "/" << cfg.opt_batch << "/" << cfg.max_batch << "\n"; - Profiler profiler("VisionTransformerProfiler"); + auto images = collectImages(image_dir); + if (images.empty()) { + std::cerr << "no images found under: " << image_dir << "\n"; + return -1; + } + // Cap batch to engine's max profile. + int64_t infer_batch = std::min(static_cast(images.size()), cfg.max_batch); + std::cout << "[infer] found " << images.size() << " image(s), using batch=" << infer_batch << "\n"; + + std::vector<__half> input_buf; + input_buf.reserve(infer_batch * 3 * cfg.img_size * cfg.img_size); + for (int64_t i = 0; i < infer_batch; ++i) { + auto img = cv::imread(images[i], cv::IMREAD_COLOR); + if (img.empty()) { + std::cerr << "cannot read image: " << images[i] << "\n"; + return -1; + } + auto one = preprocess_img(img, false, mean, stdv, 1, toI32(cfg.img_size), toI32(cfg.img_size)); + input_buf.insert(input_buf.end(), one.begin(), one.end()); + } + + // Match the engine's declared input dtype. The engine input is FP16 when + // built with `-s ... fp16` (default) and FP32 when built with `... fp32`. + const auto in_dtype = engine->getTensorDataType(NAMES[0]); + std::vector input_buf_f32; + const void* input_ptr = nullptr; + if (in_dtype == DataType::kHALF) { + input_ptr = input_buf.data(); + } else if (in_dtype == DataType::kFLOAT) { + input_buf_f32.resize(input_buf.size()); + for (std::size_t k = 0; k < input_buf.size(); ++k) + input_buf_f32[k] = __half2float(input_buf[k]); + input_ptr = input_buf_f32.data(); + } else { + std::cerr << "unsupported engine input dtype\n"; + return -1; + } - // Warmup: run a few iterations without profiling. for (int i = 0; i < 5; ++i) { - (void)doInference(*context, input.data(), N); + (void)doInference(*context, input_ptr, infer_batch, cfg); } - // Profiled runs - context->setProfiler(&profiler); - for (int i = 0; i < 20; ++i) { - auto start = std::chrono::system_clock::now(); - auto prob = doInference(*context, input.data(), N); - auto end = std::chrono::system_clock::now(); - auto period = std::chrono::duration_cast(end - start); - std::cout << period.count() << "us\n"; - - for (const auto& vector : prob) { - int idx = 0; - for (auto v : vector) { - std::cout << std::setprecision(4) << v << ", " << std::flush; - if (++idx > 20) { - std::cout << "\n====\n"; - break; - } + auto firstProb = doInference(*context, input_ptr, infer_batch, cfg); + auto labels = (cfg.num_classes == 1000) ? loadLabels() : std::map{}; + for (int64_t b = 0; b < infer_batch; ++b) { + std::cout << "[sample " << b << "] " << images[b] << "\n"; + std::vector sample(firstProb[0].begin() + b * cfg.num_classes, + firstProb[0].begin() + (b + 1) * cfg.num_classes); + printFirstOutputs("vit", sample.data(), sample.size()); + int _top = 0; + for (auto& [idx, logits] : topk(sample, 3)) { + std::cout << " Top: " << _top++ << " idx: " << idx << ", logits: " << logits; + if (!labels.empty()) { + std::cout << ", label: " << labels[idx]; } + std::cout << "\n"; } + } - if (i == 19) { - std::cout << "prediction result: \n"; - auto labels = loadImagenetLabelMap(LABELS_PATH); - int _top = 0; - for (auto& [idx, logits] : topk(prob[0], 3)) { - std::cout << "Top: " << _top++ << " idx: " << idx << ", logits: " << logits - << ", label: " << labels[idx] << "\n"; - } - std::cout << profiler << "\n"; - } + std::vector latencies; + latencies.reserve(kBenchmarkRuns); + for (int i = 0; i < kBenchmarkRuns; ++i) { + auto start = std::chrono::steady_clock::now(); + (void)doInference(*context, input_ptr, infer_batch, cfg); + auto end = std::chrono::steady_clock::now(); + auto period = std::chrono::duration_cast(end - start); + latencies.push_back(static_cast(period.count()) / 1000.0); + } + printBenchmark("vit", latencies, infer_batch); + if (enable_profile) { + std::cout << profiler; } return 0; }