haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

lenet5: support trt8

Browse Source
This commit is contained in:
wang-xinyu 2022-10-09 11:30:31 +08:00
parent 87f897bced
commit 0c78095dd6
4 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
lenet/lenet.cpp
View File

@ -136,7 +136,7 @@ ICudaEngine* createLenetEngine(unsigned int maxBatchSize, IBuilder* builder, IBu
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(1 << 20);
config->setMaxWorkspaceSize(16 << 20);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// Don't need the network any more

3
lenet/logging.h
View File

@ -25,6 +25,7 @@
#include <ostream>
#include <sstream>
#include <string>
#include "macros.h"
using Severity = nvinfer1::ILogger::Severity;
@ -236,7 +237,7 @@ public:
//! 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
void log(Severity severity, const char* msg) TRT_NOEXCEPT override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}

12
lenet/macros.h Normal file
View File

@ -0,0 +1,12 @@
#ifndef __MACROS_H
#define __MACROS_H
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif
#endif // __MACROS_H

2
tutorials/getting_started.md
View File

@ -44,7 +44,7 @@ lenet out: tensor([[0.0950, 0.0998, 0.1101, 0.0975, 0.0966, 0.1097, 0.0948, 0.10
Clone the wang-xinyu/tensorrtx in your machine. Enter lenet folder, copy lenet5.wts generated above, and cmake&make c++ code.
And of course you should install cuda/cudnn/tensorrt first.
And of course you should install cuda/cudnn/tensorrt first. You might need to adapt the tensorrt path in CMakeLists.txt if you install tensorrt from tar package.
```
git clone https://github.com/wang-xinyu/tensorrtx