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

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
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/sampleProgressMonitor/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_progress_monitor sampleProgressMonitor.cpp)
target_link_libraries(sample_progress_monitor PRIVATE trt_samples_common TRT_SAMPLES::tensorrt)
add_dependencies(tensorrt_samples sample_progress_monitor)
installLibraries(
TARGETS sample_progress_monitor
OPTIONAL
COMPONENT internal
)
samples/sampleProgressMonitor/README.md
+182
View File
@@ -0,0 +1,182 @@
# Progress Monitor API usage example based off sampleMNIST in TensorRT
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
- [Progress bar display](#progress-bar-display)
- [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, sampleProgressMonitor, shows an example of how to use the progress monitor API based on sampleOnnxMNIST ([documentation](https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html#onnx_mnist_sample)).
This sample demonstrates the usage of `IProgressMonitor` to report the status of TRT engine-building operations.
## How does this sample work?
This sample uses a Onnx model that was trained on the [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md).
Specifically, this sample performs the following steps:
- Performs the basic setup and initialization of TensorRT using the Onnx parser
- [Imports a trained Onnx model using Onnx parser](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#import_onnx_c)
- Preprocesses the input and stores the result in a managed buffer
- Builds an engine using incremental progress reporting
- [Serializes and deserializes the engines](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#serial_model_c)
- [Uses the engines to perform inference on an input image](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#perform_inference_c)
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 to be that in the image.
### Progress bar display
This sample implements an `IProgressMonitor` to display progress while building a TensorRT engine. Each long-running step of the process can define a new progress phase, nesting them as necessary.
1. Phase entry - The `IProgressMonitor::phaseBegin` callback determines an appropriate nesting level for the new phase and updates the terminal display.
2. Phase progress - The `IProgressMonitor::stepComplete` callback increments the progress bar for the selected phase and updates the terminal display. This sample always returns `true` from `stepComplete` in order to progress the build unconditionally. If you wish to cancel a build in progress, such as in response to user input, you can return `false` from this function to stop the build early.
3. Phase completion - The `IProgressMonitor::phaseEnd` callback removes the line corresponding to the completed phase and updates the terminal display.
The progress bars are drawn using virtual terminal escape sequences to manipulate the terminal's cursor and clear lines.
## Prerequisites
1. Preparing sample data
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 to perform inference on the digit:
```bash
./sample_progress_monitor [-h] [--datadir=/path/to/data/dir/] [--useDLA=N]
```
For example:
```bash
./sample_progress_monitor --datadir $TRT_DATADIR/mnist
```
This sample reads the `mnist.onnx` file to build the network:
**Note:** By default, the 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.
**Note:** The sample should be run from a terminal. It uses xterm-style escape sequences to animate its output, and is not designed to be redirected to a file.
3. Verify that the sample ran successfully. If the sample runs successfully you should see animated progress bars during the network build phase and output similar to the following:
```
&&&& RUNNING TensorRT.sample_progress_monitor [TensorRT v8700] # ./sample_progress_monitor
[I] Building and running a GPU inference engine for MNIST.
[I] [TRT] [MemUsageChange] Init CUDA: CPU +14, GPU +0, now: CPU 19, GPU 1217 (MiB)
[I] [TRT] [MemUsageChange] Init builder kernel library: CPU +1450, GPU +266, now: CPU 1545, GPU 1483 (MiB)
[I] [TRT] ----------------------------------------------------------------
[I] [TRT] Input filename: ../../../../data/samples/mnist/mnist.onnx
[I] [TRT] ONNX IR version: 0.0.3
[I] [TRT] Opset version: 8
[I] [TRT] Producer name: CNTK
[I] [TRT] Producer version: 2.5.1
[I] [TRT] Domain: ai.cntk
[I] [TRT] Model version: 1
[I] [TRT] Doc string:
[I] [TRT] ----------------------------------------------------------------
[W] [TRT] onnx2trt_utils.cpp:374: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[I] [TRT] Graph optimization time: 0.00293778 seconds.
[I] [TRT] Local timing cache in use. Profiling results in this builder pass will not be stored.
[=======---] Building engine 3/4
[----------] Building engine from subgraph 0/1
[----------] Computing profile costs 0/1
[=======---] Timing graph nodes 11/15
[===-------] Finding fastest tactic for Times212 12/37
[==========] Measuring tactic time 4/4
```
After the TensorRT network has been constructed, you should see output similar to the following. An ASCII rendering of the input image with digit 3:
```
&&&& RUNNING TensorRT.sample_progress_monitor # ./sample_progress_monitor
[I] Building and running a GPU inference engine for MNIST
[I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@#-:.-=@@@@@@@@@@@@@@
@@@@@%= . *@@@@@@@@@@@@@
@@@@% .:+%%% *@@@@@@@@@@@@@
@@@@+=#@@@@@# @@@@@@@@@@@@@@
@@@@@@@@@@@% @@@@@@@@@@@@@@
@@@@@@@@@@@: *@@@@@@@@@@@@@@
@@@@@@@@@@- .@@@@@@@@@@@@@@@
@@@@@@@@@: #@@@@@@@@@@@@@@@
@@@@@@@@: +*%#@@@@@@@@@@@@
@@@@@@@% :+*@@@@@@@@
@@@@@@@@#*+--.:: +@@@@@@
@@@@@@@@@@@@@@@@#=:. +@@@@@
@@@@@@@@@@@@@@@@@@@@ .@@@@@
@@@@@@@@@@@@@@@@@@@@#. #@@@@
@@@@@@@@@@@@@@@@@@@@# @@@@@
@@@@@@@@@%@@@@@@@@@@- +@@@@@
@@@@@@@@#-@@@@@@@@*. =@@@@@@
@@@@@@@@ .+%%%%+=. =@@@@@@@
@@@@@@@@ =@@@@@@@@
@@@@@@@@*=: :--*@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
[I] Output:
Prob 1 0.0000 Class 1:
Prob 2 0.0000 Class 2:
Prob 3 1.0000 Class 3: **********
Prob 4 0.0000 Class 4:
Prob 5 0.0000 Class 5:
Prob 6 0.0000 Class 6:
Prob 7 0.0000 Class 7:
Prob 8 0.0000 Class 8:
Prob 9 0.0000 Class 9:
&&&& PASSED TensorRT.sample_progress_monitor # ./sample_progress_monitor
```
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. For example:
```
Usage: ./sample_progress_monitor [-h or --help] [-d or --datadir=<path to data directory>] [--useDLACore=<int>]
--help Display help information
--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/)
--useDLACore=N Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, where n is the number of DLA engines on the platform.
```
# Additional resources
The following resources provide a deeper understanding about sampleProgressMonitor:
**MNIST**
- [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md)
**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 C++ API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#c_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.
**May 2023**
- This `README.md` file was created and reviewed.
# Known issues
There are no known issues in this sample.
samples/sampleProgressMonitor/sampleProgressMonitor.cpp
+564
View File
@@ -0,0 +1,564 @@
/*
* 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.
*/
//! \file sampleProgressMonitor.cpp
//! \brief This file contains the implementation of the Progress Monitor sample.
//!
//! It demonstrates the usage of IProgressMonitor for displaying engine build progress on the user's terminal.
//! It builds a TensorRT engine by importing a trained MNIST ONNX model and runs inference on an input image of a
//! digit.
//! It can be run with the following command line:
//! Command: ./sample_progress_monitor [-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 "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "parserOnnxConfig.h"
#include <algorithm>
#include <cmath>
#include <cuda_runtime_api.h>
#include <iomanip>
#include <iostream>
#include <random>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
using namespace nvinfer1;
std::string const gSampleName = "TensorRT.sample_progress_monitor";
//!
//! \brief The ConsoleProgressMonitor class displays a simple progress graph for each step of the build process.
//!
class ConsoleProgressMonitor : public IProgressMonitor
{
public:
void phaseStart(char const* phaseName, char const* parentPhase, int32_t nbSteps) noexcept final
{
PhaseEntry newPhase;
newPhase.title = phaseName;
newPhase.nbSteps = nbSteps;
PhaseIter iParent = mPhases.end();
if (parentPhase)
{
iParent = findPhase(parentPhase);
newPhase.nbIndents = 1 + iParent->nbIndents;
do
{
++iParent;
} while (iParent != mPhases.end() && iParent->nbIndents >= newPhase.nbIndents);
}
mPhases.insert(iParent, newPhase);
redraw();
}
bool stepComplete(char const* phaseName, int32_t step) noexcept final
{
PhaseIter const iPhase = findPhase(phaseName);
iPhase->steps = step;
redraw();
return true;
}
void phaseFinish(char const* phaseName) noexcept final
{
PhaseIter const iPhase = findPhase(phaseName);
iPhase->active = false;
redraw();
mPhases.erase(iPhase);
}
private:
struct PhaseEntry
{
std::string title;
int32_t steps{0};
int32_t nbSteps{0};
int32_t nbIndents{0};
bool active{true};
};
using PhaseIter = std::vector<PhaseEntry>::iterator;
std::vector<PhaseEntry> mPhases;
static int32_t constexpr kPROGRESS_INNER_WIDTH = 10;
void redraw()
{
auto const moveToStartOfLine = []() { std::cout << "\x1b[0G"; };
auto const clearCurrentLine = []() { std::cout << "\x1b[2K"; };
moveToStartOfLine();
int32_t inactivePhases = 0;
for (PhaseEntry const& phase : mPhases)
{
clearCurrentLine();
if (phase.nbIndents > 0)
{
for (int32_t indent = 0; indent < phase.nbIndents; ++indent)
{
std::cout << ' ';
}
}
if (phase.active)
{
std::cout << progressBar(phase.steps, phase.nbSteps) << ' ' << phase.title << ' ' << phase.steps << '/'
<< phase.nbSteps << std::endl;
}
else
{
// Don't draw anything at this time, but prepare to emit blank lines later.
// This ensures that stale phases are removed from display rather than lingering.
++inactivePhases;
}
}
for (int32_t phase = 0; phase < inactivePhases; ++phase)
{
clearCurrentLine();
std::cout << std::endl;
}
// Move (mPhases.size()) lines up so that logger output can overwrite the progress bars.
std::cout << "\x1b[" << mPhases.size() << "A";
}
std::string progressBar(int32_t steps, int32_t nbSteps) const
{
std::ostringstream bar;
bar << '[';
int32_t const completedChars
= static_cast<int32_t>(kPROGRESS_INNER_WIDTH * steps / static_cast<float>(nbSteps));
for (int32_t i = 0; i < completedChars; ++i)
{
bar << '=';
}
for (int32_t i = completedChars; i < kPROGRESS_INNER_WIDTH; ++i)
{
bar << '-';
}
bar << ']';
return bar.str();
}
PhaseIter findPhase(std::string const& title)
{
return std::find_if(mPhases.begin(), mPhases.end(),
[title](PhaseEntry const& phase) { return phase.title == title && phase.active; });
}
};
//!
//! \brief The SampleProgressMonitor class implements the SampleProgressReporter sample.
//!
//! \details It creates the network using a trained ONNX MNIST classification model.
//!
class SampleProgressMonitor
{
public:
explicit SampleProgressMonitor(samplesCommon::OnnxSampleParams const& params)
: mParams(params)
{
}
//!
//! \brief Builds the network engine.
//!
bool build(IProgressMonitor* monitor);
//!
//! \brief Runs the TensorRT inference engine for this sample.
//!
bool infer();
private:
//!
//! \brief uses a Onnx parser to create the MNIST Network and marks the output layers.
//!
bool constructNetwork(std::unique_ptr<nvinfer1::IBuilder>& builder,
std::unique_ptr<nvinfer1::INetworkDefinition>& network, std::unique_ptr<nvinfer1::IBuilderConfig>& config,
std::unique_ptr<nvonnxparser::IParser>& parser);
//!
//! \brief Reads the input and mean data, preprocesses, and stores the result in a managed buffer.
//!
bool processInput(
samplesCommon::BufferManager const& buffers, std::string const& inputTensorName, int32_t inputFileIdx) const;
//!
//! \brief Verifies that the output is correct and prints it.
//!
bool verifyOutput(samplesCommon::BufferManager const& buffers, std::string const& outputTensorName,
int32_t groundTruthDigit) const;
std::unique_ptr<IRuntime> mRuntime{};
std::shared_ptr<nvinfer1::ICudaEngine> mEngine{nullptr}; //!< The TensorRT engine used to run the network.
samplesCommon::OnnxSampleParams 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 (mEngine).
//!
//! \return true if the engine was created successfully and false otherwise.
//!
bool SampleProgressMonitor::build(IProgressMonitor* monitor)
{
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
if (!builder)
{
return false;
}
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
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(builder, network, config, parser);
if (!constructed)
{
return false;
}
config->setProgressMonitor(monitor);
samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore, true /*GPUFallback*/);
if (!mRuntime)
{
mRuntime = std::unique_ptr<IRuntime>(createInferRuntime(sample::gLogger.getTRTLogger()));
}
if (!mRuntime)
{
return false;
}
// CUDA stream used for profiling by the builder.
auto profileStream = samplesCommon::makeCudaStream();
if (!profileStream)
{
return false;
}
config->setProfileStream(*profileStream);
std::unique_ptr<nvinfer1::ITimingCache> timingCache{};
// Load timing cache
if (!mParams.timingCacheFile.empty())
{
timingCache
= samplesCommon::buildTimingCacheFromFile(sample::gLogger.getTRTLogger(), *config, mParams.timingCacheFile);
}
std::unique_ptr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
if (!plan)
{
return false;
}
if (timingCache != nullptr && !mParams.timingCacheFile.empty())
{
samplesCommon::updateTimingCacheFile(
sample::gLogger.getTRTLogger(), mParams.timingCacheFile, timingCache.get(), *builder);
}
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(mRuntime->deserializeCudaEngine(plan->data(), plan->size()));
if (!mEngine)
{
return false;
}
ASSERT(network->getNbInputs() == 1);
mInputDims = network->getInput(0)->getDimensions();
ASSERT(mInputDims.nbDims == 4);
return true;
}
//!
//! \brief Reads the input and mean data, preprocesses, and stores the result in a managed buffer.
//!
bool SampleProgressMonitor::processInput(
samplesCommon::BufferManager const& buffers, std::string const& inputTensorName, int32_t inputFileIdx) const
{
int32_t const inputH = mInputDims.d[2];
int32_t const inputW = mInputDims.d[3];
std::vector<uint8_t> fileData(inputH * inputW);
samplesCommon::readPGMFile(samplesCommon::locateFile(std::to_string(inputFileIdx) + ".pgm", mParams.dataDirs),
fileData.data(), inputH, inputW);
// Print ASCII representation of digit.
sample::gLogInfo << "Input:\n";
for (int32_t i = 0; i < inputH * inputW; i++)
{
sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
}
sample::gLogInfo << std::endl;
float* hostInputBuffer = static_cast<float*>(buffers.getHostBuffer(inputTensorName));
for (int32_t i = 0; i < inputH * inputW; i++)
{
hostInputBuffer[i] = 1.0F - static_cast<float>(fileData[i]) / 255.0F;
}
return true;
}
//!
//! \brief Verifies that the output is correct and prints it.
//!
bool SampleProgressMonitor::verifyOutput(
samplesCommon::BufferManager const& buffers, std::string const& outputTensorName, int32_t groundTruthDigit) const
{
float* prob = static_cast<float*>(buffers.getHostBuffer(outputTensorName));
int32_t constexpr kDIGITS = 10;
std::for_each(prob, prob + kDIGITS, [](float& n) { n = exp(n); });
float const sum = std::accumulate(prob, prob + kDIGITS, 0.F);
std::for_each(prob, prob + kDIGITS, [sum](float& n) { n = n / sum; });
auto max_ele = std::max_element(prob, prob + kDIGITS);
float const val = *max_ele;
int32_t const idx = max_ele - prob;
// Print histogram of the output probability distribution.
sample::gLogInfo << "Output:\n";
for (int32_t i = 0; i < kDIGITS; i++)
{
sample::gLogInfo << " 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)), '*')
<< std::endl;
}
sample::gLogInfo << std::endl;
return (idx == groundTruthDigit && val > 0.9F);
}
//!
//! \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 SampleProgressMonitor::constructNetwork(std::unique_ptr<nvinfer1::IBuilder>& builder,
std::unique_ptr<nvinfer1::INetworkDefinition>& network, std::unique_ptr<nvinfer1::IBuilderConfig>& config,
std::unique_ptr<nvonnxparser::IParser>& parser)
{
auto parsed = parser->parseFromFile(samplesCommon::locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
static_cast<int32_t>(sample::gLogger.getReportableSeverity()));
if (!parsed)
{
return false;
}
samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);
return true;
}
//!
//! \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 SampleProgressMonitor::infer()
{
// Create RAII buffer manager object.
samplesCommon::BufferManager buffers(mEngine);
auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
return false;
}
// Pick a random digit to try to infer.
int32_t const digit = std::invoke([] {
auto device = std::random_device();
return std::uniform_int_distribution<int>{0, 9}(device);
});
// Read the input data into the managed buffers.
// There should be just 1 input tensor.
ASSERT(mParams.inputTensorNames.size() == 1);
if (!processInput(buffers, mParams.inputTensorNames[0], digit))
{
return false;
}
// Create CUDA stream for the execution of this inference.
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// Asynchronously copy data from host input buffers to device input buffers
buffers.copyInputToDeviceAsync(stream);
for (int32_t i = 0, e = mEngine->getNbIOTensors(); i < e; i++)
{
auto const& name = mEngine->getIOTensorName(i);
context->setTensorAddress(name, buffers.getDeviceBuffer(name));
}
// Asynchronously enqueue the inference work
if (!context->enqueueV3(stream))
{
return false;
}
// Asynchronously copy data from device output buffers to host output buffers.
buffers.copyOutputToHostAsync(stream);
// Wait for the work in the stream to complete.
CHECK(cudaStreamSynchronize(stream));
// Release stream.
CHECK(cudaStreamDestroy(stream));
// Check and print the output of the inference.
// There should be just one output tensor.
ASSERT(mParams.outputTensorNames.size() == 1);
bool outputCorrect = verifyOutput(buffers, mParams.outputTensorNames[0], digit);
return outputCorrect;
}
//!
//! \brief Initializes members of the params struct using the command line args
//!
samplesCommon::OnnxSampleParams initializeSampleParams(samplesCommon::Args const& 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.dlaCore = args.useDLACore;
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_progress_monitor [-h or --help] [-d or --datadir=<path to data directory>] "
"[--useDLACore=<int>] [--timingCacheFile=<path to timing cache file>]\n";
std::cout << "--help Display help information\n";
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 << "--useDLACore=N Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, "
"where n is the number of DLA engines on the platform."
<< std::endl;
std::cout << "--timingCacheFile Specify path to a timing cache file. If it does not already exist, it will be "
<< "created." << std::endl;
}
int32_t main(int32_t argc, char** argv)
{
samplesCommon::Args args;
bool const 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::Logger::defineTest(gSampleName, argc, argv);
sample::Logger::reportTestStart(sampleTest);
samplesCommon::OnnxSampleParams params = initializeSampleParams(args);
SampleProgressMonitor sampleProgressMonitor(params);
{
sample::gLogInfo << "Building and running a GPU inference engine for MNIST." << std::endl;
ConsoleProgressMonitor progressMonitor;
if (!sampleProgressMonitor.build(&progressMonitor))
{
return sample::Logger::reportFail(sampleTest);
}
if (!sampleProgressMonitor.infer())
{
return sample::Logger::reportFail(sampleTest);
}
}
return sample::Logger::reportPass(sampleTest);
}