wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 25 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/sampleDynamicReshape/CMakeLists.txt
+25
View File
@@ -0,0 +1,25 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
add_executable(sample_dynamic_reshape sampleDynamicReshape.cpp)
target_link_libraries(sample_dynamic_reshape PRIVATE trt_samples_common TRT_SAMPLES::tensorrt)
add_dependencies(tensorrt_samples sample_dynamic_reshape)
installLibraries(
TARGETS sample_dynamic_reshape
OPTIONAL
COMPONENT internal
)
samples/sampleDynamicReshape/README.md
+280
View File
@@ -0,0 +1,280 @@
# Digit Recognition With Dynamic Shapes In TensorRT
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
* [Creating the preprocessing network](#creating-the-preprocessing-network)
* [Parsing the ONNX MNIST model](#parsing-the-onnx-mnist-model)
* [Building engines](#building-engines)
* [Running inference](#running-inference)
* [TensorRT API layers and ops](#tensorrt-api-layers-and-ops)
- [Preparing sample data](#preparing-sample-data)
- [Running the sample](#running-the-sample)
* [Sample `--help` options](#sample-help-options)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
This sample, sampleDynamicReshape, demonstrates how to use dynamic input dimensions in TensorRT. It creates an engine that takes a dynamically shaped input and resizes it to be consumed by an ONNX MNIST model that expects a fixed size input. For more information, see [Working With Dynamic Shapes](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#work_dynamic_shapes) in the TensorRT Developer Guide.
## How does this sample work?
This sample creates an engine for resizing an input with dynamic dimensions to a size that an ONNX MNIST model can consume.
Specifically, this sample:
- Creates a network with dynamic input dimensions to act as a preprocessor for the model
- Parses an ONNX MNIST model to create a second network
- Builds engines for both networks
- Runs inference using both engines
### Creating the preprocessing network
First, create a network with full dims support:
`auto preprocessorNetwork = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));`
Next, add an input layer that accepts an input with a dynamic shape, followed by a resize layer that will reshape the input to the shape the model expects:
```
auto input = preprocessorNetwork->addInput("input", nvinfer1::DataType::kFLOAT, Dims4{-1, 1, -1, -1});
auto resizeLayer = preprocessorNetwork->addResize(*input);
resizeLayer->setOutputDimensions(mPredictionInputDims);
preprocessorNetwork->markOutput(*resizeLayer->getOutput(0));
```
The -1 dimensions denote dimensions that will be supplied at runtime.
### Parsing the ONNX MNIST model
First, create an empty full-dims network, and parser:
```
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
auto parser = nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger());
```
Next, parse the model file to populate the network:
```
parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(), static_cast<int>(sample::gLogger.getReportableSeverity()));
```
### Building engines
When building the preprocessor engine, also provide an optimization profile so that TensorRT knows which input shapes to optimize for:
```
auto preprocessorConfig = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
auto profile = builder->createOptimizationProfile();
```
`OptProfileSelector::kOPT` specifies the dimensions that the profile will be optimized for, whereas `OptProfileSelector::kMIN` and `OptProfileSelector::kMAX` specify the minimum and maximum dimensions for which the profile will be valid:
```
profile->setDimensions(input->getName(), OptProfileSelector::kMIN, Dims4{1, 1, 1, 1});
profile->setDimensions(input->getName(), OptProfileSelector::kOPT, Dims4{1, 1, 28, 28});
profile->setDimensions(input->getName(), OptProfileSelector::kMAX, Dims4{1, 1, 56, 56});
preprocessorConfig->addOptimizationProfile(profile);
```
Run engine build with config:
```
auto preprocessorPlan = std::unique_ptr<nvinfer1::IHostMemory>(
builder->buildSerializedNetwork(*preprocessorNetwork, *preprocessorConfig));
if (!preprocessorPlan)
{
sample::gLogError << "Preprocessor serialized engine build failed." << std::endl;
return false;
}
mPreprocessorEngine = std::unique_ptr<nvinfer1::ICudaEngine>(
runtime->deserializeCudaEngine(preprocessorPlan->data(), preprocessorPlan->size()));
if (!mPreprocessorEngine)
{
sample::gLogError << "Preprocessor engine deserialization failed." << std::endl;
return false;
}
```
For the MNIST model, attach a Softmax layer to the end of the network, set softmax axis to 1 since network output has shape [1, 10] in full dims mode and replace the existing network output with the Softmax:
```
auto softmax = network->addSoftMax(*network->getOutput(0));
softmax->setAxes(1 << 1);
network->unmarkOutput(*network->getOutput(0));
network->markOutput(*softmax->getOutput(0));
```
Finally, build as normal:
```
auto predictionPlan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
if (!predictionPlan)
{
sample::gLogError << "Prediction serialized engine build failed." << std::endl;
return false;
}
mPredictionEngine = std::unique_ptr<nvinfer1::ICudaEngine>(
runtime->deserializeCudaEngine(predictionPlan->data(), predictionPlan->size()));
if (!mPredictionEngine)
{
sample::gLogError << "Prediction engine deserialization failed." << std::endl;
return false;
}
```
### Running inference
During inference, first copy the input buffer to the device:
```
CHECK(cudaMemcpy(mInput.deviceBuffer.data(), mInput.hostBuffer.data(), mInput.hostBuffer.nbBytes(), cudaMemcpyHostToDevice));
```
Since the preprocessor engine accepts dynamic shapes, specify the actual shape of the current input to the execution context:
`mPreprocessorContext->setInputShape(inputTensorName, inputDims);`, where inputTensorName is the name of the input tensor on binding index 0.
Next, run the preprocessor using the `executeV2` function. The example writes the output of the preprocessor engine directly to the input device buffer of the MNIST engine:
```
std::vector<void*> preprocessorBindings = {mInput.deviceBuffer.data(), mPredictionInput.data()};
bool status = mPreprocessorContext->executeV2(preprocessorBindings.data());
```
Then, run the MNIST engine:
```
std::vector<void*> predicitonBindings = {mPredictionInput.data(), mOutput.deviceBuffer.data()};
status = mPredictionContext->executeV2(predicitonBindings.data());
```
Finally, copy the output back to the host:
```
CHECK(cudaMemcpy(mOutput.hostBuffer.data(), mOutput.deviceBuffer.data(), mOutput.deviceBuffer.nbBytes(), cudaMemcpyDeviceToHost));
```
### TensorRT API layers and ops
In this sample, the following layers are used. For more information about these layers, see the [TensorRT Developer Guide: Layers](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#layers) documentation.
[Resize layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#resize-layer)
The IResizeLayer implements the resize operation on an input tensor.
## Prerequisites
1. See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README.
## Running the sample
1. Compile the sample by following build instructions in [TensorRT README](https://github.com/NVIDIA/TensorRT/).
2. Run the sample.
```bash
./sample_dynamic_reshape [-h or --help] [-d or --datadir=<path to data directory>] [--useDLACore=<int>]
```
For example:
```bash
./sample_dynamic_reshape --datadir $TRT_DATADIR/mnist
```
3. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
```
&&&& RUNNING TensorRT.sample_dynamic_reshape # ./sample_dynamic_reshape
----------------------------------------------------------------
Input filename: ../../../../../data/samples/mnist/mnist.onnx
ONNX IR version: 0.0.3
Opset version: 8
Producer name: CNTK
Producer version: 2.5.1
Domain: ai.cntk
Model version: 1
Doc string:
----------------------------------------------------------------
[W] [TRT] onnx2trt_utils.cpp:214: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[W] [TRT] onnx2trt_utils.cpp:214: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[I] [TRT] Detected 1 inputs and 1 output network tensors.
[I] [TRT] Detected 1 inputs and 1 output network tensors.
[I] Profile dimensions in preprocessor engine:
[I] Minimum = (1, 1, 1, 1)
[I] Optimum = (1, 1, 28, 28)
[I] Maximum = (1, 1, 56, 56)
[I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@*. .*@@@@@@@@@@@
@@@@@@@@@@*. +@@@@@@@@@@
@@@@@@@@@@. :#+ %@@@@@@@@@
@@@@@@@@@@.:@@@+ +@@@@@@@@@
@@@@@@@@@@.:@@@@: +@@@@@@@@@
@@@@@@@@@@=%@@@@: +@@@@@@@@@
@@@@@@@@@@@@@@@@# +@@@@@@@@@
@@@@@@@@@@@@@@@@* +@@@@@@@@@
@@@@@@@@@@@@@@@@: +@@@@@@@@@
@@@@@@@@@@@@@@@@: +@@@@@@@@@
@@@@@@@@@@@@@@@* .@@@@@@@@@@
@@@@@@@@@@%**%@. *@@@@@@@@@@
@@@@@@@@%+. .: .@@@@@@@@@@@
@@@@@@@@= .. :@@@@@@@@@@@
@@@@@@@@: *@@: :@@@@@@@@@@@
@@@@@@@% %@* *@@@@@@@@@@
@@@@@@@% ++ ++ .%@@@@@@@@@
@@@@@@@@- +@@- +@@@@@@@@@
@@@@@@@@= :*@@@# .%@@@@@@@@
@@@@@@@@@+*@@@@@%. %@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
[I] Output:
[I] Prob 0 0.0000 Class 0:
[I] Prob 1 0.0000 Class 1:
[I] Prob 2 1.0000 Class 2: **********
[I] Prob 3 0.0000 Class 3:
[I] Prob 4 0.0000 Class 4:
[I] Prob 5 0.0000 Class 5:
[I] Prob 6 0.0000 Class 6:
[I] Prob 7 0.0000 Class 7:
[I] Prob 8 0.0000 Class 8:
[I] Prob 9 0.0000 Class 9:
&&&& PASSED TensorRT.sample_dynamic_reshape # ./sample_dynamic_reshape
```
This output shows that the sample ran successfully; `PASSED`.
### Sample `--help` options
To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
# Additional resources
The following resources provide a deeper understanding of dynamic shapes.
**ONNX**
- [GitHub: ONNX](https://github.com/onnx/onnx)
- [GitHub: ONNX-TensorRT open source parser](https://github.com/onnx/onnx-tensorrt)
**Models**
- [MNIST - Handwritten Digit Recognition](https://github.com/onnx/models/tree/main/validated/vision/classification/mnist)
- [GitHub: ONNX Models](https://github.com/onnx/models)
**Documentation**
- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
# License
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
# Changelog
October 2025
Migrate to strongly typed APIs.
February 2020
This is the second release of the `README.md` file and sample.
# Known issues
There are no known issues in this sample.
samples/sampleDynamicReshape/sampleDynamicReshape.cpp
+540
View File
@@ -0,0 +1,540 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
//!
//! sampleDynamicReshape.cpp
//! This file contains the implementation of the dynamic reshape MNIST sample. It creates a network
//! using the MNIST ONNX model, and uses a second engine to resize inputs to the shape the model
//! expects.
//! It can be run with the following command:
//! Command: ./sample_dynamic_reshape [-h or --help [-d=/path/to/data/dir or --datadir=/path/to/data/dir]
//!
// Define TRT entrypoints used in common code
#define DEFINE_TRT_ENTRYPOINTS 1
#include "BatchStream.h"
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "parserOnnxConfig.h"
#include "NvInfer.h"
#include <cuda_runtime_api.h>
#include <random>
using namespace nvinfer1;
const std::string gSampleName = "TensorRT.sample_dynamic_reshape";
//! \brief The SampleDynamicReshape class implements the dynamic reshape sample.
//!
//! \details This class builds one engine that resizes a given input to the correct size, and a
//! second engine based on an ONNX MNIST model that generates a prediction.
//!
class SampleDynamicReshape
{
public:
SampleDynamicReshape(const samplesCommon::OnnxSampleParams& params)
: mParams(params)
{
}
//!
//! \brief Builds both engines.
//!
bool build();
//!
//! \brief Prepares the model for inference by creating execution contexts and allocating buffers.
//!
bool prepare();
//!
//! \brief Runs inference using TensorRT on a random image.
//!
bool infer();
private:
[[nodiscard]] bool buildPreprocessorEngine(
nvinfer1::IBuilder& builder, nvinfer1::IRuntime& runtime, cudaStream_t profileStream);
[[nodiscard]] bool buildPredictionEngine(
nvinfer1::IBuilder& builder, nvinfer1::IRuntime& runtime, cudaStream_t profileStream);
[[nodiscard]] Dims loadPGMFile(const std::string& fileName);
[[nodiscard]] bool validateOutput(int digit);
samplesCommon::OnnxSampleParams mParams; //!< The parameters for the sample.
nvinfer1::Dims mPredictionInputDims; //!< The dimensions of the input of the MNIST model.
nvinfer1::Dims mPredictionOutputDims; //!< The dimensions of the output of the MNIST model.
std::unique_ptr<nvinfer1::IRuntime> mRuntime{nullptr};
// Engine plan files used for inference. One for resizing inputs, another for prediction.
std::unique_ptr<nvinfer1::ICudaEngine> mPreprocessorEngine{nullptr}, mPredictionEngine{nullptr};
std::unique_ptr<nvinfer1::IExecutionContext> mPreprocessorContext{nullptr}, mPredictionContext{nullptr};
samplesCommon::ManagedBuffer mInput{}; //!< Host and device buffers for the input.
samplesCommon::DeviceBuffer mPredictionInput{}; //!< Device buffer for the output of the preprocessor, i.e. the
//!< input to the prediction model.
samplesCommon::ManagedBuffer mOutput{}; //!< Host buffer for the output
};
//!
//! \brief Builds the two engines required for inference.
//!
//! \details This function creates one TensorRT engine for resizing inputs to the correct sizes,
//! then creates a TensorRT network by parsing the ONNX model and builds
//! an engine that will be used to run inference (mPredictionEngine).
//!
//! \return false if error in build preprocessor or predict engine.
//!
bool SampleDynamicReshape::build()
{
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
if (!builder)
{
sample::gLogError << "Create inference builder failed." << std::endl;
return false;
}
mRuntime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
if (!mRuntime)
{
sample::gLogError << "Runtime object creation failed." << std::endl;
return false;
}
// This function will also set mPredictionInputDims and mPredictionOutputDims,
// so it needs to be called before building the preprocessor.
try
{
// CUDA stream used for profiling by the builder.
auto profileStream = samplesCommon::makeCudaStream();
if (!profileStream)
{
return false;
}
bool result = buildPredictionEngine(*builder, *mRuntime, *profileStream)
&& buildPreprocessorEngine(*builder, *mRuntime, *profileStream);
return result;
}
catch (std::runtime_error& e)
{
sample::gLogError << e.what() << std::endl;
return false;
}
}
//!
//! \brief Builds an engine for preprocessing (mPreprocessorEngine).
//!
//! \return false if error in build preprocessor engine.
//!
bool SampleDynamicReshape::buildPreprocessorEngine(
nvinfer1::IBuilder& builder, nvinfer1::IRuntime& runtime, cudaStream_t profileStream)
{
// Create the preprocessor engine using a network that supports full dimensions (createNetworkV2).
auto preprocessorNetwork = std::unique_ptr<INetworkDefinition>(
builder.createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
if (!preprocessorNetwork)
{
sample::gLogError << "Create network failed." << std::endl;
return false;
}
// Reshape a dynamically shaped input to the size expected by the model, (1, 1, 28, 28).
auto input = preprocessorNetwork->addInput("input", nvinfer1::DataType::kFLOAT, Dims4{-1, 1, -1, -1});
auto resizeLayer = preprocessorNetwork->addResize(*input);
resizeLayer->setOutputDimensions(mPredictionInputDims);
preprocessorNetwork->markOutput(*resizeLayer->getOutput(0));
// Finally, configure and build the preprocessor engine.
auto preprocessorConfig = std::unique_ptr<IBuilderConfig>{builder.createBuilderConfig()};
if (!preprocessorConfig)
{
sample::gLogError << "Create builder config failed." << std::endl;
return false;
}
// Create an optimization profile so that we can specify a range of input dimensions.
auto profile = builder.createOptimizationProfile();
// This profile will be valid for all images whose size falls in the range of [(1, 1, 1, 1), (1, 1, 56, 56)]
// but TensorRT will optimize for (1, 1, 28, 28)
// We do not need to check the return of setDimension and addOptimizationProfile here as all dims are explicitly set
profile->setDimensions(input->getName(), OptProfileSelector::kMIN, Dims4{1, 1, 1, 1});
profile->setDimensions(input->getName(), OptProfileSelector::kOPT, Dims4{1, 1, 28, 28});
profile->setDimensions(input->getName(), OptProfileSelector::kMAX, Dims4{1, 1, 56, 56});
preprocessorConfig->addOptimizationProfile(profile);
std::unique_ptr<nvinfer1::ITimingCache> timingCache{};
// Load timing cache
if (!mParams.timingCacheFile.empty())
{
timingCache = samplesCommon::buildTimingCacheFromFile(
sample::gLogger.getTRTLogger(), *preprocessorConfig, mParams.timingCacheFile);
}
auto preprocessorPlan = std::unique_ptr<nvinfer1::IHostMemory>(
builder.buildSerializedNetwork(*preprocessorNetwork, *preprocessorConfig));
if (!preprocessorPlan)
{
sample::gLogError << "Preprocessor serialized engine build failed." << std::endl;
return false;
}
if (timingCache != nullptr && !mParams.timingCacheFile.empty())
{
samplesCommon::updateTimingCacheFile(
sample::gLogger.getTRTLogger(), mParams.timingCacheFile, timingCache.get(), builder);
}
mPreprocessorEngine = std::unique_ptr<nvinfer1::ICudaEngine>(
runtime.deserializeCudaEngine(preprocessorPlan->data(), preprocessorPlan->size()));
if (!mPreprocessorEngine)
{
sample::gLogError << "Preprocessor engine deserialization failed." << std::endl;
return false;
}
auto const tensorName = mPreprocessorEngine->getIOTensorName(0);
sample::gLogInfo << "Profile dimensions in preprocessor engine:" << std::endl;
sample::gLogInfo << " Minimum = " << mPreprocessorEngine->getProfileShape(tensorName, 0, OptProfileSelector::kMIN)
<< std::endl;
sample::gLogInfo << " Optimum = " << mPreprocessorEngine->getProfileShape(tensorName, 0, OptProfileSelector::kOPT)
<< std::endl;
sample::gLogInfo << " Maximum = " << mPreprocessorEngine->getProfileShape(tensorName, 0, OptProfileSelector::kMAX)
<< std::endl;
return true;
}
//!
//! \brief Builds an engine for prediction (mPredictionEngine).
//!
//! \details This function builds an engine for the MNIST model, and updates mPredictionInputDims and
//! mPredictionOutputDims according to the dimensions specified by the model. The preprocessor reshapes inputs to
//! mPredictionInputDims.
//!
//! \return false if error in build prediction engine.
//!
bool SampleDynamicReshape::buildPredictionEngine(
nvinfer1::IBuilder& builder, nvinfer1::IRuntime& runtime, cudaStream_t profileStream)
{
// Create a network using the parser.
auto network = std::unique_ptr<INetworkDefinition>(
builder.createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
if (!network)
{
sample::gLogError << "Create network failed." << std::endl;
return false;
}
auto const parser
= std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
if (parser == nullptr)
{
throw std::runtime_error("Failed to create ONNX parser");
}
bool parsingSuccess
= parser->parseFromFile(samplesCommon::locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
static_cast<int>(sample::gLogger.getReportableSeverity()));
if (!parsingSuccess)
{
sample::gLogError << "Failed to parse model." << std::endl;
return false;
}
// Attach a softmax layer to the end of the network.
auto softmax = network->addSoftMax(*network->getOutput(0));
// Set softmax axis to 1 since network output has shape [1, 10] in full dims mode
softmax->setAxes(1 << 1);
network->unmarkOutput(*network->getOutput(0));
network->markOutput(*softmax->getOutput(0));
// Get information about the inputs/outputs directly from the model.
mPredictionInputDims = network->getInput(0)->getDimensions();
mPredictionOutputDims = network->getOutput(0)->getDimensions();
// Create a builder config
auto config = std::unique_ptr<IBuilderConfig>(builder.createBuilderConfig());
if (!config)
{
sample::gLogError << "Create builder config failed." << std::endl;
return false;
}
config->setProfileStream(profileStream);
// Build the prediction engine.
std::unique_ptr<nvinfer1::ITimingCache> timingCache{};
// Load timing cache
if (!mParams.timingCacheFile.empty())
{
timingCache
= samplesCommon::buildTimingCacheFromFile(sample::gLogger.getTRTLogger(), *config, mParams.timingCacheFile);
}
// Build the prediction engine.
auto predictionPlan = std::unique_ptr<nvinfer1::IHostMemory>(builder.buildSerializedNetwork(*network, *config));
if (!predictionPlan)
{
sample::gLogError << "Prediction serialized engine build failed." << std::endl;
return false;
}
if (timingCache != nullptr && !mParams.timingCacheFile.empty())
{
samplesCommon::updateTimingCacheFile(
sample::gLogger.getTRTLogger(), mParams.timingCacheFile, timingCache.get(), builder);
}
mPredictionEngine = std::unique_ptr<nvinfer1::ICudaEngine>(
runtime.deserializeCudaEngine(predictionPlan->data(), predictionPlan->size()));
if (!mPredictionEngine)
{
sample::gLogError << "Prediction engine deserialization failed." << std::endl;
return false;
}
return true;
}
//!
//! \brief Prepares the model for inference by creating an execution context and allocating buffers.
//!
//! \details This function sets up the sample for inference. This involves allocating buffers for the inputs and
//! outputs, as well as creating TensorRT execution contexts for both engines. This only needs to be called a single
//! time.
//!
//! \return false if error in build preprocessor or predict context.
//!
bool SampleDynamicReshape::prepare()
{
mPreprocessorContext = std::unique_ptr<IExecutionContext>(mPreprocessorEngine->createExecutionContext());
if (!mPreprocessorContext)
{
sample::gLogError << "Preprocessor context build failed." << std::endl;
return false;
}
mPredictionContext = std::unique_ptr<IExecutionContext>(mPredictionEngine->createExecutionContext());
if (!mPredictionContext)
{
sample::gLogError << "Prediction context build failed." << std::endl;
return false;
}
// Since input dimensions are not known ahead of time, we only allocate the output buffer and preprocessor output
// buffer.
mPredictionInput.resize(mPredictionInputDims);
mOutput.hostBuffer.resize(mPredictionOutputDims);
mOutput.deviceBuffer.resize(mPredictionOutputDims);
return true;
}
//!
//! \brief Runs inference for this sample
//!
//! \details This function is the main execution function of the sample.
//! It runs inference for using a random image from the MNIST dataset as an input.
//!
bool SampleDynamicReshape::infer()
{
// Load a random PGM file into a host buffer, then copy to device.
std::random_device rd{};
std::default_random_engine generator{rd()};
std::uniform_int_distribution<int> digitDistribution{0, 9};
int digit = digitDistribution(generator);
Dims inputDims = loadPGMFile(samplesCommon::locateFile(std::to_string(digit) + ".pgm", mParams.dataDirs));
mInput.deviceBuffer.resize(inputDims);
CHECK(cudaMemcpy(
mInput.deviceBuffer.data(), mInput.hostBuffer.data(), mInput.hostBuffer.nbBytes(), cudaMemcpyHostToDevice));
// Set the input size for the preprocessor
CHECK_RETURN_W_MSG(mPreprocessorContext->setInputShape("input", inputDims), false, "Invalid binding dimensions.");
// We can only run inference once all dynamic input shapes have been specified.
if (!mPreprocessorContext->allInputDimensionsSpecified())
{
return false;
}
// Run the preprocessor to resize the input to the correct shape
std::vector<void*> preprocessorBindings = {mInput.deviceBuffer.data(), mPredictionInput.data()};
// For engines using full dims, we can use executeV2, which does not include a separate batch size parameter.
bool status = mPreprocessorContext->executeV2(preprocessorBindings.data());
if (!status)
{
return false;
}
// Next, run the model to generate a prediction.
std::vector<void*> predicitonBindings = {mPredictionInput.data(), mOutput.deviceBuffer.data()};
status = mPredictionContext->executeV2(predicitonBindings.data());
if (!status)
{
return false;
}
// Copy the outputs back to the host and verify the output.
CHECK(cudaMemcpy(mOutput.hostBuffer.data(), mOutput.deviceBuffer.data(), mOutput.deviceBuffer.nbBytes(),
cudaMemcpyDeviceToHost));
return validateOutput(digit);
}
//!
//! \brief Loads a PGM file into mInput and returns the dimensions of the loaded image.
//!
//! \details This function loads the specified PGM file into the input host buffer.
//!
Dims SampleDynamicReshape::loadPGMFile(const std::string& fileName)
{
std::ifstream infile(fileName, std::ifstream::binary);
ASSERT(infile.is_open() && "Attempting to read from a file that is not open.");
std::string magic;
int h, w, max;
infile >> magic >> h >> w >> max;
infile.seekg(1, infile.cur);
Dims4 inputDims{1, 1, h, w};
size_t vol = samplesCommon::volume(inputDims);
std::vector<uint8_t> fileData(vol);
infile.read(reinterpret_cast<char*>(fileData.data()), vol);
// Print an ascii representation
sample::gLogInfo << "Input:\n";
for (size_t i = 0; i < vol; i++)
{
sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % w) ? "" : "\n");
}
sample::gLogInfo << std::endl;
// Normalize and copy to the host buffer.
mInput.hostBuffer.resize(inputDims);
float* hostDataBuffer = static_cast<float*>(mInput.hostBuffer.data());
std::transform(fileData.begin(), fileData.end(), hostDataBuffer,
[](uint8_t x) { return 1.0 - static_cast<float>(x / 255.0); });
return inputDims;
}
//!
//! \brief Checks whether the model prediction (in mOutput) is correct.
//!
bool SampleDynamicReshape::validateOutput(int digit)
{
const float* bufRaw = static_cast<const float*>(mOutput.hostBuffer.data());
std::vector<float> prob(bufRaw, bufRaw + mOutput.hostBuffer.size());
int curIndex{0};
for (const auto& elem : prob)
{
sample::gLogInfo << " Prob " << curIndex << " " << std::fixed << std::setw(5) << std::setprecision(4) << elem
<< " "
<< "Class " << curIndex << ": " << std::string(int(std::floor(elem * 10 + 0.5F)), '*')
<< std::endl;
++curIndex;
}
int predictedDigit = std::max_element(prob.begin(), prob.end()) - prob.begin();
return digit == predictedDigit;
}
//!
//! \brief Initializes members of the params struct using the command line args
//!
samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
{
samplesCommon::OnnxSampleParams params;
if (args.dataDirs.empty()) // Use default directories if user hasn't provided directory paths
{
params.dataDirs.push_back("data/mnist/");
params.dataDirs.push_back("data/samples/mnist/");
}
else // Use the data directory provided by the user
{
params.dataDirs = args.dataDirs;
}
params.onnxFileName = "mnist.onnx";
params.inputTensorNames.push_back("Input3");
params.outputTensorNames.push_back("Plus214_Output_0");
params.timingCacheFile = args.timingCacheFile;
return params;
}
//!
//! \brief Prints the help information for running this sample
//!
void printHelpInfo()
{
std::cout << "Usage: ./sample_dynamic_reshape [-h or --help] [-d or --datadir=<path to data directory>] "
"[--timingCacheFile=<path to timing cache file>]"
<< std::endl;
std::cout << "--help, -h Display help information" << std::endl;
std::cout << "--datadir Specify path to a data directory, overriding the default. This option can be used "
"multiple times to add multiple directories. If no data directories are given, the default is to use "
"(data/samples/mnist/, data/mnist/)"
<< std::endl;
std::cout << "--timingCacheFile Specify path to a timing cache file. If it does not already exist, it will be "
<< "created." << std::endl;
}
int main(int argc, char** argv)
{
samplesCommon::Args args;
bool argsOK = samplesCommon::parseArgs(args, argc, argv);
if (!argsOK)
{
sample::gLogError << "Invalid arguments" << std::endl;
printHelpInfo();
return EXIT_FAILURE;
}
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
auto sampleTest = sample::gLogger.defineTest(gSampleName, argc, argv);
sample::gLogger.reportTestStart(sampleTest);
SampleDynamicReshape sample{initializeSampleParams(args)};
if (!sample.build())
{
return sample::gLogger.reportFail(sampleTest);
}
if (!sample.prepare())
{
return sample::gLogger.reportFail(sampleTest);
}
if (!sample.infer())
{
return sample::gLogger.reportFail(sampleTest);
}
return sample::gLogger.reportPass(sampleTest);
}