Inceptionv4 (#505)

* add: structure inception v4

* fix: camelcase naming

* fix and optimize

* add: doInference def

* define inception layers

* fix: spacing

* add: inception network def

* fix layers

* create net

* validate

* add README

* update readme

* typo

* create dir inception
This commit is contained in:
makaveli 2021-04-26 08:36:52 +05:30 committed by GitHub
parent 046dd5b417
commit 4d630ed599
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 1437 additions and 0 deletions

View File

@ -0,0 +1,35 @@
cmake_minimum_required(VERSION 2.6)
project(InceptionV4)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -pthread -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
file(GLOB SOURCE_FILES "*.h" "*.cpp")
add_executable(inceptionv4 ${SOURCE_FILES})
target_link_libraries(inceptionv4 nvinfer)
target_link_libraries(inceptionv4 cudart)
target_link_libraries(inceptionv4 ${OpenCV_LIBS})
add_definitions(-O2 -pthread)

View File

@ -0,0 +1,36 @@
# Inception v4
Inception v4 model architecture from "Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning" <https://arxiv.org/abs/1602.07261v2>.
For the details, you can refer to [rwightman/pytorch-image-models](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/inception_v4.py)
Following tricks are used in this inception:
- For pooling layer with padding, we need pay attention to see if padding is included or excluded while calculating average number. Pytorch includes padding while doing avgPool by default, but Tensorrt doesn't. So for pooling layer with padding, we need `setAverageCountExcludesPadding(false)` in tensorrt.
- Batchnorm layer, implemented by scale layer.
```
// 1. generate inception.wts from [BlueMirrors/torchtrtz](https://github.com/BlueMirrors/torchtrtz/blob/main/generate_weights.py)
// 2. put inception.wts into tensorrtx/inceptionV4
// 3. build and run
cd tensorrtx/inception/inceptionV4
mkdir build
cd build
cmake ..
make
sudo ./inceptionV4 -s // serialize model to plan file i.e. 'inceptionV4.engine'
sudo ./inceptionV4 -d // deserialize plan file and run inference
// 4. see if the output is same as rwightman/pytorch-image-models/inceptionv4
```

View File

@ -0,0 +1,235 @@
# include "inception_v4.h"
namespace trtx {
InceptionV4::InceptionV4(const InceptionV4Params &params)
: mParams(params)
, mContext(nullptr)
, mEngine(nullptr)
{
}
/**
* Builds the tensorrt engine and serializes it.
**/
bool InceptionV4::serializeEngine()
{
// load weights
weightMap = loadWeights(mParams.weightsFile);
// create builder
IBuilder* builder = createInferBuilder(gLogger);
assert(builder);
// create builder config
IBuilderConfig* config = builder -> createBuilderConfig();
assert(config);
// create engine
bool created = buildEngine(builder, config);
if(!created)
{
std::cerr << "Engine creation failed. Check logs." << std::endl;
return false;
}
// serilaize engine
assert(mEngine != nullptr);
IHostMemory* modelStream{nullptr};
modelStream = mEngine -> serialize();
assert(modelStream != nullptr);
// destroy
config -> destroy();
builder -> destroy();
// write serialized engine to file
std::ofstream trtFile(mParams.trtEngineFile);
if(!trtFile){
std::cerr << "Unable to open engine file." << std::endl;
return false;
}
trtFile.write(reinterpret_cast<const char*>(modelStream -> data()), modelStream -> size());
std::cout << "Engine serialized and saved." << std::endl;
// clean
modelStream -> destroy();
return true;
}
bool InceptionV4::buildEngine(IBuilder *builder, IBuilderConfig *config) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape { 1, 1, 32, 32 } with name INPUT_BLOB_NAME
ITensor* data = network->addInput(mParams.inputTensorName, dt, Dims3{3, mParams.inputH, mParams.inputW});
assert(data);
Weights emptywts{DataType::kFLOAT, nullptr, 0};
float shval[3] = {(0.485 - 0.5) / 0.5, (0.456 - 0.5) / 0.5, (0.406 - 0.5) / 0.5};
float scval[3] = {0.229 / 0.5, 0.224 / 0.5, 0.225 / 0.5};
float pval[3] = {1.0, 1.0, 1.0};
Weights shift{DataType::kFLOAT, shval, 3};
Weights scale{DataType::kFLOAT, scval, 3};
Weights power{DataType::kFLOAT, pval, 3};
IScaleLayer* scale1 = network->addScale(*data, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale1);
IActivationLayer* relu0 = basicConv2d(network, weightMap, *scale1 -> getOutput(0), 32, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, "features.0");
relu0 = basicConv2d(network, weightMap, *relu0 -> getOutput(0), 32, DimsHW{ 3, 3 }, 1, DimsHW{ 0, 0 }, "features.1");
relu0 = basicConv2d(network, weightMap, *relu0 -> getOutput(0), 64, DimsHW{ 3, 3 }, 1, DimsHW{ 1, 1 }, "features.2");
auto cat0 = mixed_3a(network, weightMap, *relu0 -> getOutput(0), "features.3");
cat0 = mixed_4a(network, weightMap, *cat0 -> getOutput(0), "features.4");
cat0 = mixed_5a(network, weightMap, *cat0 -> getOutput(0), "features.5");
cat0 = inceptionA(network, weightMap, *cat0 -> getOutput(0), "features.6");
cat0 = inceptionA(network, weightMap, *cat0 -> getOutput(0), "features.7");
cat0 = inceptionA(network, weightMap, *cat0 -> getOutput(0), "features.8");
cat0 = inceptionA(network, weightMap, *cat0 -> getOutput(0), "features.9");
cat0 = reductionA(network, weightMap, *cat0 -> getOutput(0), "features.10");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.11");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.12");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.13");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.14");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.15");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.16");
cat0 = inceptionB(network, weightMap, *cat0 -> getOutput(0), "features.17");
cat0 = reductionB(network, weightMap, *cat0 -> getOutput(0), "features.18");
cat0 = inceptionC(network, weightMap, *cat0 -> getOutput(0), "features.19");
cat0 = inceptionC(network, weightMap, *cat0 -> getOutput(0), "features.20");
cat0 = inceptionC(network, weightMap, *cat0 -> getOutput(0), "features.21");
IPoolingLayer* pool2 = network->addPoolingNd(*cat0->getOutput(0), PoolingType::kAVERAGE, DimsHW{8, 8});
assert(pool2);
IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool2->getOutput(0), 1000, weightMap["last_linear.weight"], weightMap["last_linear.bias"]);
assert(fc1);
fc1->getOutput(0)->setName(mParams.outputTensorName);
std::cout << "set name out" << std::endl;
network->markOutput(*fc1->getOutput(0));
// Build engine
builder->setMaxBatchSize(mParams.batchSize);
config->setMaxWorkspaceSize(1 << 28);
if (mParams.fp16)
config->setFlag(BuilderFlag::kFP16);
mEngine = builder->buildEngineWithConfig(*network, *config);
std::cout << "build out" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
if (mEngine == nullptr) return false;
return true;
}
bool InceptionV4::deserializeCudaEngine() {
if (mContext != nullptr && mEngine != nullptr)
{
return true;
}
if (mEngine == nullptr)
{
char* trtModelStream{nullptr};
size_t size{0};
// open file
std::ifstream f(mParams.trtEngineFile, std::ios::binary);
if (f.good())
{
// get size
f.seekg(0, f.end);
size = f.tellg();
f.seekg(0, f.beg);
trtModelStream = new char[size];
// read data as a block
f.read(trtModelStream, size);
f.close();
}
if (trtModelStream == nullptr)
{
return false;
}
// deserialize
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime);
mEngine = runtime -> deserializeCudaEngine(trtModelStream, size, 0);
assert(mEngine != nullptr);
// clean up
runtime -> destroy();
delete[] trtModelStream;
}
std::cout << "deserialized engine successfully." << std::endl;
// create execution context
mContext = mEngine -> createExecutionContext();
assert(mContext != nullptr);
return true;
}
void InceptionV4::doInference(float* input, float* output, int batchSize) {
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(mEngine -> 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 = mEngine->getBindingIndex(mParams.inputTensorName);
const int outputIndex = mEngine->getBindingIndex(mParams.outputTensorName);
// Create GPU buffers on device
CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * mParams.inputH * mParams.inputW * sizeof(float)));
CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batchSize * 1000 * sizeof(float)));
// Create stream
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * mParams.inputH * mParams.inputW * sizeof(float), cudaMemcpyHostToDevice, stream));
mContext->enqueue(batchSize, buffers, stream, nullptr);
CUDA_CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * 1000 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CUDA_CHECK(cudaFree(buffers[inputIndex]));
CUDA_CHECK(cudaFree(buffers[outputIndex]));
}
/**
* Cleans up any state created in the InceptionV4Trt class
**/
bool InceptionV4::cleanUp()
{
if (mContext != nullptr)
mContext -> destroy();
if (mEngine != nullptr)
mEngine -> destroy();
return true;
}
}

