Nvidia TensorRT官方例程源代码

#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  datasamples
     目录           0  2017-02-20 16:27  datasamplesgooglenet
     文件    41481639  2016-08-12 09:31  datasamplesgooglenetgooglenet.caffemodel
     文件       36261  2016-08-12 09:31  datasamplesgooglenetgooglenet.prototxt
     目录           0  2017-02-20 16:27  datasamplesmnist
     文件         797  2016-08-12 09:31  datasamplesmnist.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist1.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist2.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist3.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist4.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist5.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist6.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist7.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist8.pgm
     文件         797  2016-08-12 09:31  datasamplesmnist9.pgm
     文件     1725135  2016-08-12 09:31  datasamplesmnistmnist.caffemodel
     文件        1806  2016-08-12 09:31  datasamplesmnistmnist.prototxt
     文件        3147  2016-08-12 09:31  datasamplesmnistmnist_mean.binaryproto
     文件     7685475  2016-08-12 09:31  datasamplesmnistmnistgie.wts
     目录           0  2017-02-20 16:27  giexec
     文件        8655  2016-08-12 09:31  giexecgiexec.cpp
     文件       10099  2016-08-12 09:31  giexecgiexec.vcxproj
     文件         116  2016-08-12 09:31  giexecMakefile
     文件        3489  2016-08-12 09:31  giexecMakefile.giexec
     目录           0  2017-02-20 16:27  sampleGoogleNet
     文件         136  2016-08-12 09:31  sampleGoogleNetMakefile
     文件        3489  2016-08-12 09:31  sampleGoogleNetMakefile.sample_googlenet
     文件        6092  2016-08-12 09:31  sampleGoogleNetsampleGoogleNet.cpp
     文件       10126  2016-08-12 09:31  sampleGoogleNetsampleGoogleNet.vcxproj
     目录           0  2017-02-20 16:27  sampleMNIST
............此处省略9个文件信息