資源簡介
Nvidia TensorRT官方例程源代碼,從TX1上拷貝下來的。
代碼片段和文件信息
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?
#include?“NvInfer.h“
#include?“NvCaffeParser.h“
using?namespace?nvinfer1;
using?namespace?nvcaffeparser1;
#define?CHECK(status) \
{ \
if?(status?!=?0) \
{ \
std::cout?<“Cuda?failure:?“?< abort(); \
} \
}
struct?Params
{
std::string?deployFile?modelFile?engine;
std::vector?outputs;
int?device{?0?}?batchSize{?1?}?workspaceSize{?16?}?iterations{?10?}?avgRuns{?10?};
bool?half2{?false?}?verbose{?false?}?hostTime{?false?};
}?gParams;
std::vector?gInputs;
//?Logger?for?GIE?info/warning/errors
class?Logger?:?public?ILogger
{
void?log(Severity?severity?const?char*?msg)?override
{
//?suppress?info-level?messages
if?(severity?!=?Severity::kINFO?||?gParams.verbose)
std::cout?< }
}?gLogger;
ICudaEngine*?caffeToGIEModel()
{
//?create?the?builder
IBuilder*?builder?=?createInferBuilder(gLogger);
//?parse?the?caffe?model?to?populate?the?network?then?set?the?outputs
INetworkDefinition*?network?=?builder->createNetwork();
ICaffeParser*?parser?=?createCaffeParser();
const?IBlobNameToTensor*?blobNameToTensor?=?parser->parse(gParams.deployFile.c_str()
??gParams.modelFile.c_str()
??*network
??gParams.half2???DataType::kHALF:DataType::kFLOAT);
if?(!blobNameToTensor)
return?nullptr;
for?(int?i?=?0?n?=?network->getNbInputs();?i? gInputs.push_back(network->getInput(i)->getName());
//?specify?which?tensors?are?outputs
for?(auto&?s?:?gParams.outputs)
{
if?(blobNameToTensor->find(s.c_str())?==?nullptr)
{
std::cout?<“could?not?find?output?blob?“?< return?nullptr;
}
network->markOutput(*blobNameToTensor->find(s.c_str()));
}
//?Build?the?engine
builder->setMaxBatchSize(gParams.batchSize);
builder->setMaxWorkspaceSize(gParams.workspaceSize<<20);
builder->setHalf2Mode(gParams.half2);
ICudaEngine*?engine?=?builder->buildCudaEngine(*network);
if?(engine?==?nullptr)
std::cout?<“could?not?build?engine“?<
parser->destroy();
network->destroy();
builder->destroy();
shutdownProtobufLibrary();
return?engine;
}
void?createMemory(const?ICudaEngine&?engine?std::vector&?buffers?const?std::string&?name)
{
size_t?bindingIndex?=?engine.getBindingIndex(name.c_str());
assert(bindingIndex? Dims3?dimensions?=?engine.getBindingDimensions((int)bindingIndex);
size_t?eltCount?=?dimensions.c*dimensions.h*dimensions.w*gParams.batchSize?memSize?=?eltCount?*?sizeof(float);
float*?localMem?=?new?float[eltCount];
for?(size_t?i?=?0;?i? localMem[i]?=?(float(rand())?/?RAND_MAX)?*?2?-?1;
void*?deviceMem;
CHECK(cudaMalloc(&
?屬性????????????大小?????日期????時間???名稱
-----------?---------??----------?-----??----
?????目錄???????????0??2017-02-20?16:27??data\
?????目錄???????????0??2017-02-20?16:27??data\samples\
?????目錄???????????0??2017-02-20?16:27??data\samples\googlenet\
?????文件????41481639??2016-08-12?09:31??data\samples\googlenet\googlenet.caffemodel
?????文件???????36261??2016-08-12?09:31??data\samples\googlenet\googlenet.prototxt
?????目錄???????????0??2017-02-20?16:27??data\samples\mnist\
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\0.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\1.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\2.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\3.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\4.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\5.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\6.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\7.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\8.pgm
?????文件?????????797??2016-08-12?09:31??data\samples\mnist\9.pgm
?????文件?????1725135??2016-08-12?09:31??data\samples\mnist\mnist.caffemodel
?????文件????????1806??2016-08-12?09:31??data\samples\mnist\mnist.prototxt
?????文件????????3147??2016-08-12?09:31??data\samples\mnist\mnist_mean.binaryproto
?????文件?????7685475??2016-08-12?09:31??data\samples\mnist\mnistgie.wts
?????目錄???????????0??2017-02-20?16:27??giexec\
?????文件????????8655??2016-08-12?09:31??giexec\giexec.cpp
?????文件???????10099??2016-08-12?09:31??giexec\giexec.vcxproj
?????文件?????????116??2016-08-12?09:31??giexec\Makefile
?????文件????????3489??2016-08-12?09:31??giexec\Makefile.giexec
?????目錄???????????0??2017-02-20?16:27??sampleGoogleNet\
?????文件?????????136??2016-08-12?09:31??sampleGoogleNet\Makefile
?????文件????????3489??2016-08-12?09:31??sampleGoogleNet\Makefile.sample_googlenet
?????文件????????6092??2016-08-12?09:31??sampleGoogleNet\sampleGoogleNet.cpp
?????文件???????10126??2016-08-12?09:31??sampleGoogleNet\sampleGoogleNet.vcxproj
?????目錄???????????0??2017-02-20?16:27??sampleMNIST\
............此處省略9個文件信息
- 上一篇:從零開始學微信小程序開發
- 下一篇:操作系統概念-第七版中文
評論
共有 條評論