资源简介
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: “ << status; \
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 << msg << std::endl;
}
} 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 < n; 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 “ << s << std::endl;
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“ << std::endl;
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 < buffers.size());
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 < eltCount; 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个文件信息
- 上一篇:从零开始学微信小程序开发
- 下一篇:操作系统概念-第七版中文
评论
共有 条评论