关于TensorRT的介绍可以参考: http://blog.csdn.net/fengbingchun/article/details/78469551
以下是参考TensorRT 2.1.2中的sampleMNISTAPI.cpp文件改写的实现对手写数字0-9识别的测试代码,各个文件内容如下:
common.hpp:
#ifndef FBC_TENSORRT_TEST_COMMON_HPP_
#define FBC_TENSORRT_TEST_COMMON_HPP_
#include
#include
#include
template< typename T >
static inline int check_Cuda(T result, const char * const func, const char * const file, const int line)
{
if (result) {
fprintf(stderr, "Error CUDA: at %s: %d, error code=%d, func: %s\n", file, line, static_cast(result), func);
cudaDeviceReset(); // Make sure we call CUDA Device Reset before exiting
return -1;
}
}
template< typename T >
static inline int check(T result, const char * const func, const char * const file, const int line)
{
if (result) {
fprintf(stderr, "Error: at %s: %d, error code=%d, func: %s\n", file, line, static_cast(result), func);
return -1;
}
}
#define checkCudaErrors(val) check_Cuda((val), __FUNCTION__, __FILE__, __LINE__)
#define checkErrors(val) check((val), __FUNCTION__, __FILE__, __LINE__)
#define CHECK(x) { \
if (x) {} \
else { fprintf(stderr, "Check Failed: %s, file: %s, line: %d\n", #x, __FILE__, __LINE__); return -1; } \
}
// Logger for GIE info/warning/errors
class Logger : public nvinfer1::ILogger
{
void log(Severity severity, const char* msg) override
{
// suppress info-level messages
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
};
#endif // FBC_TENSORRT_TEST_COMMON_HPP_
mnist_api.cpp:
#include
#include
#include
#include
测试图像如下:
执行结果如下:(与 http://blog.csdn.net/fengbingchun/article/details/78552908 中结果一致)

测试代码编译步骤如下(ReadMe.txt):
在Linux下通过CMake编译TensorRT_Test中的测试代码步骤:
1. 将终端定位到CUDA_Test/prj/linux_tensorrt_cmake,依次执行如下命令:
$ mkdir build
$ cd build
$ cmake ..
$ make (生成TensorRT_Test执行文件)
$ ln -s ../../../test_data/models ./ (将models目录软链接到build目录下)
$ ln -s ../../../test_data/images ./ (将images目录软链接到build目录下)
$ ./TensorRT_Test
2. 对于有需要用OpenCV参与的读取图像的操作,需要先将对应文件中的图像路径修改为Linux支持的路径格式
GitHub: https://github.com/fengbingchun/CUDA_Test