View File

@ -0,0 +1,58 @@
#ifndef TRTX_INCEPTION_NETWORK_H
#define TRTX_INCEPTION_NETWORK_H
#include <memory>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "logging.h"
#include "utils.h"
#include "layers_api.h"
static Logger gLogger;
using namespace trtxlayers;
namespace trtx {
struct InceptionV4Params
{
/* data */
int32_t batchSize{1}; // Number of inputs in a batch
bool int8{false}; // Allow runnning the network in Int8 mode.
bool fp16{false}; // Allow running the network in FP16 mode.
const char* inputTensorName = "data";
const char* outputTensorName = "prob";
int inputW; // The input width of the network.
int inputH; // The input height of the the network.
int outputSize; // THe output size of the network.
std::string weightsFile; // Weights file filename.
std::string trtEngineFile; // trt engine file name
};
class InceptionV4 {
public:
InceptionV4(const InceptionV4Params &enginecfg);
~InceptionV4() {};
bool serializeEngine(); // create & serialize netowrk Engine
bool deserializeCudaEngine();
void doInference(float* input, float* output, int batchSize);
bool cleanUp();
private:
bool buildEngine(IBuilder *builder, IBuilderConfig *config);
// Runs the Tensorrt network inference engine on a sample.
private:
InceptionV4Params mParams;
ICudaEngine* mEngine; // The tensorrt engine used to run the network.
std::map<std::string, Weights> weightMap; // The weight value map.
IExecutionContext* mContext; // The TensorRT execution context to run inference.
std::string inception;
DataType dt{DataType::kFLOAT};
};
}
#endif

