class Logger : public ILogger //1.创建logger
{
void log(Severity severity, const char* msg) override
{
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}} gLogger;
void createngine() //生成 engine
{
//创建builder
IBuilder* builder = createInferBuilder(gLogger);
//创建network
INetworkDefinition* network = builder->createNetworkV2(0U);
//创建parser
IUffParser* parser = createUffParser();
//注册网络输入输出 (输入名和输出名可在.uff文件中找到)
parser->registerInput("inputs", DimsCHW(1, 576 , 960), UffInputOrder::kNCHW);
parser->registerOutput("generator/deconv/BiasAdd");
//使用parser,填充network (DataType::kFLOAT(fp32) DataType::kHALF(fp16))
parser->parse("/home/xxxx.uff",
*network,
nvinfer1::DataType::kHALF);
//创建config并设置最大的工作空间
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 30);
//创建engine
ICudaEngine* engine = builder->buildEngineWithConfig(*network,*config);
//序列化保存engine到文件
IHostMemory *serializedModel = engine->serialize();
ofstream p("xxxx.engine");
p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());
//释放资源
parser->destroy();
network->destroy();
builder->destroy();
engine->destroy();
serializedModel->destroy();
config->destroy();
return;
}
tensorrt使用parser解析UFF模型
最新推荐文章于 2023-01-14 15:27:03 发布