added MLP with proper documentation (#874)

* [MLP]: create using core TensortRT python APIs

* [MLP]: README.md added with chart

* [MLP]: C++ TensorRT APIs added

* [MLP]: Updated Docs

* [MLP]: basic .wts added with APIs

* [MLP]: Updated minor details

* [MLP]: fix
This commit is contained in:
Subtain Malik 2022-01-10 08:34:23 +05:00 committed by GitHub
parent e6b22917e6
commit f2ac76df5a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 1159 additions and 0 deletions

View File

@ -53,6 +53,7 @@ Following models are implemented.
|Name | Description |
|-|-|
|[mlp](./mlp) | the very basic model for starters, properly documented |
|[lenet](./lenet) | the simplest, as a "hello world" of this project |
|[alexnet](./alexnet)| easy to implement, all layers are supported in tensorrt |
|[googlenet](./googlenet)| GoogLeNet (Inception v1) |

24
mlp/CMakeLists.txt Normal file
View File

@ -0,0 +1,24 @@
cmake_minimum_required(VERSION 3.14) # change the version, if asked by compiler
project(mlp)
set(CMAKE_CXX_STANDARD 14)
# include and link dirs of tensorrt, you need adapt them if yours are different
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
# include and link dirs of cuda for inference
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# create link for executable files
add_executable(mlp mlp.cpp)
# perform linking with nvinfer libraries
target_link_libraries(mlp nvinfer)
# link with cuda libraries for Inference
target_link_libraries(mlp cudart)
add_definitions(-O2 -pthread)

57
mlp/README.md Normal file
View File

@ -0,0 +1,57 @@
# MLP
MLP is the most basic net in this tensorrtx project for starters. You can learn the basic procedures of building
TensorRT app from the provided APIs. The process of building a TensorRT engine explained in the chart below.
![TensorRT Image](https://user-images.githubusercontent.com/33795294/148565279-795b12da-5243-4e7e-881b-263eb7658683.jpg)
## Helper Files
`logging.h` : A logger file for using NVIDIA TRT API (mostly same for all models)
`mlp.wts` : Converted weight file (simple file, you can open and check it)
## TensorRT C++ API
```
// 1. generate mlp.wts from https://github.com/wang-xinyu/pytorchx/tree/master/mlp -- or use the given .wts file
// 2. put mlp.wts into tensorrtx/mlp (if using the generated weights)
// 3. build and run
cd tensorrtx/mlp
mkdir build
cd build
cmake ..
make
sudo ./mlp -s // serialize model to plan file i.e. 'mlp.engine'
sudo ./mlp -d // deserialize plan file and run inference
```
## TensorRT Python API
```
# 1. Generate mlp.wts from https://github.com/wang-xinyu/pytorchx/tree/master/mlp -- or use the given .wts file
# 2. Put mlp.wts into tensorrtx/mlp (if using the generated weights)
# 3. Install Python dependencies (tensorrt/pycuda/numpy)
# 4. Run
cd tensorrtx/mlp
python mlp.py -s # serialize model to plan file, i.e. 'mlp.engine'
python mlp.py -d # deserialize plan file and run inference
```
## Note
It also supports the latest CUDA-11.4 and TensorRT-8.2.x

503
mlp/logging.h Normal file
View File

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

323
mlp/mlp.cpp Normal file
View File

@ -0,0 +1,323 @@
#include "NvInfer.h" // TensorRT library
#include "iostream" // Standard input/output library
#include "logging.h" // logging file -- by NVIDIA
#include <map> // for weight maps
#include <fstream> // for file-handling
#include <chrono> // for timing the execution
// provided by nvidia for using TensorRT APIs
using namespace nvinfer1;
// Logger from TRT API
static Logger gLogger;
const int INPUT_SIZE = 1;
const int OUTPUT_SIZE = 1;
/** ////////////////////////////
// DEPLOYMENT RELATED /////////
////////////////////////////*/
std::map<std::string, Weights> loadWeights(const std::string file) {
/**
* Parse the .wts file and store weights in dict format.
*
* @param file path to .wts file
* @return weight_map: dictionary containing weights and their values
*/
std::cout << "[INFO]: Loading weights..." << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open Weight file
std::ifstream input(file);
assert(input.is_open() && "[ERROR]: Unable to load weight file...");
// Read number of weights
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
// Loop through number of line, actually the number of weights & biases
while (count--) {
// TensorRT weights
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of weights
std::string w_name;
input >> w_name >> std::dec >> size;
wt.type = DataType::kFLOAT;
uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x) {
// Change hex values to uint32 (for higher values)
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
// Add weight values against its name (key)
weightMap[w_name] = wt;
}
return weightMap;
}
ICudaEngine *createMLPEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt) {
/**
* Create Multi-Layer Perceptron using the TRT Builder and Configurations
*
* @param maxBatchSize: batch size for built TRT model
* @param builder: to build engine and networks
* @param config: configuration related to Hardware
* @param dt: datatype for model layers
* @return engine: TRT model
*/
std::cout << "[INFO]: Creating MLP using TensorRT..." << std::endl;
// Load Weights from relevant file
std::map<std::string, Weights> weightMap = loadWeights("../mlp.wts");
// Create an empty network
INetworkDefinition *network = builder->createNetworkV2(0U);
// Create an input with proper *name
ITensor *data = network->addInput("data", DataType::kFLOAT, Dims3{1, 1, 1});
assert(data);
// Add layer for MLP
IFullyConnectedLayer *fc1 = network->addFullyConnected(*data, 1,
weightMap["linear.weight"],
weightMap["linear.bias"]);
assert(fc1);
// set output with *name
fc1->getOutput(0)->setName("out");
// mark the output
network->markOutput(*fc1->getOutput(0));
// Set configurations
builder->setMaxBatchSize(1);
// Set workspace size
config->setMaxWorkspaceSize(1 << 20);
// Build CUDA Engine using network and configurations
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
assert(engine != nullptr);
// Don't need the network any more
// free captured memory
network->destroy();
// Release host memory
for (auto &mem: weightMap) {
free((void *) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream) {
/**
* Create engine using TensorRT APIs
*
* @param maxBatchSize: for the deployed model configs
* @param modelStream: shared memory to store serialized model
*/
// Create builder with the help of logger
IBuilder *builder = createInferBuilder(gLogger);
// Create hardware configs
IBuilderConfig *config = builder->createBuilderConfig();
// Build an engine
ICudaEngine *engine = createMLPEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// serialize the engine into binary stream
(*modelStream) = engine->serialize();
// free up the memory
engine->destroy();
builder->destroy();
}
void performSerialization() {
/**
* Serialization Function
*/
// Shared memory object
IHostMemory *modelStream{nullptr};
// Write model into stream
APIToModel(1, &modelStream);
assert(modelStream != nullptr);
std::cout << "[INFO]: Writing engine into binary..." << std::endl;
// Open the file and write the contents there in binary format
std::ofstream p("../mlp.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return;
}
p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
// Release the memory
modelStream->destroy();
std::cout << "[INFO]: Successfully created TensorRT engine..." << std::endl;
std::cout << "\n\tRun inference using `./mlp -d`" << std::endl;
}
/** ////////////////////////////
// INFERENCE RELATED //////////
////////////////////////////*/
void doInference(IExecutionContext &context, float *input, float *output, int batchSize) {
/**
* Perform inference using the CUDA context
*
* @param context: context created by engine
* @param input: input from the host
* @param output: output to save on host
* @param batchSize: batch size for TRT model
*/
// Get engine from the context
const ICudaEngine &engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void *buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex("data");
const int outputIndex = engine.getBindingIndex("out");
// Create GPU buffers on device -- allocate memory for input and output
cudaMalloc(&buffers[inputIndex], batchSize * INPUT_SIZE * sizeof(float));
cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float));
// create CUDA stream for simultaneous CUDA operations
cudaStream_t stream;
cudaStreamCreate(&stream);
// copy input from host (CPU) to device (GPU) in stream
cudaMemcpyAsync(buffers[inputIndex], input, batchSize * INPUT_SIZE * sizeof(float), cudaMemcpyHostToDevice, stream);
// execute inference using context provided by engine
context.enqueue(batchSize, buffers, stream, nullptr);
// copy output back from device (GPU) to host (CPU)
cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
stream);
// synchronize the stream to prevent issues
// (block CUDA and wait for CUDA operations to be completed)
cudaStreamSynchronize(stream);
// Release stream and buffers (memory)
cudaStreamDestroy(stream);
cudaFree(buffers[inputIndex]);
cudaFree(buffers[outputIndex]);
}
void performInference() {
/**
* Get inference using the pre-trained model
*/
// stream to write model
char *trtModelStream{nullptr};
size_t size{0};
// read model from the engine file
std::ifstream file("../mlp.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
// create a runtime (required for deserialization of model) with NVIDIA's logger
IRuntime *runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
// deserialize engine for using the char-stream
ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
// create execution context -- required for inference executions
IExecutionContext *context = engine->createExecutionContext();
assert(context != nullptr);
float out[1]; // array for output
float data[1]; // array for input
for (float &i: data)
i = 12.0; // put any value for input
// time the execution
auto start = std::chrono::system_clock::now();
// do inference using the parameters
doInference(*context, data, out, 1);
// time the execution
auto end = std::chrono::system_clock::now();
std::cout << "\n[INFO]: Time taken by execution: "
<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
// free the captured space
context->destroy();
engine->destroy();
runtime->destroy();
std::cout << "\nInput:\t" << data[0];
std::cout << "\nOutput:\t";
for (float i: out) {
std::cout << i;
}
std::cout << std::endl;
}
int checkArgs(int argc, char **argv) {
/**
* Parse command line arguments
*
* @param argc: argument count
* @param argv: arguments vector
* @return int: a flag to perform operation
*/
if (argc != 2) {
std::cerr << "[ERROR]: Arguments not right!" << std::endl;
std::cerr << "./mlp -s // serialize model to plan file" << std::endl;
std::cerr << "./mlp -d // deserialize plan file and run inference" << std::endl;
return -1;
}
if (std::string(argv[1]) == "-s") {
return 1;
} else if (std::string(argv[1]) == "-d") {
return 2;
}
return -1;
}
int main(int argc, char **argv) {
int args = checkArgs(argc, argv);
if (args == 1)
performSerialization();
else if (args == 2)
performInference();
return 0;
}

248
mlp/mlp.py Normal file
View File

@ -0,0 +1,248 @@
import argparse
import os
import numpy as np
import struct
# required for the model creation
import tensorrt as trt
# required for the inference using TRT engine
import pycuda.autoinit
import pycuda.driver as cuda
# Sizes of input and output for TensorRT model
INPUT_SIZE = 1
OUTPUT_SIZE = 1
# path of .wts (weight file) and .engine (model file)
WEIGHT_PATH = "./mlp.wts"
ENGINE_PATH = "./mlp.engine"
# input and output names are must for the TRT model
INPUT_BLOB_NAME = 'data'
OUTPUT_BLOB_NAME = 'out'
# A logger provided by NVIDIA-TRT
gLogger = trt.Logger(trt.Logger.INFO)
################################
# DEPLOYMENT RELATED ###########
################################
def load_weights(file_path):
"""
Parse the .wts file and store weights in dict format
:param file_path:
:return weight_map: dictionary containing weights and their values
"""
print(f"[INFO]: Loading weights: {file_path}")
assert os.path.exists(file_path), '[ERROR]: Unable to load weight file.'
weight_map = {}
with open(file_path, "r") as f:
lines = [line.strip() for line in f]
# count for total # of weights
count = int(lines[0])
assert count == len(lines) - 1
# Loop through counts and get the exact num of values against weights
for i in range(1, count + 1):
splits = lines[i].split(" ")
name = splits[0]
cur_count = int(splits[1])
# len of splits must be greater than current weight counts
assert cur_count + 2 == len(splits)
# loop through all weights and unpack from the hexadecimal values
values = []
for j in range(2, len(splits)):
# hex string to bytes to float
values.append(struct.unpack(">f", bytes.fromhex(splits[j])))
# store in format of { 'weight.name': [weights_val0, weight_val1, ..] }
weight_map[name] = np.array(values, dtype=np.float32)
return weight_map
def create_mlp_engine(max_batch_size, builder, config, dt):
"""
Create Multi-Layer Perceptron using the TRT Builder and Configurations
:param max_batch_size: batch size for built TRT model
:param builder: to build engine and networks
:param config: configuration related to Hardware
:param dt: datatype for model layers
:return engine: TRT model
"""
print("[INFO]: Creating MLP using TensorRT...")
# load weight maps from the file
weight_map = load_weights(WEIGHT_PATH)
# build an empty network using builder
network = builder.create_network()
# add an input to network using the *input-name
data = network.add_input(INPUT_BLOB_NAME, dt, (1, 1, INPUT_SIZE))
assert data
# add the layer with output-size (number of outputs)
linear = network.add_fully_connected(input=data,
num_outputs=OUTPUT_SIZE,
kernel=weight_map['linear.weight'],
bias=weight_map['linear.bias'])
assert linear
# set the name for output layer
linear.get_output(0).name = OUTPUT_BLOB_NAME
# mark this layer as final output layer
network.mark_output(linear.get_output(0))
# set the batch size of current builder
builder.max_batch_size = max_batch_size
# create the engine with model and hardware configs
engine = builder.build_engine(network, config)
# free captured memory
del network
del weight_map
# return engine
return engine
def api_to_model(max_batch_size):
"""
Create engine using TensorRT APIs
:param max_batch_size: for the deployed model configs
:return:
"""
# Create Builder with logger provided by TRT
builder = trt.Builder(gLogger)
# Create configurations from Engine Builder
config = builder.create_builder_config()
# Create MLP Engine
engine = create_mlp_engine(max_batch_size, builder, config, trt.float32)
assert engine
# Write the engine into binary file
print("[INFO]: Writing engine into binary...")
with open(ENGINE_PATH, "wb") as f:
# write serialized model in file
f.write(engine.serialize())
# free the memory
del engine
del builder
################################
# INFERENCE RELATED ############
################################
def perform_inference(input_val):
"""
Get inference using the pre-trained model
:param input_val: a number as an input
:return:
"""
def do_inference(inf_context, inf_host_in, inf_host_out):
"""
Perform inference using the CUDA context
:param inf_context: context created by engine
:param inf_host_in: input from the host
:param inf_host_out: output to save on host
:return:
"""
inference_engine = inf_context.engine
# Input and output bindings are required for inference
assert inference_engine.num_bindings == 2
# allocate memory in GPU using CUDA bindings
device_in = cuda.mem_alloc(inf_host_in.nbytes)
device_out = cuda.mem_alloc(inf_host_out.nbytes)
# create bindings for input and output
bindings = [int(device_in), int(device_out)]
# create CUDA stream for simultaneous CUDA operations
stream = cuda.Stream()
# copy input from host (CPU) to device (GPU) in stream
cuda.memcpy_htod_async(device_in, inf_host_in, stream)
# execute inference using context provided by engine
inf_context.execute_async(bindings=bindings, stream_handle=stream.handle)
# copy output back from device (GPU) to host (CPU)
cuda.memcpy_dtoh_async(inf_host_out, device_out, stream)
# synchronize the stream to prevent issues
# (block CUDA and wait for CUDA operations to be completed)
stream.synchronize()
# create a runtime (required for deserialization of model) with NVIDIA's logger
runtime = trt.Runtime(gLogger)
assert runtime
# read and deserialize engine for inference
with open(ENGINE_PATH, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
assert engine
# create execution context -- required for inference executions
context = engine.create_execution_context()
assert context
# create input as array
data = np.array([input_val], dtype=np.float32)
# capture free memory for input in GPU
host_in = cuda.pagelocked_empty((INPUT_SIZE), dtype=np.float32)
# copy input-array from CPU to Flatten array in GPU
np.copyto(host_in, data.ravel())
# capture free memory for output in GPU
host_out = cuda.pagelocked_empty(OUTPUT_SIZE, dtype=np.float32)
# do inference using required parameters
do_inference(context, host_in, host_out)
print(f'\n[INFO]: Predictions using pre-trained model..\n\tInput:\t{input_val}\n\tOutput:\t{host_out[0]:.4f}')
def get_args():
"""
Parse command line arguments
:return arguments: parsed arguments
"""
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-s', action='store_true')
arg_parser.add_argument('-d', action='store_true')
arguments = vars(arg_parser.parse_args())
# check for the arguments
if not (arguments['s'] ^ arguments['d']):
print("[ERROR]: Arguments not right!\n")
print("\tpython mlp.py -s # serialize model to engine file")
print("\tpython mlp.py -d # deserialize engine file and run inference")
exit()
return arguments
if __name__ == "__main__":
args = get_args()
if args['s']:
api_to_model(max_batch_size=1)
print("[INFO]: Successfully created TensorRT engine...")
print("\n\tRun inference using `python mlp.py -d`\n")
else:
perform_inference(input_val=4.0)

3
mlp/mlp.wts Normal file
View File

@ -0,0 +1,3 @@
2
linear.weight 1 3fff7e32
linear.bias 1 3c138a5a