View File

@ -0,0 +1,315 @@
#include "layers_api.h"
namespace trtxlayers {
IScaleLayer* addBatchNorm2d(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname,
float eps
)
{
float *gamma = (float*)weightMap[lname + ".weight"].values;
float *beta = (float*)weightMap[lname + ".bias"].values;
float *mean = (float*)weightMap[lname + ".running_mean"].values;
float *var = (float*)weightMap[lname + ".running_var"].values;
int len = weightMap[lname + ".running_var"].count;
std::cout << "len " << len << std::endl;
float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IActivationLayer* basicConv2d(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
int outch,
DimsHW ksize,
int s,
DimsHW p,
std::string lname
)
{
// empty wts for bias
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// add conv -> bn -> relu
IConvolutionLayer* conv = network -> addConvolutionNd(input, outch, ksize, weightMap[lname + ".conv.weight"], emptywts);
assert(conv);
conv -> setStrideNd(DimsHW{s, s});
conv -> setPaddingNd(p);
IScaleLayer* bn = addBatchNorm2d(network, weightMap, *conv -> getOutput(0), lname + ".bn", 1e-3);
IActivationLayer* relu = network -> addActivation(*bn -> getOutput(0), ActivationType::kRELU);
assert(relu);
return relu;
}
IConcatenationLayer* mixed_3a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// branch 0
IPoolingLayer* pool = network -> addPoolingNd(input, PoolingType::kMAX, DimsHW{3, 3});
assert(pool);
pool -> setStrideNd(DimsHW{2, 2});
// branch 1
IActivationLayer* relu = basicConv2d(network, weightMap, input, 96, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".conv");
// concatenate two branches
ITensor* inputTensors[] = { pool -> getOutput(0), relu -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 2);
assert(cat);
return cat;
}
IConcatenationLayer* mixed_4a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// branch 0
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 64, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch0.0");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 96, DimsHW{ 3, 3 }, 1, DimsHW{ 0, 0 }, lname + ".branch0.1");
// branch 1
IActivationLayer* relu2 = basicConv2d(network, weightMap, input, 64, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch1.0");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 64, DimsHW{ 1, 7 }, 1, DimsHW{ 0, 3 }, lname + ".branch1.1");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 64, DimsHW{ 7, 1 }, 1, DimsHW{ 3, 0 }, lname + ".branch1.2");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 96, DimsHW{ 3, 3 }, 1, DimsHW{ 0, 0 }, lname + ".branch1.3");
// concatenate two branches
ITensor* inputTensors[] = { relu1 -> getOutput(0), relu2 -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 2);
assert(cat);
return cat;
}
IConcatenationLayer* mixed_5a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
std::cout<<"mixed_5a"<<std::endl;
//branch 0
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 192, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".conv");
//branch 1
IPoolingLayer* pool1 = network -> addPoolingNd(input, PoolingType::kMAX, DimsHW{ 3, 3 });
assert(pool1);
pool1 -> setStrideNd(DimsHW{ 2, 2 });
// concatenate branches
ITensor* inputTensors[] = { relu1 -> getOutput(0), pool1 -> getOutput(0)};
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 2);
assert(cat);
std::cout<<"mixed_5a done"<<std::endl;
return cat;
}
IConcatenationLayer* inceptionA(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// branch 0
IActivationLayer* relu0 = basicConv2d(network, weightMap, input, 96, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch0");
// branch 1
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 64, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname +".branch1.0");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 96, DimsHW{ 3, 3 }, 1, DimsHW{ 1, 1 }, lname+".branch1.1");
// branch 2
IActivationLayer* relu2 = basicConv2d(network, weightMap, input, 64, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname+".branch2.0");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 96, DimsHW{ 3, 3 }, 1, DimsHW{ 1, 1 }, lname+".branch2.1");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 96, DimsHW{ 3, 3 }, 1, DimsHW{ 1, 1 }, lname+".branch2.2");
// branch 3
IPoolingLayer* pool1 = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{3, 3});
assert(pool1);
pool1->setStrideNd(DimsHW{1, 1});
pool1->setPaddingNd(DimsHW{1, 1});
pool1->setAverageCountExcludesPadding(false);
IActivationLayer* relu3 = basicConv2d(network, weightMap, *pool1 -> getOutput(0), 96, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname+".branch3.1");
// concatenate all branches outputs
ITensor* inputTensors[] = { relu0 -> getOutput(0), relu1 -> getOutput(0), relu2 -> getOutput(0), relu3 -> getOutput(0)};
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 4);
assert(cat);
return cat;
}
IConcatenationLayer* reductionA(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// features 10 branch 0
IActivationLayer* relu0 = basicConv2d(network, weightMap, input, 384, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".branch0");
// branch 1
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 192, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch1.0");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 224, DimsHW{ 3, 3 }, 1, DimsHW{ 1, 1 }, lname + ".branch1.1");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 256, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".branch1.2");
// branch 2
IPoolingLayer* pool1 = network -> addPoolingNd(input, PoolingType::kMAX, DimsHW{ 3, 3 });
assert(pool1);
pool1 -> setStrideNd(DimsHW{ 2, 2 });
// concatenate
ITensor* inputTensors[] = { relu0 -> getOutput(0), relu1 -> getOutput(0), pool1 -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 3);
assert(cat);
return cat;
}
IConcatenationLayer* inceptionB(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// features 11 branch 0
IActivationLayer* relu0 = basicConv2d(network, weightMap, input, 384, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch0");
// branch 1
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 192, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch1.0");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 224, DimsHW{ 1, 7 }, 1, DimsHW{ 0, 3 }, lname + ".branch1.1");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 256, DimsHW{ 7, 1 }, 1, DimsHW{ 3, 0 }, lname + ".branch1.2");
// branch 2
IActivationLayer* relu2 = basicConv2d(network, weightMap, input, 192, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch2.0");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 192, DimsHW{ 7, 1 }, 1, DimsHW{ 3, 0 }, lname + ".branch2.1");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 224, DimsHW{ 1, 7 }, 1, DimsHW{ 0, 3 }, lname + ".branch2.2");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 224, DimsHW{ 7, 1 }, 1, DimsHW{ 3, 0 }, lname + ".branch2.3");
relu2 = basicConv2d(network, weightMap, *relu2 -> getOutput(0), 256, DimsHW{ 1, 7 }, 1, DimsHW{ 0, 3 }, lname + ".branch2.4");
// branch 3
IPoolingLayer* pool0 = network -> addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{ 3, 3 });
assert(pool0);
pool0 -> setStrideNd(DimsHW{ 1, 1 });
pool0 -> setPaddingNd(DimsHW{ 1, 1 });
pool0 -> setAverageCountExcludesPadding(false);
IActivationLayer* relu3 = basicConv2d(network, weightMap, *pool0 -> getOutput(0), 128, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch3.1");
// concatenate branches
ITensor* inputTensors[] = { relu0 -> getOutput(0), relu1 -> getOutput(0), relu2 -> getOutput(0), relu3 -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 4);
assert(cat);
return cat;
}
IConcatenationLayer* reductionB(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// features 18 branch 0
IActivationLayer* relu0 = basicConv2d(network, weightMap, input, 192, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch0.0");
relu0 = basicConv2d(network, weightMap, *relu0 -> getOutput(0), 192, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".branch0.1");
// branch 1
IActivationLayer* relu1 = basicConv2d(network, weightMap, input, 256, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch1.0");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 256, DimsHW{ 1, 7 }, 1, DimsHW{ 0, 3 }, lname + ".branch1.1");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 320, DimsHW{ 7, 1 }, 1, DimsHW{ 3, 0 }, lname + ".branch1.2");
relu1 = basicConv2d(network, weightMap, *relu1 -> getOutput(0), 320, DimsHW{ 3, 3 }, 2, DimsHW{ 0, 0 }, lname + ".branch1.3");
// branch 2
IPoolingLayer* pool1 = network -> addPoolingNd(input, PoolingType::kMAX, DimsHW{ 3, 3 });
assert(pool1);
pool1 -> setStrideNd(DimsHW{ 2, 2 });
// concatenate
ITensor* inputTensors[] = { relu0 -> getOutput(0), relu1 -> getOutput(0), pool1 -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 3);
assert(cat);
return cat;
}
IConcatenationLayer* inceptionC(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
)
{
// features 19 branch 0
IActivationLayer* relu0 = basicConv2d(network, weightMap, input, 256, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch0");
// branch 1
IActivationLayer* relu1_0 = basicConv2d(network, weightMap, input, 384, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch1_0");
IActivationLayer* relu1_1a = basicConv2d(network, weightMap, *relu1_0 -> getOutput(0), 256, DimsHW{ 1, 3 }, 1, DimsHW{ 0, 1 }, lname + ".branch1_1a");
IActivationLayer* relu1_1b = basicConv2d(network, weightMap, *relu1_0 -> getOutput(0), 256, DimsHW{ 3, 1 }, 1, DimsHW{ 1, 0 }, lname + ".branch1_1b");
ITensor* inputTensors1[] = { relu1_1a -> getOutput(0), relu1_1b -> getOutput(0) };
IConcatenationLayer* cat1 = network -> addConcatenation(inputTensors1, 2);
assert(cat1);
// branch 2
IActivationLayer* relu2_0 = basicConv2d(network, weightMap, input, 384, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch2_0");
IActivationLayer* relu2_1 = basicConv2d(network, weightMap, *relu2_0 -> getOutput(0), 448, DimsHW{ 3, 1 }, 1, DimsHW{ 1, 0 }, lname + ".branch2_1");
IActivationLayer* relu2_2 = basicConv2d(network, weightMap, *relu2_1 -> getOutput(0), 512, DimsHW{ 1, 3 }, 1, DimsHW{ 0, 1 }, lname + ".branch2_2");
IActivationLayer* relu2_3a = basicConv2d(network, weightMap, *relu2_2 -> getOutput(0), 256, DimsHW{ 1, 3 }, 1, DimsHW{ 0, 1 }, lname + ".branch2_3a");
IActivationLayer* relu2_3b = basicConv2d(network, weightMap, *relu2_2 -> getOutput(0), 256, DimsHW{ 3, 1 }, 1, DimsHW{ 1, 0 }, lname + ".branch2_3b");
ITensor* inputTensors2[] = { relu2_3a -> getOutput(0), relu2_3b -> getOutput(0) };
IConcatenationLayer* cat2 = network -> addConcatenation(inputTensors2, 2);
assert(cat2);
// branch 3
IPoolingLayer* pool3 = network -> addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{ 3, 3 });
assert(pool3);
pool3 -> setStrideNd(DimsHW{ 1, 1 });
pool3 -> setPaddingNd(DimsHW{ 1, 1 });
pool3 -> setAverageCountExcludesPadding(false);
IActivationLayer* relu3 = basicConv2d(network, weightMap, *pool3 -> getOutput(0), 256, DimsHW{ 1, 1 }, 1, DimsHW{ 0, 0 }, lname + ".branch3.1");
// concatenate
ITensor* inputTensors[] = { relu0 -> getOutput(0), cat1 -> getOutput(0), cat2 -> getOutput(0), relu3 -> getOutput(0) };
IConcatenationLayer* cat = network -> addConcatenation(inputTensors, 4);
assert(cat);
return cat;
}
}

View File

@ -0,0 +1,93 @@
#ifndef TRTX_LAYERS_API_H
#define TRTX_LAYERS_API_H
#include <map>
#include <math.h>
#include <assert.h>
#include <iostream>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
using namespace nvinfer1;
namespace trtxlayers {
// Declare your layers here
IScaleLayer* addBatchNorm2d(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname,
float eps
);
IActivationLayer* basicConv2d(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
int outch,
DimsHW ksize,
int s,
DimsHW p,
std::string lname
);
IConcatenationLayer* mixed_3a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* mixed_4a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* mixed_5a(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* inceptionA(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* reductionA(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* inceptionB(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* reductionB(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
IConcatenationLayer* inceptionC(
INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
std::string lname
);
}
#endif // TRTX_LAYERS_API_H

View File

@ -0,0 +1,507 @@
/*
* Copyright (c) 2021, 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)
, mPrefix(other.mPrefix)
, mShouldLog(other.mShouldLog)
{
}
~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) 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

View File

@ -0,0 +1,88 @@
#include "inception_v4.h"
/**
* Initializes Inception class params in the
* InceptionV4Params structure.
**/
trtx::InceptionV4Params initializeParams()
{
trtx::InceptionV4Params params;
params.batchSize = 1;
params.fp16 = false;
params.inputH = 299;
params.inputW = 299;
params.outputSize = 1000;
// change weights file name here
params.weightsFile = "../inceptionV4.wts";
// change engine file name here
params.trtEngineFile = "inceptionV4.engine";
return params;
}
int main(int argc, char** argv){
if (argc != 2) {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./inception -s // serialize model to plan file" << std::endl;
std::cerr << "./inception -d // deserialize plan file and run inference" << std::endl;
return -1;
}
trtx::InceptionV4Params params = initializeParams();
trtx::InceptionV4 inceptionV4(params);
if (std::string(argv[1]) == "-s") {
// check if engine exists already
std::ifstream f(params.trtEngineFile, std::ios::binary);
// if engine does not exists build, serialize and save
if(!f.good())
{
std::cout << "Building network ..." << std::endl;
f.close();
inceptionV4.serializeEngine();
}
return 1;
}
else if(std::string(argv[1]) == "-d")
{
// deserialize
inceptionV4.deserializeCudaEngine();
}
// create data
float data[3 * params.inputH * params.inputW];
for(int i=0; i<3*params.inputH*params.inputW; i++)
{
data[i] = 1.0;
}
// run inference
float prob[params.outputSize];
for(int i=0; i<100; i++)
{
auto start = std::chrono::system_clock::now();
inceptionV4.doInference(data, prob, 1);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
// cleanup
bool cleaned = inceptionV4.cleanUp();
std::cout << "\nOutput:\n\n";
for (unsigned int i = 0; i < params.outputSize; i++)
{
std::cout << prob[i] << ", ";
if (i % 10 == 0) std::cout << i / 10 << std::endl;
}
std::cout << std::endl;
return 0;
}

View File

@ -0,0 +1,43 @@
# include "utils.h"
// Load weights from files.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string input) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}

View File

@ -0,0 +1,27 @@
# ifndef TRTX_UTILS_H
# define TRTX_UTILS_H
#include <map>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "assert.h"
#include <fstream>
#include <iostream>
#include <memory>
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr)\
{\
cudaError_t error_code = callstr;\
if (error_code != cudaSuccess) {\
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
assert(0);\
}\
}
#endif // CUDA_CHECK
using namespace nvinfer1;
std::map<std::string, Weights> loadWeights(const std::string input);
#endif // TRTX_UTILS_H