TensorRT序列化后结果保存

今天主要工作是对序列化后的结果进行保存,这部分的工作也是依靠论坛的力量和nvidia官方给予的信息:

主要参考的网页如下:

https://github.com/dusty-nv/jetson-inference/blob/master/tensorNet.cpp#L244

https://devtalk.nvidia.com/default/topic/1030534/jetson-tx2/segmentation-fault-core-dumped-while-doing-tensorrt-optimization-of-lenet/

https://devtalk.nvidia.com/default/topic/1038294/tensorrt/saving-a-serialized-model-to-disk/

主要的代码我粘贴在本网页中:



    //cache_path是保存序列化后的结果,这里填写路径
    std::stringstream gieModelStream;
	gieModelStream.seekg(0, gieModelStream.beg);

	
	std::ifstream cache( cache_path );

	if( !cache )
	{
       //如果没有保存序列化的结果,就需要进行第一次加载,生成trtmodel;	
		if( !ProfileModel(prototxt_path, model_path, output_blobs, maxBatchSize, gieModelStream) )
		{
			printf("failed to load %s\n", model_path);
			return 0;
		}
	   
        //保存序列化的结果;
		std::ofstream outFile;
		outFile.open(cache_path);
		outFile << gieModelStream.rdbuf();
		outFile.close();
		gieModelStream.seekg(0, gieModelStream.beg);
		printf("completed writing cache to %s\n", cache_path);
	}
	else
	{
       //使用保存的序列化结果,无需重新加载和序列化;
		printf(loading network profile from cache... %s\n", cache_path);
		gieModelStream << cache.rdbuf();
		cache.close();

		// test for half FP16 support
		nvinfer1::IBuilder* builder = CREATE_INFER_BUILDER(gLogger);
		
		if( builder != NULL )
		{
            //这块确定本次算法是否使用半精度(FP16)进行计算
			mEnableFP16 = !mOverride16 && builder->platformHasFastFp16();
			printf("platform %s FP16 support.\n", mEnableFP16 ? "has" : "does not have");
			builder->destroy();	
		}
	}

	

     /*
	 * create runtime inference engine execution context
	 */
	nvinfer1::IRuntime* infer = CREATE_INFER_RUNTIME(gLogger);
	
	if( !infer )
	{
		printf(LOG_GIE "failed to create InferRuntime\n");
		return 0;
	}
	

	// support for stringstream deserialization was deprecated in TensorRT v2
	// instead, read the stringstream into a memory buffer and pass that to TRT.
	gieModelStream.seekg(0, std::ios::end);
	const int modelSize = gieModelStream.tellg();
	gieModelStream.seekg(0, std::ios::beg);

	void* modelMem = malloc(modelSize);

	if( !modelMem )
	{
		printf(LOG_GIE "failed to allocate %i bytes to deserialize model\n", modelSize);
		return 0;
	}

	gieModelStream.read((char*)modelMem, modelSize);
	nvinfer1::ICudaEngine* engine = infer->deserializeCudaEngine(modelMem, modelSize, NULL);
	free(modelMem);
	if( !engine )
	{
		printf(LOG_GIE "failed to create CUDA engine\n");
		return 0;
	}
	
	nvinfer1::IExecutionContext* context = engine->createExecutionContext();
	
	if( !context )
	{
		printf(LOG_GIE "failed to create execution context\n");
		return 0;
	}

	if( mEnableDebug )
	{
		printf(LOG_GIE "enabling context debug sync.\n");
		context->setDebugSync(true);
	}

	if( mEnableProfiler )
		context->setProfiler(&gProfiler);
	
	mInfer   = infer;
	mEngine  = engine;
	mContext = context;

上面的代码是摘抄的https://github.com/dusty-nv/jetson-inference/blob/master/tensorNet.cpp#L244中的代码,需要的基本功能进行了备注,若有其他不懂问题,可以留言,也可以进参考代码中进行查阅;

评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

猫猫与橙子

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值