java调用onnx模型_使用TensorRT对caffe和pytorch onnx模型进行fp32和fp16推理

tensorrt fp32 fp16 tutorial with caffe pytorch minist model

Series

Code Example

include headers

#include

#include

#include

#include

#include

#include

#include

#include

#include

#include

#include "NvCaffeParser.h"

#include "NvOnnxConfig.h"

#include "NvOnnxParser.h"

#include "NvInfer.h"

#include "common.h"

using namespace nvinfer1;

using namespace nvcaffeparser1;

static Logger gLogger;

// Attributes of MNIST Caffe model

static const int INPUT_H = 28;

static const int INPUT_W = 28;

static const int OUTPUT_SIZE = 10;

//const char* INPUT_BLOB_NAME = "data";

const char* OUTPUT_BLOB_NAME = "prob";

const std::string mnist_data_dir = "data/mnist/";

// Simple PGM (portable greyscale map) reader

void readPGMFile(const std::string& fileName, uint8_t buffer[INPUT_H * INPUT_W])

{

readPGMFile(fileName, buffer, INPUT_H, INPUT_W);

}

caffe model to tensorrt

void caffeToTRTModel(const std::string& deployFilepath, // Path of Caffe prototxt file

const std::string& modelFilepath, // Path of Caffe model file

const std::vector<:string>& outputs, // Names of network outputs

unsigned int maxBatchSize, // Note: Must be at least as large as the batch we want to run with

IHostMemory*& trtModelStream) // Output buffer for the TRT model

{

// Create builder

IBuilder* builder = createInferBuilder(gLogger);

// Parse caffe model to populate network, then set the outputs

std::cout << "Reading Caffe prototxt: " << deployFilepath << "\n";

std::cout << "Reading Caffe model: " << modelFilepath << "\n";

INetworkDefinition* network = builder->createNetwork();

ICaffeParser* parser = createCaffeParser();

bool useFp16 = builder->platformHasFastFp16();

std::cout << "platformHasFastFp16: " << useFp16 << "\n";

bool useInt8 = builder->platformHasFastInt8();

std::cout << "platformHasFastInt8: " << useInt8 << "\n";

// create a 16-bit model if it's natively supported

DataType modelDataType = useFp16 ? DataType::kHALF : DataType::kFLOAT;

const IBlobNameToTensor* blobNameToTensor = parser->parse(deployFilepath.c_str(),

modelFilepath.c_str(),

*network,

modelDataType);

// Specify output tensors of network

// ERROR: Network must have at least one output

for (auto& s : outputs){

std::cout