mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Yolo3Detection.cpp doesn't build yet,issues with dependecies of preprocessing and postprocessing on IPluginFactory

Browse Source
This commit is contained in:
perseusdg
2021-08-21 22:29:36 +05:30
parent eba78e7e78
commit 3d8b1ac494
7 changed files with 65 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/pluginsRT/YoloRT.h
+4
View File
@@ -2,6 +2,7 @@
#include "../kernels.h" #include "../kernels.h"
#define YOLORT_CLASSNAME_W 256 #define YOLORT_CLASSNAME_W 256
class YoloRT : public IPluginV2 { class YoloRT : public IPluginV2 {
public: public:
@@ -49,6 +50,7 @@ public:
classesNames[1] = std::string(tmp); classesNames[1] = std::string(tmp);
} }
assert(buf == bufCheck + length); assert(buf == bufCheck + length);
} }
~YoloRT() { ~YoloRT() {
@@ -56,6 +58,7 @@ public:
} }
int getNbOutputs() const NOEXCEPT override { int getNbOutputs() const NOEXCEPT override {
return 1; return 1;
} }
@@ -186,6 +189,7 @@ public:
float nms_thresh; float nms_thresh;
int nms_kind; int nms_kind;
int new_coords; int new_coords;
int NUM=0;
std::vector<std::string> classesNames; std::vector<std::string> classesNames;
dnnType *mask; dnnType *mask;
include/tkDNN/utils.h
+2
View File
@@ -147,4 +147,6 @@ static inline bool isCudaPointer(void *data) {
cudaPointerAttributes attr; cudaPointerAttributes attr;
return cudaPointerGetAttributes(&attr, data) == 0; return cudaPointerGetAttributes(&attr, data) == 0;
} }
#endif //UTILS_H #endif //UTILS_H
src/Int8BatchStream.cpp
+10 -10
View File
@@ -8,14 +8,14 @@
BatchStream::BatchStream(tk::dnn::dataDim_t dim, int batchSize, int maxBatches, const std::string& fileimglist, const std::string& filelabellist) { BatchStream::BatchStream(tk::dnn::dataDim_t dim, int batchSize, int maxBatches, const std::string& fileimglist, const std::string& filelabellist) {
mBatchSize = batchSize; mBatchSize = batchSize;
mMaxBatches = maxBatches; mMaxBatches = maxBatches;
mDims = nvinfer1::DimsNCHW{ dim.n, dim.c, dim.h, dim.w }; mDims = nvinfer1::Dims4{ dim.n, dim.c, dim.h, dim.w };
mHeight = dim.h; mHeight = dim.h;
mWidth = dim.w; mWidth = dim.w;
mImageSize = mDims.c()*mDims.h()*mDims.w(); mImageSize = mDims.d[1]*mDims.d[2]*mDims.d[3];
mBatch.resize(mBatchSize*mImageSize, 0); mBatch.resize(mBatchSize*mImageSize, 0);
mLabels.resize(mBatchSize, 0); mLabels.resize(mBatchSize, 0);
mFileBatch.resize(mDims.n()*mImageSize, 0); mFileBatch.resize(mDims.d[0]*mImageSize, 0);
mFileLabels.resize(mDims.n(), 0); mFileLabels.resize(mDims.d[0], 0);
mFileImgList = fileimglist; mFileImgList = fileimglist;
readInListFile(fileimglist, mListImg); readInListFile(fileimglist, mListImg);
mFileLabelList = filelabellist; mFileLabelList = filelabellist;
@@ -27,7 +27,7 @@ BatchStream::BatchStream(tk::dnn::dataDim_t dim, int batchSize, int maxBatches,
void BatchStream::reset(int firstBatch) { void BatchStream::reset(int firstBatch) {
mBatchCount = 0; mBatchCount = 0;
mFileCount = 0; mFileCount = 0;
mFileBatchPos = mDims.n(); mFileBatchPos = mDims.d[0];
skip(firstBatch); skip(firstBatch);
} }
@@ -37,11 +37,11 @@ bool BatchStream::next() {
return false; return false;
for (int csize = 1, batchPos = 0; batchPos < mBatchSize; batchPos += csize, mFileBatchPos += csize) { for (int csize = 1, batchPos = 0; batchPos < mBatchSize; batchPos += csize, mFileBatchPos += csize) {
assert(mFileBatchPos > 0 && mFileBatchPos <= mDims.n()); assert(mFileBatchPos > 0 && mFileBatchPos <= mDims.d[0]);
if (mFileBatchPos == mDims.n() && !update()) if (mFileBatchPos == mDims.d[0] && !update())
return false; return false;
csize = std::min(mBatchSize - batchPos, mDims.n() - mFileBatchPos); csize = std::min(mBatchSize - batchPos, mDims.d[0] - mFileBatchPos);
std::copy_n(getFileBatch() + mFileBatchPos * mImageSize, csize * mImageSize, getBatch() + batchPos * mImageSize); std::copy_n(getFileBatch() + mFileBatchPos * mImageSize, csize * mImageSize, getBatch() + batchPos * mImageSize);
std::copy_n(getFileLabels() + mFileBatchPos, csize, getLabels() + batchPos); std::copy_n(getFileLabels() + mFileBatchPos, csize, getLabels() + batchPos);
} }
@@ -50,8 +50,8 @@ bool BatchStream::next() {
} }
void BatchStream::skip(int skipCount) { void BatchStream::skip(int skipCount) {
if (mBatchSize >= mDims.n() && mBatchSize%mDims.n() == 0 && mFileBatchPos == mDims.n()) { if (mBatchSize >= mDims.d[0] && mBatchSize%mDims.d[0] == 0 && mFileBatchPos == mDims.d[0]) {
mFileCount += skipCount * mBatchSize / mDims.n(); mFileCount += skipCount * mBatchSize / mDims.d[0];
return; return;
} }
src/Int8Calibrator.cpp
+4 -4
View File
@@ -9,12 +9,12 @@ Int8EntropyCalibrator::Int8EntropyCalibrator(BatchStream& stream, int firstBatch
mInputBlobName(inputBlobName.c_str()), mInputBlobName(inputBlobName.c_str()),
mReadCache(readCache) { mReadCache(readCache) {
nvinfer1::Dims4 dims = mStream.getDims(); nvinfer1::Dims4 dims = mStream.getDims();
mInputCount = mStream.getBatchSize() * dims.c() * dims.h() * dims.w(); mInputCount = mStream.getBatchSize() + dims.d[1]*dims.d[2]*dims.d[3];
checkCuda(cudaMalloc(&mDeviceInput, mInputCount * sizeof(float))); checkCuda(cudaMalloc(&mDeviceInput, mInputCount * sizeof(float)));
mStream.reset(firstBatch); mStream.reset(firstBatch);
} }
bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings) { bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings) NOEXCEPT {
if (!mStream.next()) if (!mStream.next())
return false; return false;
@@ -24,7 +24,7 @@ bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int
return true; return true;
} }
const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) { const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) NOEXCEPT {
mCalibrationCache.clear(); mCalibrationCache.clear();
assert(!mCalibTableFilePath.empty()); assert(!mCalibTableFilePath.empty());
std::ifstream input(mCalibTableFilePath, std::ios::binary); std::ifstream input(mCalibTableFilePath, std::ios::binary);
@@ -38,7 +38,7 @@ const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) {
return length ? &mCalibrationCache[0] : nullptr; return length ? &mCalibrationCache[0] : nullptr;
} }
void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length) { void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length) NOEXCEPT {
assert(!mCalibTableFilePath.empty()); assert(!mCalibTableFilePath.empty());
std::ofstream output(mCalibTableFilePath, std::ios::binary); std::ofstream output(mCalibTableFilePath, std::ios::binary);
output.write(reinterpret_cast<const char*>(cache), length); output.write(reinterpret_cast<const char*>(cache), length);
src/NetworkRT.cpp
+39 -41
View File
@@ -15,7 +15,7 @@ using namespace nvinfer1;
// Logger for info/warning/errors // Logger for info/warning/errors
class Logger : public ILogger { class Logger : public ILogger {
void log(Severity severity, const char* msg) override { void log(Severity severity, const char* msg) NOEXCEPT override {
#ifdef DEBUG #ifdef DEBUG
std::cout <<"TENSORRT LOG: "<< msg << std::endl; std::cout <<"TENSORRT LOG: "<< msg << std::endl;
#endif #endif
@@ -39,7 +39,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
#if NV_TENSORRT_MAJOR >= 5 #if NV_TENSORRT_MAJOR >= 5
std::cout<<"DLAs: "<<builderRT->getNbDLACores()<<"\n"; std::cout<<"DLAs: "<<builderRT->getNbDLACores()<<"\n";
#endif #endif
networkRT = builderRT->createNetwork(); networkRT = builderRT->createNetworkV2(0U);
#if NV_TENSORRT_MAJOR >= 6 #if NV_TENSORRT_MAJOR >= 6
configRT = builderRT->createBuilderConfig(); configRT = builderRT->createBuilderConfig();
#endif #endif
@@ -59,22 +59,21 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
dtRT = DataType::kFLOAT; dtRT = DataType::kFLOAT;
builderRT->setMaxBatchSize(net->maxBatchSize); builderRT->setMaxBatchSize(net->maxBatchSize);
builderRT->setMaxWorkspaceSize(1 << 30); configRT->setMaxWorkspaceSize(1 << 30);
if(net->fp16 && builderRT->platformHasFastFp16()) { if(net->fp16 && builderRT->platformHasFastFp16()) {
dtRT = DataType::kHALF; dtRT = DataType::kHALF;
builderRT->setHalf2Mode(true); #if NV_TENSORRT_MAJOR >= 6
#if NV_TENSORRT_MAJOR >= 6
configRT->setFlag(BuilderFlag::kFP16); configRT->setFlag(BuilderFlag::kFP16);
#endif #endif
} }
#if NV_TENSORRT_MAJOR >= 5 #if NV_TENSORRT_MAJOR >= 5
if(net->dla && builderRT->getNbDLACores() > 0) { if(net->dla && builderRT->getNbDLACores() > 0) {
dtRT = DataType::kHALF; dtRT = DataType::kHALF;
builderRT->setFp16Mode(true); configRT->setFlag(BuilderFlag::kFP16);
builderRT->allowGPUFallback(true); configRT->setFlag(BuilderFlag::kGPU_FALLBACK);
builderRT->setDefaultDeviceType(DeviceType::kDLA); configRT->setDefaultDeviceType(DeviceType::kDLA);
builderRT->setDLACore(0); configRT->setDLACore(0);
} }
#endif #endif
#if NV_TENSORRT_MAJOR >= 6 #if NV_TENSORRT_MAJOR >= 6
@@ -104,7 +103,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
// add input layer // add input layer
ITensor *input = networkRT->addInput("data", DataType::kFLOAT, ITensor *input = networkRT->addInput("data", DataType::kFLOAT,
DimsCHW{ dim.c, dim.h, dim.w}); Dims3{ dim.c, dim.h, dim.w});
checkNULL(input); checkNULL(input);
//add other layers //add other layers
@@ -368,8 +367,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX_FIXEDSIZE) if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX_FIXEDSIZE)
{ {
IPlugin *plugin = new MaxPoolFixedSizeRT(l->output_dim.c, l->output_dim.h, l->output_dim.w, l->output_dim.n, l->strideH, l->strideW, l->winH, l->winH-1); IPluginV2 *plugin = new MaxPoolFixedSizeRT(l->output_dim.c, l->output_dim.h, l->output_dim.w, l->output_dim.n, l->strideH, l->strideW, l->winH, l->winH-1);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -413,20 +412,20 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
return lRT; return lRT;
} }
else if(l->act_mode == CUDNN_ACTIVATION_CLIPPED_RELU) { else if(l->act_mode == CUDNN_ACTIVATION_CLIPPED_RELU) {
IPlugin *plugin = new ActivationReLUCeiling(l->ceiling); IPluginV2 *plugin = new ActivationReLUCeiling(l->ceiling);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
else if(l->act_mode == ACTIVATION_MISH) { else if(l->act_mode == ACTIVATION_MISH) {
IPlugin *plugin = new ActivationMishRT(); IPluginV2 *plugin = new ActivationMishRT();
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
else if(l->act_mode == ACTIVATION_LOGISTIC) { else if(l->act_mode == ACTIVATION_LOGISTIC) {
IPlugin *plugin = new ActivationLogisticRT(); IPluginV2 *plugin = new ActivationLogisticRT();
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -460,8 +459,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
} }
if(l->groups > 1){ if(l->groups > 1){
IPlugin *plugin = new RouteRT(l->groups, l->group_id); IPluginV2 *plugin = new RouteRT(l->groups, l->group_id);
IPluginLayer *lRT = networkRT->addPlugin(tens, l->layers_n, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(tens, l->layers_n, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -472,8 +471,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) { ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) {
IPlugin *plugin = new FlattenConcatRT(); IPluginV2 *plugin = new FlattenConcatRT();
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -481,8 +480,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) {
ILayer* NetworkRT::convert_layer(ITensor *input, Reshape *l) { ILayer* NetworkRT::convert_layer(ITensor *input, Reshape *l) {
// std::cout<<"convert Reshape\n"; // std::cout<<"convert Reshape\n";
IPlugin *plugin = new ReshapeRT(l->output_dim); IPluginV2 *plugin = new ReshapeRT(l->output_dim);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -494,7 +493,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Resize *l) {
checkNULL(lRT); checkNULL(lRT);
Dims d{}; Dims d{};
lRT->setResizeMode(ResizeMode(l->mode)); lRT->setResizeMode(ResizeMode(l->mode));
lRT->setOutputDimensions(DimsCHW{l->output_dim.c, l->output_dim.h, l->output_dim.w}); lRT->setOutputDimensions(Dims3{l->output_dim.c, l->output_dim.h, l->output_dim.w});
return lRT; return lRT;
} }
@@ -502,8 +501,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
//std::cout<<"convert Reorg\n"; //std::cout<<"convert Reorg\n";
//std::cout<<"New plugin REORG\n"; //std::cout<<"New plugin REORG\n";
IPlugin *plugin = new ReorgRT(l->stride); IPluginV2 *plugin = new ReorgRT(l->stride);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -512,8 +511,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Region *l) {
//std::cout<<"convert Region\n"; //std::cout<<"convert Region\n";
//std::cout<<"New plugin REGION\n"; //std::cout<<"New plugin REGION\n";
IPlugin *plugin = new RegionRT(l->classes, l->coords, l->num); IPluginV2 *plugin = new RegionRT(l->classes, l->coords, l->num);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -534,11 +533,11 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
else else
{ {
// plugin version // plugin version
IPlugin *plugin = new ShortcutRT(l->backLayer->output_dim, l->mul); IPluginV2 *plugin = new ShortcutRT(l->backLayer->output_dim, l->mul);
ITensor **inputs = new ITensor*[2]; ITensor **inputs = new ITensor*[2];
inputs[0] = input; inputs[0] = input;
inputs[1] = back_tens; inputs[1] = back_tens;
IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(inputs, 2, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -548,8 +547,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Yolo *l) {
//std::cout<<"convert Yolo\n"; //std::cout<<"convert Yolo\n";
//std::cout<<"New plugin YOLO\n"; //std::cout<<"New plugin YOLO\n";
IPlugin *plugin = new YoloRT(l->classes, l->num, l, l->n_masks, l->scaleXY, l->nms_thresh, l->nsm_kind, l->new_coords); IPluginV2 *plugin = new YoloRT(l->classes, l->num, l, l->n_masks, l->scaleXY, l->nms_thresh, l->nsm_kind, l->new_coords);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -558,8 +557,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Upsample *l) {
//std::cout<<"convert Upsample\n"; //std::cout<<"convert Upsample\n";
//std::cout<<"New plugin UPSAMPLE\n"; //std::cout<<"New plugin UPSAMPLE\n";
IPlugin *plugin = new UpsampleRT(l->stride); IPluginV2 *plugin = new UpsampleRT(l->stride);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT); checkNULL(lRT);
return lRT; return lRT;
} }
@@ -574,10 +573,10 @@ ILayer* NetworkRT::convert_layer(ITensor *input, DeformConv2d *l) {
inputs[1] = preconv->getOutput(0); inputs[1] = preconv->getOutput(0);
//std::cout<<"New plugin DEFORMABLE\n"; //std::cout<<"New plugin DEFORMABLE\n";
IPlugin *plugin = new DeformableConvRT(l->chunk_dim, l->kernelH, l->kernelW, l->strideH, l->strideW, l->paddingH, l->paddingW, IPluginV2 *plugin = new DeformableConvRT(l->chunk_dim, l->kernelH, l->kernelW, l->strideH, l->strideW, l->paddingH, l->paddingW,
l->deformableGroup, l->input_dim.n, l->input_dim.c, l->input_dim.h, l->input_dim.w, l->deformableGroup, l->input_dim.n, l->input_dim.c, l->input_dim.h, l->input_dim.w,
l->output_dim.n, l->output_dim.c, l->output_dim.h, l->output_dim.w, l); l->output_dim.n, l->output_dim.c, l->output_dim.h, l->output_dim.w, l);
IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin); IPluginV2Layer *lRT = networkRT->addPluginV2(inputs, 2, *plugin);
checkNULL(lRT); checkNULL(lRT);
lRT->setName( ("Deformable" + std::to_string(l->id)).c_str() ); lRT->setName( ("Deformable" + std::to_string(l->id)).c_str() );
delete[](inputs); delete[](inputs);
@@ -646,16 +645,15 @@ bool NetworkRT::deserialize(const char *filename) {
file.close(); file.close();
} }
pluginFactory = new PluginFactory();
runtimeRT = createInferRuntime(loggerRT); runtimeRT = createInferRuntime(loggerRT);
engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory); engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size);
//if (gieModelStream) delete [] gieModelStream; //if (gieModelStream) delete [] gieModelStream;
return true; return true;
} }
/*
IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) { IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) {
const char * buf = reinterpret_cast<const char*>(serialData),*bufCheck = buf; const char * buf = reinterpret_cast<const char*>(serialData),*bufCheck = buf;
@@ -897,5 +895,5 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
FatalError("Cant deserialize Plugin"); FatalError("Cant deserialize Plugin");
return NULL; return NULL;
} }
*/
}} }}
src/Yolo3Detection.cpp
+1
View File
@@ -14,6 +14,7 @@ bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes, c
tk::dnn::dataDim_t idim = netRT->input_dim; tk::dnn::dataDim_t idim = netRT->input_dim;
idim.n = nBatches; idim.n = nBatches;
if(netRT->pluginFactory->n_yolos < 2 ) { if(netRT->pluginFactory->n_yolos < 2 ) {
FatalError("this is not yolo3"); FatalError("this is not yolo3");
} }
tests/mnist/test_mnistRT.cpp
+5 -5
View File
@@ -15,7 +15,7 @@ using namespace nvinfer1;
// Logger for info/warning/errors // Logger for info/warning/errors
class Logger : public ILogger class Logger : public ILogger
{ {
void log(Severity severity, const char* msg) override void log(Severity severity, const char* msg) NOEXCEPT override
{ {
// suppress info-level messages // suppress info-level messages
if (severity != Severity::kINFO) if (severity != Severity::kINFO)
@@ -66,11 +66,11 @@ int main() {
std::cout<<"\n==== TensorRT ====\n"; std::cout<<"\n==== TensorRT ====\n";
// create the builder // create the builder
IBuilder* builder = nvinfer1::createInferBuilder(gLogger); IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork(); IBuilderConfig* config = builder->createBuilderConfig();
INetworkDefinition* network = builder->createNetworkV2(0U);
DataType dt = DataType::kFLOAT; DataType dt = DataType::kFLOAT;
// Create input of shape { 1, 1, 28, 28 } with name referenced by "data" // Create input of shape { 1, 1, 28, 28 } with name referenced by "data"
auto input = network->addInput("data", dt, DimsCHW{ 1, 28, 28}); auto input = network->addInput("data", dt, Dims3{ 1, 28, 28});
assert(input != nullptr); assert(input != nullptr);
tk::dnn::Conv2d *c0 = &l0; tk::dnn::Conv2d *c0 = &l0;
@@ -126,7 +126,7 @@ int main() {
// Build the engine // Build the engine
builder->setMaxBatchSize(1); builder->setMaxBatchSize(1);
builder->setMaxWorkspaceSize(1 << 20); config->setMaxWorkspaceSize(1 << 20);
auto engine = builder->buildCudaEngine(*network); auto engine = builder->buildCudaEngine(*network);
// we don't need the network any more // we don't need the network any more