• 大小: 41.65MB
    文件类型: .zip
    金币: 1
    下载: 0 次
    发布日期: 2023-07-10
  • 语言: 其他
  • 标签: TensorRT  

资源简介

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个文件信息

评论

共有 条评论