专注Java领域技术
我们一直在努力

在Caffe中调用TensorRT提供的MNIST model

原文始发于:在Caffe中调用TensorRT提供的MNIST model

在TensorRT 2.1.2中提供了MNIST的model,这里拿来用Caffe的代码调用实现,原始的mnist_mean.binaryproto文件调整为了纯二进制文件mnist_tensorrt_mean.binary,测试结果与使用TensorRT调用(http://blog.csdn.net/fengbingchun/article/details/78552908)一致:

测试代码如下:

#include "funset.hpp" #include <memory> #include <fstream> #include <tuple> #include "common.hpp"  int mnist_tensorrt_predict() { #ifdef CPU_ONLY 	caffe::Caffe::set_mode(caffe::Caffe::CPU); #else 	caffe::Caffe::set_mode(caffe::Caffe::GPU); #endif  	const std::string deploy_file{ "E:/GitCode/Caffe_Test/test_data/model/mnist/mnist_tensorrt.prototxt" }; 	const std::string model_filename{ "E:/GitCode/Caffe_Test/test_data/model/mnist/mnist_tensorrt.caffemodel" }; 	const std::string mean_file{ "E:/GitCode/Caffe_Test/test_data/model/mnist/mnist_tensorrt_mean.binary" }; 	const std::string image_path{ "E:/GitCode/Caffe_Test/test_data/images/handwritten_digits/" };  	caffe::Net<float> caffe_net(deploy_file, caffe::TEST); 	caffe_net.CopyTrainedLayersFrom(model_filename);  	// print net info 	fprintf(stdout, "input blob num: %d, output blob num: %dn", caffe_net.num_inputs(), caffe_net.num_outputs()); 	const boost::shared_ptr<caffe::Blob<float> > blob_by_name = caffe_net.blob_by_name("data"); 	int image_num = blob_by_name->num(); 	int image_channel = blob_by_name->channels(); 	int image_height = blob_by_name->height(); 	int image_width = blob_by_name->width(); 	fprintf(stdout, "inpub blob shape(num, channels, height, width): %d, %d, %d, %dn", 		image_num, image_channel, image_height, image_width);  	fprintf(stdout, "layer names: "); 	for (int i = 0; i < caffe_net.layer_names().size(); ++i) { 		fprintf(stdout, "  %s  ", caffe_net.layer_names()[i].c_str()); 	} 	fprintf(stdout, "nblob names: "); 	for (int i = 0; i < caffe_net.blob_names().size(); ++i) { 		fprintf(stdout, "  %s  ", caffe_net.blob_names()[i].c_str()); 	} 	fprintf(stdout, "nlayer types: "); 	for (int i = 0; i < caffe_net.layers().size(); ++i) { 		fprintf(stdout, "  %s  ", caffe_net.layers()[i]->type()); 	} 	const std::vector<caffe::Blob<float>*> output_blobs = caffe_net.output_blobs(); 	fprintf(stdout, "noutput blobs num: %d, blob(num, channel, heihgt, width): %d, %d, %d, %dn", 		output_blobs.size(), output_blobs[0]->num(), output_blobs[0]->channels(), output_blobs[0]->height(), output_blobs[0]->width());  	const int image_size{ image_num * image_channel * image_height * image_width }; 	std::unique_ptr<float[]> mean_values(new float[image_size]); 	std::ifstream in(mean_file.c_str(), std::ios::in | std::ios::binary); 	if (!in.is_open()) { 		fprintf(stderr, "read mean file fail: %sn", mean_file.c_str()); 		return -1; 	} 	in.read((char*)mean_values.get(), image_size * sizeof(float)); 	in.close();  	const std::vector<int> target{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 }; 	typedef std::tuple<int, float> result; 	std::vector<result> results;  	for (const auto& num : target) { 		std::string str = std::to_string(num); 		str += ".png"; 		str = image_path + str;  		cv::Mat mat = cv::imread(str.c_str(), 0); 		if (!mat.data) { 			fprintf(stderr, "load image error: %sn", str.c_str()); 			return -1; 		}  		cv::resize(mat, mat, cv::Size(image_width, image_height)); 		mat.convertTo(mat, CV_32FC1);  		float* p = (float*)mat.data; 		for (int i = 0; i < image_size; ++i) { 			p[i] -= mean_values.get()[i]; 		}  		const std::vector<caffe::Blob<float>*>& blob_input = caffe_net.input_blobs(); 		blob_input[0]->set_cpu_data((float*)mat.data);  		const std::vector<caffe::Blob<float>*>& output_blob_ = caffe_net.Forward(nullptr); 		const float* output = output_blob_[0]->cpu_data();  		float tmp{ -1.f }; 		int pos{ -1 };  		for (int j = 0; j < output_blobs[0]->count(); j++) { 			if (tmp < output[j]) { 				pos = j; 				tmp = output[j]; 			} 		}  		result ret = std::make_tuple(pos, tmp); 		results.push_back(ret); 	}  	for (auto i = 0; i < target.size(); i++) 		fprintf(stdout, "actual digit is: %d, result digit is: %d, probability: %fn", 			target[i], std::get<0>(results[i]), std::get<1>(results[i]));  	fprintf(stdout, "predict finishn");  	return 0; }

测试图像如下:

在Caffe中调用TensorRT提供的MNIST model

执行结果如下:

在Caffe中调用TensorRT提供的MNIST model

GitHub: https://github.com/fengbingchun/Caffe_Test

赞(0) 打赏
未经允许不得转载:Java小咖秀 » 在Caffe中调用TensorRT提供的MNIST model
免责声明

抢沙发

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址

专注Java技术 100年

联系我们联系我们

你默默的关注就是最好的打赏~

支付宝扫一扫打赏

微信扫一扫打赏