wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 30 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/sampleSafeMNIST/CMakeLists.txt
+78
View File
@@ -0,0 +1,78 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2025-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.
#
if(NOT TRT_SAFETY_INFERENCE_ONLY)
add_executable(sample_mnist_safe_build sampleSafeMNISTBuild.cpp)
target_link_libraries(sample_mnist_safe_build PRIVATE
trt_samples_common
TRT_SAMPLES::tensorrt
TRTSAFE::nvinfer_safe_shared
)
# Link ONNX parser if available
if(TARGET nvonnxparser)
target_link_libraries(sample_mnist_safe_build PRIVATE nvonnxparser)
endif()
if(TRT_OUT_DIR)
set_target_properties(sample_mnist_safe_build
PROPERTIES
RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
)
endif()
add_dependencies(tensorrt_samples sample_mnist_safe_build)
installLibraries(
TARGETS sample_mnist_safe_build
OPTIONAL
COMPONENT internal
)
endif()
add_executable(sample_mnist_safe_infer sampleSafeMNISTInfer.cpp)
if(TRT_SAFETY_INFERENCE_ONLY)
target_link_libraries(sample_mnist_safe_infer PRIVATE trt_global_definitions tensorrt_headers)
target_include_directories(sample_mnist_safe_infer PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/../common
)
# Explicitly link CTK libraries
if(CUDAToolkit_LIBRARY_DIR)
target_link_directories(sample_mnist_safe_infer PRIVATE
${CUDAToolkit_LIBRARY_DIR}
)
endif()
else()
target_link_libraries(sample_mnist_safe_infer PRIVATE
trt_samples_common
TRTSAFE::nvinfer_safe_shared
)
endif()
if(TRT_OUT_DIR)
set_target_properties(sample_mnist_safe_infer
PROPERTIES
RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
)
endif()
add_dependencies(tensorrt_samples sample_mnist_safe_infer)
installLibraries(
TARGETS sample_mnist_safe_infer
OPTIONAL
COMPONENT internal
)
samples/sampleSafeMNIST/README.md
+221
View File
@@ -0,0 +1,221 @@
# “Hello World” For TensorRT Safety
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
* [TensorRT API layers and ops](#tensorrt-api-layers-and-ops)
- [Running the sample](#running-the-sample)
* [Tool command line arguments](#tool-command-line-arguments)
* [When to use remoteAutoTuningConfig](#when-to-use-remoteautotuningconfig)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
This sample, sampleSafeMNIST, consists of two parts; build and infer. The build part of this sample demonstrates how to use the builder `IBuilderConfig::setEngineCapability()` flag for safety. The inference part of this sample demonstrates how to use the safe graph.
The build part builds a safe version of a TensorRT engine and saves it into a binary file, then the infer part loads the prebuilt safe engine and performs inference on an input image.
## How does this sample work?
This sample uses an ONNX model that was trained on the [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md).
Specifically, this sample:
- Build (sample_mnist_safe_build):
- Performs the basic setup and initialization of TensorRT
- [Imports a trained ONNX model using ONNX parser](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/c-api-docs.html#importing-a-model-using-the-onnx-parser)
- Preprocesses the input and stores the result in a managed buffer
- [Builds a safe engine](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/c-api-docs.html#building-an-engine)
- Infer (sample_mnist_safe_infer):
- Create a safe graph for setting up tensors and executing inference on a built network.
To verify whether the engine is operating correctly, this sample picks a 28x28 image of a digit at random and runs inference on it using the engine it created. The output of the network is a probability distribution on the digit, showing which digit is likely that in the image.
### TensorRT API layers and ops
In this sample, the following layers are used. For more information about these layers, see the [TensorRT API: Layers](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html) documentation.
[Activation layer](https://docs.nvidia.com/deeplearning/tensorrt/operators/docs/Activation.html)
The Activation layer implements element-wise activation functions. Specifically, this sample uses the Activation layer with the type `kRELU`.
[Convolution layer](https://docs.nvidia.com/deeplearning/tensorrt/operators/docs/Convolution.html)
The Convolution layer computes a 2D (channel, height, and width) convolution, with or without bias.
## Running the sample
1. Download the [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md) to read images from the ubyte file. The images need to be saved into `.pgm` format and renamed as `<label>.pgm`.
2. Put the images into the `data/mnist` directory together with the existing ONNX network `safe_mnist.onnx`.
3. Compile the sample by following the build instructions in the [TensorRT README](https://github.com/NVIDIA/TensorRT/). This will build the sample binaries, including `sample_mnist_safe_build` and `sample_mnist_safe_infer`.
4. The compile options are summarized in the following table.
| Compile Option | Default |Description|
| ------------------------------- | ------- |---------- |
|TRT_SAFETY_INFERENCE_ONLY | OFF |When enabled, build the infer part only, skip compiling the builder part.|
5. Run the sample to build a TensorRT safe engine.
```
./sample_mnist_safe_build [--datadir=/path/to/data/dir/] [--remoteAutoTuningConfig=<config>] [--cpuOnly]
```
This sample generates `safe_mnist.engine`, which is a binary file that contains the serialized engine data.
This sample reads ONNX model to build the network:
- `safe_mnist.onnx` - The ONNX model that contains the network design.
**Note:** By default, this sample expects these files to be in either the `data/samples/mnist/` or `data/mnist/` directories. The list of default directories can be changed by adding one or more paths with `--datadir=/new/path/` as a command line argument.
6. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
```
&&&& RUNNING TensorRT.sample_mnist_safe_build # ./sample_mnist_safe_build
[I] Building a GPU inference engine for MNIST
[I] [TRT] Detected 1 input and 1 output network tensors.
&&&& PASSED TensorRT.sample_mnist_safe_build # ./sample_mnist_safe_build
```
This output shows that the sample ran successfully; `PASSED`.
7. Run the sample to perform inference on the digit:
`./sample_mnist_safe_infer`
**Note:** This sample expects `./sample_mnist_safe_build` has been run to generate a safe engine file. It loads input image from `data/samples/mnist` directory, and walks back 10 directories to locate the image.
8. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following; ASCII rendering of the input image with digit 3:
```
&&&& RUNNING TensorRT.sample_mnist_safe_infer # ./sample_mnist_safe_infer
[I] Running a GPU inference engine for MNIST
[I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@#-:.-=@@@@@@@@@@@@@@
@@@@@%= . *@@@@@@@@@@@@@@@@@
@@@@% .:+%%% *@@@@@@@@@@@@@@
@@@@+=#@@@@@# @@@@@@@@@@@@@@
@@@@@@@@@@@% @@@@@@@@@@@@@@@
@@@@@@@@@@@: *@@@@@@@@@@@@@@
@@@@@@@@@@- .@@@@@@@@@@@@@@@
@@@@@@@@@: #@@@@@@@@@@@@@@@@
@@@@@@@@: +*%#@@@@@@@@@@@@@@
@@@@@@@% :+*@@@@@@@@@@@@@@@@
@@@@@@@@#*+--.:: +@@@@@@@@@@
@@@@@@@@@@@@@@@@#=:. +@@@@@@
@@@@@@@@@@@@@@@@@@@@ .@@@@@@
@@@@@@@@@@@@@@@@@@@@#. #@@@@
@@@@@@@@@@@@@@@@@@@@# @@@@@@
@@@@@@@@@%@@@@@@@@@@- +@@@@@
@@@@@@@@#-@@@@@@@@*. =@@@@@@
@@@@@@@@ .+%%%%+=. =@@@@@@@@
@@@@@@@@ =@@@@@@@@@@@@@@@@@@
@@@@@@@@*=: :--*@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
[I] Output:
[I] Prob 0 0.0000 Class 0:
[I] Prob 1 0.0000 Class 1:
[I] Prob 2 0.0000 Class 2:
[I] Prob 3 1.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_safe_mnist_infer # ./sample_mnist_safe_infer
```
This output shows that the sample ran successfully; `PASSED`.
### Tool command line arguments
To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
```bash
sample_mnist_safe_build --help
sample_mnist_safe_infer --help
```
### When to use remoteAutoTuningConfig
The `--remoteAutoTuningConfig` parameter is designed for **cross-platform development scenarios** where you need to:
**Primary Use Case - Cross-Platform Building:**
- **Build on Host Platform**: Compile and build TensorRT engines on a development machine (e.g., Linux x86_64)
- **Auto-tune on Target Platform**: Perform kernel auto-tuning on the actual deployment target (e.g., QNX aarch64)
Use `--cpuOnly` with `--remoteAutoTuningConfig` to build the engine without a local GPU on the build host:
```bash
./sample_mnist_safe_build --remoteAutoTuningConfig=<config> --cpuOnly
```
**Typical Scenarios:**
- **QNX Development**: Building engines on Linux development machines but deploying on QNX automotive platforms
**Important Technical Limitation:**
- **QNX Safety Devices**: QNX safety platforms do **NOT** support engine building operations. All engine construction must be performed on development platforms (Linux/QNX standard), making remote auto-tuning essential for safety deployments.
## Additional resources
The following resources provide a deeper understanding about sampleSafeMNIST.
**Dataset**
- [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md)
**Documentation**
- [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
Jun. 2019
This is the first release of the `README.md` file and sample.
Dec. 2019
Switch the sample to use ONNX model, and update the content of `README.md`.
Jun. 2020
This sample was updated to fit TensorRT API changes since version 6.3. Please see [TensorRT API](http://docs.nvidia.com/deeplearning/sdk/tensorrt-api/index.html).
Sep. 2020
This sample was updated to fit TensorRT API changes since version 6.4.
Mar. 2022
This sample was updated for DriveOS 6.0 and later releases.
Jun. 2023
This sample was updated to remove deprecated APIs of ICudaEngine and IExecutionContext.
Jan. 2024
Update static linking description
Feb. 2025
This sample was updated for TRT 10.x and later releases.
Jul. 2025
This sample was updated for the TRT 10.13.1 safety release.
Dec. 2025
This sample was updated to use the CMake-based build system.
Apr. 2026
This sample was updated to add the `--cpuOnly` build option for remote auto-tuning workflows without requiring a local GPU on the build host.
## Known issues
There are no known issues in this sample.
samples/sampleSafeMNIST/sampleSafeMNISTBuild.cpp
+389
View File
@@ -0,0 +1,389 @@
/*
* 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.
*/
//! \file sampleSafeMNISTBuild.cpp
//! \brief This file contains the implementation of the MNIST sample.
//!
//! It builds a TensorRT safe engine by importing a trained MNIST ONNX model.
//! It can be run with the following command line:
//! Command: ./sample_mnist_safe_build [-h or --help] [-d=/path/to/data/dir or --datadir=/path/to/data/dir]
#define DEFINE_TRT_ENTRYPOINTS 1
#define DEFINE_TRT_ONNX_PARSER_ENTRYPOINT 0
#define DEFINE_TRT_BUILDER_ENTRYPOINT 0
#define DEFINE_TRT_REFITTER_ENTRYPOINT 0
#define DEFINE_TRT_RUNTIME_ENTRYPOINT 0
#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "parserOnnxConfig.h"
#include "safeCommon.h"
#include "sampleUtils.h"
#include "NvInfer.h"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cuda_runtime_api.h>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
using namespace nvinfer1;
std::string const gSampleName = "TensorRT.sample_mnist_safe_build";
//!
//! \brief The SampleSafeMNISTBuildArgs struct stores the additional arguments required by the sample
//!
struct SampleSafeMNISTBuildArgs : public samplesCommon::Args
{
std::string engineFileName{"safe_mnist.engine"};
bool verbose{false};
std::string remoteAutoTuningConfig{};
int32_t maxAuxStreams{0};
bool cpuOnly{false};
};
//!
//! \brief This function parses arguments specific to the sample
//!
bool parseSampleSafeMNISTBuildArgs(SampleSafeMNISTBuildArgs& args, int32_t argc, char* argv[])
{
using namespace samplesSafeCommon;
for (int32_t i = 1; i < argc; ++i)
{
std::string const arg = argv[i];
if (parseBool(arg, "help", 'h'))
{
args.help = true;
}
else if (parseBool(arg, "verbose"))
{
args.verbose = true;
}
else if (parseBool(arg, "cpuOnly"))
{
args.cpuOnly = true;
}
else if (auto const value = parseString(arg, "saveEngine"))
{
if (value->empty())
{
sample::gLogError << "Engine filename cannot be empty\n";
return false;
}
args.engineFileName = std::move(*value);
}
else if (auto const value = parseString(arg, "remoteAutoTuningConfig"))
{
if (value->empty())
{
sample::gLogError << "Remote auto tuning config cannot be empty\n";
return false;
}
args.remoteAutoTuningConfig = std::move(*value);
}
else if (auto const value = parseString(arg, "datadir", 'd'))
{
if (value->empty())
{
sample::gLogError << "Data directory path cannot be empty\n";
return false;
}
args.dataDirs.push_back(sample::normalizeDirectoryPath(*value));
}
else if (auto const value = parseString(arg, "maxAuxStreams"))
{
args.maxAuxStreams = std::stoi(*value);
if (args.maxAuxStreams < 0)
{
sample::gLogError << "Number of auxiliary streams must be >= 0, got: " << arg << "\n";
return false;
}
}
else
{
sample::gLogError << "Invalid Argument: " << arg << "\n";
return false;
}
}
return true;
}
//!
//! \brief The SampleSafeMNISTBuildParams struct stores the additional parameters required by the sample
//!
struct SampleSafeMNISTBuildParams : public samplesCommon::OnnxSampleParams
{
std::string engineFileName{};
std::string remoteAutoTuningConfig{};
int32_t maxAuxStreams{0};
};
//!
//! \brief Initialize members of the params struct using the command line args.
//!
SampleSafeMNISTBuildParams initializeSampleParams(SampleSafeMNISTBuildArgs const& args)
{
SampleSafeMNISTBuildParams 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 = "safe_mnist.onnx";
params.engineFileName = args.engineFileName;
params.remoteAutoTuningConfig = args.remoteAutoTuningConfig;
params.maxAuxStreams = args.maxAuxStreams;
return params;
}
//!
//! \brief The SampleSafeMNIST class implements the MNIST sample.
//!
//! \details It creates the network using a trained ONNX MNIST classification model.
//!
class SampleSafeMNIST
{
public:
SampleSafeMNIST(SampleSafeMNISTBuildParams const& params)
: mParams(params)
{
}
//!
//! \brief Builds the network engine.
//!
bool build();
private:
//!
//! \brief Uses an ONNX parser to create the MNIST Network and marks the
//! output layers.
//!
bool constructNetwork(std::unique_ptr<nvonnxparser::IParser>& parser);
SampleSafeMNISTBuildParams mParams; //!< The parameters for the sample.
nvinfer1::Dims mInputDims; //!< The dimensions of the input to the network.
};
//!
//! \brief Creates the network, configures the builder and creates the network engine.
//!
//! \details This function creates the MNIST network by parsing the ONNX model and builds
//! the engine that will be used to run MNIST.
//!
//! \return true if the engine was created successfully and false otherwise.
//!
bool SampleSafeMNIST::build()
{
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
if (!builder)
{
return false;
}
NetworkDefinitionCreationFlags flags = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED);
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(flags));
if (!network)
{
return false;
}
auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
if (!config)
{
return false;
}
auto parser
= std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
if (!parser)
{
return false;
}
auto constructed = constructNetwork(parser);
if (!constructed)
{
return false;
}
config->setEngineCapability(EngineCapability::kSAFETY);
config->setMaxAuxStreams(mParams.maxAuxStreams);
config->setFlag(BuilderFlag::kGPU_FALLBACK);
// Set remote auto tuning config if provided
if (!mParams.remoteAutoTuningConfig.empty())
{
config->setRemoteAutoTuningConfig(mParams.remoteAutoTuningConfig.c_str());
}
auto buffer = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
if (!buffer)
{
return false;
}
ASSERT(network->getNbInputs() == 1);
mInputDims = network->getInput(0)->getDimensions();
ASSERT(mInputDims.nbDims == 4);
// Save the engine
std::string engineFile = mParams.engineFileName;
std::ofstream file(engineFile, std::ios::binary);
if (!file)
{
sample::gLogError << "Failed to open file to save engine: " << engineFile << std::endl;
return false;
}
file.write(reinterpret_cast<char const*>(buffer->data()), buffer->size());
file.close();
return true;
}
//!
//! \brief Uses an ONNX parser to create the MNIST Network and marks the
//! output layers.
//!
//! \param network Pointer to the network that will be populated with the MNIST network.
//!
//! \param builder Pointer to the engine builder.
//!
bool SampleSafeMNIST::constructNetwork(std::unique_ptr<nvonnxparser::IParser>& parser)
{
return parser->parseFromFile(samplesCommon::locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
static_cast<int32_t>(sample::gLogger.getReportableSeverity()));
}
//!
//! \brief Prints the help information for running this sample.
//!
void printHelpInfo()
{
SampleSafeMNISTBuildArgs const defArgs{};
std::cout << R"(Usage: sample_mnist_safe_build [options]
Options:
--help, -h Print this message and exit.
--datadir=DIR, -d=DIR Search for data in DIR. This option can be passed multiple times
to add multiple search directories. If omitted, default data dirs are:
data/samples/mnist/, data/mnist/
--verbose Use verbose logging.
--saveEngine=FILE Save the serialized engine into FILE (default = )"
<< defArgs.engineFileName << R"().
--remoteAutoTuningConfig=CONFIG
Set remote auto tuning configuration in the following format:
protocol://username[:password]@hostname[:port]?param1=value1&param2=value2
--maxAuxStreams=N Limit the number of auxiliary streams to N (default = )"
<< defArgs.maxAuxStreams << R"().
--cpuOnly Build the engine with CPU-only mode. Requires --remoteAutoTuningConfig.
No local GPU is required on the build machine.
Examples:
sample_mnist_safe_build \
--remoteAutoTuningConfig=ssh://user:pass@192.0.2.100:22?remote_exec_path=/opt/tensorrt/bin&remote_lib_path=/opt/tensorrt/lib
)";
}
int main(int argc, char** argv)
{
SampleSafeMNISTBuildArgs args;
bool argsOK = parseSampleSafeMNISTBuildArgs(args, argc, argv);
if (!argsOK)
{
printHelpInfo();
return EXIT_FAILURE;
}
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
if (args.verbose)
{
sample::setReportableSeverity(ILogger::Severity::kVERBOSE);
}
// Log remoteAutoTuningConfig usage
if (!args.remoteAutoTuningConfig.empty())
{
sample::gLogInfo << "Remote auto tuning config specified: "
<< sample::sanitizeRemoteAutoTuningConfig(args.remoteAutoTuningConfig) << std::endl;
sample::gLogInfo << "This is a safety sample and will build in remote mode automatically." << std::endl;
}
if (args.cpuOnly)
{
if (args.remoteAutoTuningConfig.empty())
{
sample::gLogError << "--cpuOnly requires --remoteAutoTuningConfig to be specified." << std::endl;
printHelpInfo();
return EXIT_FAILURE;
}
sample::gLogInfo << "Setting CPU-only mode" << std::endl;
if (!samplesSafeCommon::applyCpuOnlyMode())
{
return EXIT_FAILURE;
}
}
if (!args.cpuOnly && !samplesCommon::isSmSafe())
{
sample::gLogInfo << "Skip safe mode test on unsupported platforms." << std::endl;
return EXIT_SUCCESS;
}
// Create sanitized argv for logging to avoid exposing credentials in test reports
auto sanitizedArgs = sample::sanitizeArgv(argc, argv);
std::vector<char const*> sanitizedArgv;
sanitizedArgv.reserve(sanitizedArgs.size());
for (auto const& s : sanitizedArgs)
{
sanitizedArgv.push_back(s.c_str());
}
auto sampleTest
= sample::gLogger.defineTest(gSampleName, static_cast<int32_t>(sanitizedArgv.size()), sanitizedArgv.data());
sample::gLogger.reportTestStart(sampleTest);
SampleSafeMNISTBuildParams params = initializeSampleParams(args);
SampleSafeMNIST sample(params);
sample::gLogInfo << "Building a GPU inference engine for MNIST" << std::endl;
if (!sample.build())
{
return sample::gLogger.reportFail(sampleTest);
}
return sample::gLogger.reportPass(sampleTest);
}
samples/sampleSafeMNIST/sampleSafeMNISTInfer.cpp
+454
View File
@@ -0,0 +1,454 @@
/*
* 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.
*/
//! \file sampleSafeMNISTInfer.cpp
//! \brief This file contains the implementation of the MNIST sample.
//!
//! It uses the prebuilt TensorRT engine to run inference on an input image of a digit.
//! It can be run with the following command line:
//! Command: ./sample_mnist_safe_infer
#include "NvInferSafeRuntime.h"
#include "safeCommon.h"
#include "safeErrorRecorder.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cuda_runtime_api.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <numeric>
#include <random>
#include <string_view>
#include <thread>
#include <vector>
using namespace samplesSafeCommon;
namespace
{
//!
//! \brief Locate path to file by its filename. Will walk back MAX_DEPTH dirs from CWD to check for such a file path.
//!
std::string locateFile(std::string const& fileName, nvinfer2::safe::ISafeRecorder& recorder)
{
constexpr uint32_t MAX_DEPTH{10U};
std::array<std::string const, 2> const dirPatterns
= {std::string{"data/samples/mnist/"}, std::string{"data/mnist/"}};
std::string foundFile{};
for (auto const& dir : dirPatterns)
{
std::string file{dir + fileName};
bool found{false};
for (uint32_t i = 0U; i < MAX_DEPTH; i++)
{
std::ifstream checkFile(file);
found = checkFile.is_open();
if (found)
{
break;
}
file = "../" + file; // Try again in parent dir.
}
if (found)
{
foundFile = file;
break;
}
}
if (foundFile.empty())
{
safeLogError(recorder, "Could not find " + fileName + " in data/samples/mnist/ or data/mnist.");
safeLogError(recorder, "&&&& FAILED");
exit(EXIT_FAILURE);
}
return foundFile;
}
//!
//! \brief Reads the input data, preprocesses, and stores the result in a managed buffer.
//!
bool processInput(void* input, int32_t const inputFileIdx, nvinfer2::safe::ISafeRecorder& recorder)
{
std::stringstream ss;
constexpr int32_t kINPUT_H{28};
constexpr int32_t kINPUT_W{28};
// Read the digit file according to the inputFileIdx.
std::vector<uint8_t> fileData(kINPUT_H * kINPUT_W);
readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm", recorder), fileData.data(), kINPUT_H, kINPUT_W);
// Print ASCII representation of digit.
ss << "Input:\n";
for (int32_t i = 0; i < kINPUT_H * kINPUT_W; i++)
{
ss << (" .:-=+*#%@"[fileData[i] / 26U]) << (((i + 1) % kINPUT_W) ? "" : "\n");
}
safeLogInfo(recorder, ss.str());
float* hostInputBuffer = static_cast<float*>(input);
std::copy(fileData.begin(), fileData.end(), hostInputBuffer);
// Normalize to 0-1 with background at 0
std::transform(hostInputBuffer, hostInputBuffer + kINPUT_H * kINPUT_W, hostInputBuffer,
[](float v) -> float { return 1.0f - v / 255.0f; });
return true;
}
//!
//! \brief Verifies that the output is correct and prints it.
//!
bool verifyOutput(void* output, int32_t groundTruthDigit, nvinfer2::safe::ISafeRecorder& recorder)
{
float* prob = static_cast<float*>(output);
// Print histogram of the output distribution.
safeLogInfo(recorder, "Output:");
float val{0.0f};
int32_t idx{0};
constexpr int32_t kDIGITS{10};
// Calculate Softmax
float sum{0.0f};
for (int32_t i = 0; i < kDIGITS; ++i)
{
prob[i] = exp(prob[i]);
sum += prob[i];
}
for (int32_t i = 0; i < kDIGITS; ++i)
{
std::stringstream ss;
prob[i] /= sum;
if (val < prob[i])
{
val = prob[i];
idx = i;
}
ss << " Prob " << i << " " << std::fixed << std::setw(5) << std::setprecision(4) << prob[i] << " Class " << i
<< ": " << std::string(int32_t(std::floor(prob[i] * 10 + 0.5f)), '*');
safeLogInfo(recorder, ss.str());
}
return (idx == groundTruthDigit && val > 0.9f);
}
//!
//! \brief Loads the enginePlanFile from engineFile and returns it.
//!
std::vector<char> loadEnginePlanFile(std::string const& engineFile, int& size, nvinfer2::safe::ISafeRecorder& recorder)
{
std::string const& filename = engineFile;
std::vector<char> gieModelStream;
std::ifstream file(filename, std::ios::binary);
if (!file.good())
{
safeLogError(recorder, "Could not open input engine file or file is empty. File name: " + filename);
return {};
}
file.seekg(0, std::ifstream::end);
size = file.tellg();
file.seekg(0, std::ifstream::beg);
gieModelStream.resize(size);
file.read(gieModelStream.data(), size);
file.close();
return gieModelStream;
}
//!
//! \brief Returns a random digit between 0 and 9
//!
int32_t getRandomDigit()
{
std::random_device rd;
std::default_random_engine generator{rd()};
std::uniform_int_distribution<int32_t> distribution(0, 9);
return distribution(generator);
}
//!
//! \brief Structure representing memory allocation for CUDA
//!
struct CudaMemory
{
void* hostPtr = nullptr;
void* devicePtr = nullptr;
size_t size = 0;
};
//!
//! \brief Do inference
//!
void doInferenceThread(nvinfer2::safe::ITRTGraph* graph, int8_t& ret_status, nvinfer2::safe::ISafeRecorder* recorder)
{
// Initialize to success; will be set to 0 on any error.
ret_status = 1;
int64_t nbIOs{};
SAFE_API_CALL(graph->getNbIOTensors(nbIOs), *recorder);
// This sample only has one input and one output.
SAFE_ASSERT(nbIOs == 2);
CudaMemory inputCudaMemory;
CudaMemory outputCudaMemory;
// Initialize main stream
cudaStream_t stream;
CUDA_CALL(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), *recorder);
// Setup as many auxiliary streams as the graph requires - destroyed at scope end.
auto auxStreamsDeleter = samplesSafeCommon::setUpAuxStreamsOn(*graph, *recorder);
// Pick a random digit to try to infer.
int32_t digit = getRandomDigit();
// Iterate through all input/output tensors
for (int64_t i = 0; i < nbIOs; ++i)
{
// Get the tensor name for the current I/O tensor
char const* tensorName;
SAFE_API_CALL(graph->getIOTensorName(tensorName, i), *recorder);
// Get tensor descriptor which contains metadata like size and I/O mode
nvinfer2::safe::TensorDescriptor desc;
SAFE_API_CALL(graph->getIOTensorDescriptor(desc, tensorName), *recorder);
// Allocate device and host memory for this tensor
void* deviceBuf = nullptr;
void* hostBuf = nullptr;
CUDA_CALL(cudaMalloc(&deviceBuf, desc.sizeInBytes), *recorder);
CUDA_CHECK(cudaHostAlloc(&hostBuf, desc.sizeInBytes, cudaHostAllocDefault));
if (desc.ioMode == TensorIOMode::kINPUT)
{
// Read the input data into the managed buffers.
processInput(hostBuf, digit, *recorder);
// Asynchronously copy data from host input buffers to device input buffers.
CUDA_CHECK(cudaMemcpyAsync(deviceBuf, hostBuf, desc.sizeInBytes, cudaMemcpyHostToDevice, stream));
inputCudaMemory = {hostBuf, deviceBuf, desc.sizeInBytes};
}
else if (desc.ioMode == TensorIOMode::kOUTPUT)
{
CUDA_CALL(cudaMemsetAsync(deviceBuf, 0, desc.sizeInBytes, stream), *recorder);
outputCudaMemory = {hostBuf, deviceBuf, desc.sizeInBytes};
}
else
{
safeLogError(*recorder, "Unexpected tensor IO mode");
ret_status = 0;
}
SAFE_ASSERT(desc.dataType == DataType::kFLOAT);
// Create a typed array for the tensor
nvinfer2::safe::TypedArray tensor
= nvinfer2::safe::TypedArray(static_cast<float*>(deviceBuf), desc.sizeInBytes);
SAFE_API_CALL(graph->setIOTensorAddress(tensorName, tensor), *recorder);
}
cudaEvent_t inputConsumedEvent;
cudaEventCreate(&inputConsumedEvent);
SAFE_API_CALL(graph->setInputConsumedEvent(inputConsumedEvent), *recorder);
// Run the graph
SAFE_API_CALL(graph->executeAsync(stream), *recorder);
cudaEvent_t retrievedEvent;
SAFE_API_CALL(graph->getInputConsumedEvent(retrievedEvent), *recorder);
SAFE_ASSERT(retrievedEvent != nullptr);
cudaEventSynchronize(retrievedEvent);
// Synchronize the network
SAFE_API_CALL(graph->sync(), *recorder);
// Asynchronously copy data from device output buffers to host output buffers.
CUDA_CHECK(cudaMemcpyAsync(
outputCudaMemory.hostPtr, outputCudaMemory.devicePtr, outputCudaMemory.size, cudaMemcpyDeviceToHost, stream));
// Wait for the work in the stream to complete.
CUDA_CHECK(cudaStreamSynchronize(stream));
// Check and print the output of the inference.
if (!verifyOutput(outputCudaMemory.hostPtr, digit, *recorder))
{
safeLogError(*recorder, "Failed to verify output");
ret_status = 0;
}
// Release stream and buffers.
CUDA_CHECK(cudaStreamDestroy(stream));
CUDA_CHECK(cudaFreeHost(inputCudaMemory.hostPtr));
CUDA_CHECK(cudaFreeHost(outputCudaMemory.hostPtr));
CUDA_CHECK(cudaFree(inputCudaMemory.devicePtr));
CUDA_CHECK(cudaFree(outputCudaMemory.devicePtr));
}
//!
//! \brief The SampleSafeMNISTInferArgs struct stores the additional arguments required by the sample
//!
struct SampleSafeMNISTInferArgs
{
std::string engineFileName{"safe_mnist.engine"};
int32_t threads{1};
bool help{false};
};
//!
//! \brief Runs the TensorRT inference engine for this sample.
//!
//! \details This function is the main execution function of the sample. It allocates
//! the buffer, sets inputs, executes the engine, and verifies the output.
//!
bool doInference(SampleSafeMNISTInferArgs const& args)
{
int32_t const nbThreads = args.threads;
std::vector<int8_t> ret_status(nbThreads);
std::vector<std::unique_ptr<sample::SampleSafeRecorder>> recorders(nbThreads);
for (int32_t i = 0; i < nbThreads; ++i)
{
recorders[i] = std::make_unique<sample::SampleSafeRecorder>(nvinfer2::safe::Severity::kINFO, i);
}
// Load safe engine blob
int32_t engineFileSize = 0;
auto gieModelStream = loadEnginePlanFile(args.engineFileName, engineFileSize, *recorders[0]);
SAFE_ASSERT(engineFileSize != 0);
// Configure executor(s)
std::vector<nvinfer2::safe::ITRTGraph*> graphs(nbThreads);
SAFE_API_CALL(nvinfer2::safe::createTRTGraph(graphs[0], gieModelStream.data(), engineFileSize, *recorders[0], true),
*recorders[0]);
for (int32_t i = 1; i < nbThreads; ++i)
{
SAFE_API_CALL(graphs[0]->clone(graphs[i], *recorders[i]), *recorders[0]);
}
// Run the graphs in independent threads
std::vector<std::thread> threads(nbThreads);
for (int32_t i = 0; i < nbThreads; ++i)
{
threads[i] = std::thread{doInferenceThread, graphs[i], std::ref(ret_status[i]), recorders[i].get()};
}
for (int32_t i = 0; i < nbThreads; ++i)
{
threads[i].join();
if (!ret_status[i])
{
return false;
}
}
for (int32_t i = 0; i < nbThreads; ++i)
{
SAFE_API_CALL(nvinfer2::safe::destroyTRTGraph(graphs[i]), *recorders[i]);
graphs[i] = nullptr;
}
return true;
}
//!
//! \brief This function parses arguments specific to the sample
//!
bool parseSampleSafeMNISTInferArgs(SampleSafeMNISTInferArgs& args, int32_t argc, char* argv[])
{
for (int32_t i = 1; i < argc; ++i)
{
std::string const arg = argv[i];
if (auto const value = parseString(arg, "loadEngine"))
{
args.engineFileName = *value;
}
else if (auto const value = parseString(arg, "threads"))
{
args.threads = std::stoi(*value);
if (args.threads <= 0)
{
SAFE_LOG << "Number of threads must be > 0, got: " << arg << "\n";
return false;
}
}
else if (parseBool(arg, "help", 'h'))
{
args.help = true;
}
else
{
SAFE_LOG << "Invalid Argument: " << arg << "\n";
return false;
}
}
return true;
}
//!
//! \brief Prints the help information for running this sample.
//!
void printHelpInfo()
{
SampleSafeMNISTInferArgs const defArgs{};
std::cout << R"(Usage: sample_mnist_safe_infer [options]
Options:
--help, -h Print this message and exit.
--loadEngine=FILE Load serialized engine from FILE (default = )"
<< defArgs.engineFileName << R"().
--threads=N Run inference in N threads concurrently (default = )"
<< defArgs.threads << R"().
)";
}
} // namespace
int32_t main(int32_t argc, char** argv)
{
safetyCompliance::setPromgrAbility();
SampleSafeMNISTInferArgs args;
bool argsOK = parseSampleSafeMNISTInferArgs(args, argc, argv);
if (!argsOK)
{
printHelpInfo();
return EXIT_FAILURE;
}
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
// Initialize SafeCuda before any other Cuda APIs are called. This may be skipped if createInferRuntime() is called
// first as per DEEPLRN_RES_116
safetyCompliance::initSafeCuda();
if (!isSmSafe())
{
SAFE_LOG << "Skip safe mode test on unsupported platforms." << std::endl;
return EXIT_SUCCESS;
}
TestResult result = doInference(args) ? TestResult::kPASSED : TestResult::kFAILED;
reportTestResult("TensorRT.sample_mnist_safe_infer", result, argc, argv);
return EXIT_SUCCESS;
}