TensorRT(2)-基本使用：mnist手写体识别
結(jié)合 tensorRT官方給出的一個(gè)例程，介紹tensorRT的使用。
這個(gè)例程是mnist手寫(xiě)體識(shí)別。例程位于目錄：?/usr/src/tensorrt/samples/sampleMNIST
文件結(jié)構(gòu)：
|
1 2 3 4 5 |
tensorrt/samples/sampleMNIST -common.cpp -common.h -Makefile -sampleMNIST.cpp |
主要是?sampleMNIST.cpp?文件，?common.cpp?文件主要提供 讀取文件的函數(shù)和 Logger對(duì)象。
main
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
#include <algorithm> #include <assert.h> #include <cmath> #include <cuda_runtime_api.h> #include <fstream> #include <iostream> #include <sstream> #include <sys/stat.h> #include <time.h> #include "NvCaffeParser.h" #include "NvInfer.h" #include "common.h" using namespace nvinfer1; using namespace nvcaffeparser1; //定義輸入輸出大小，創(chuàng)建Logger對(duì)象 //Logger是一個(gè)日志類(lèi)，在common.h文件中定義 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; //指定輸入輸出blob，和資源文件夾 const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; const std::vector<std::string> directories{"data/samples/mnist/", "data/mnist/"}; //查找文件 std::string locateFile(const std::string& input) { return locateFile(input, directories); } //讀取圖片 // 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); } ……………… ……………… int main(int argc, char** argv) { if (argc > 1) { std::cout << "This sample builds a TensorRT engine by importing a trained MNIST Caffe model.\n"; std::cout << "It uses the engine to run inference on an input image of a digit.\n"; return EXIT_SUCCESS; } // Create TRT model from caffe model and serialize it to a stream // 創(chuàng)建tensorRT流對(duì)象 trtModelStream，這個(gè)就跟文件流中的 ifstream 是類(lèi)似的。 // trtModelStream是一塊內(nèi)存區(qū)域，用于保存序列化的plan文件。 IHostMemory* trtModelStream{nullptr}; //1. build階段：調(diào)用caffeToTRTModel函數(shù)，傳入caffe模型文件和權(quán)值文件，創(chuàng)建 Ibuilder對(duì)象，調(diào)用模型解析函數(shù)， //生成的plan文件保存在 gieModelStream 中 caffeToTRTModel("mnist.prototxt", "mnist.caffemodel", std::vector<std::string>{OUTPUT_BLOB_NAME}, 1, trtModelStream); assert(trtModelStream != nullptr); // 隨機(jī)讀取一張圖片 // Read a random digit file srand(unsigned(time(nullptr))); uint8_t fileData[INPUT_H * INPUT_W]; const int num = rand() % 10; readPGMFile(locateFile(std::to_string(num) + ".pgm", directories), fileData); //將原始圖片中的像素用二進(jìn)制文本 “.:-=+*#%@”來(lái)輸出 // Print ASCII representation of digit std::cout << "\nInput:\n" << std::endl; for (int i = 0; i < INPUT_H * INPUT_W; i++) std::cout << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % INPUT_W) ? "" : "\n"); // 加載均值文件，將讀取的圖片統(tǒng)一減去平均值。 // Parse mean file ICaffeParser* parser = createCaffeParser(); IBinaryProtoBlob* meanBlob = parser->parseBinaryProto(locateFile("mnist_mean.binaryproto", directories).c_str()); parser->destroy(); // Subtract mean from image const float* meanData = reinterpret_cast<const float*>(meanBlob->getData()); float data[INPUT_H * INPUT_W]; for (int i = 0; i < INPUT_H * INPUT_W; i++) data[i] = float(fileData[i]) - meanData[i]; meanBlob->destroy(); // Deserialize engine we serialized earlier // 創(chuàng)建運(yùn)行時(shí)環(huán)境 IRuntime對(duì)象，傳入 gLogger 用于打印信息 IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream->data(), trtModelStream->size(), nullptr); assert(engine != nullptr); trtModelStream->destroy(); //創(chuàng)建上下文環(huán)境，主要用于inference 函數(shù)中啟動(dòng)cuda核 IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); //2.deploy 階段：調(diào)用 doInference 函數(shù)，進(jìn)行推理過(guò)程 // Run inference on input data float prob[OUTPUT_SIZE]; doInference(*context, data, prob, 1); //銷(xiāo)毀無(wú)用對(duì)象 // Destroy the engine context->destroy(); engine->destroy(); runtime->destroy(); //輸出分類(lèi)結(jié)果 // Print histogram of the output distribution std::cout << "\nOutput:\n\n"; float val{0.0f}; int idx{0}; for (unsigned int i = 0; i < 10; i++) { val = std::max(val, prob[i]); if (val == prob[i]) idx = i; std::cout << i << ": " << std::string(int(std::floor(prob[i] * 10 + 0.5f)), '*') << "\n"; } std::cout << std::endl; return (idx == num && val > 0.9f) ? EXIT_SUCCESS : EXIT_FAILURE; } |
實(shí)際上從第93行創(chuàng)建 IRuntime對(duì)象時(shí)，就可以認(rèn)為是屬于deploy了。
最后輸出是這樣的：讀進(jìn)一張9，輸出一個(gè)結(jié)果：
正在上傳…重新上傳取消
其中最重要的兩個(gè)函數(shù) caffeToTRTModel( ) 和 doInference( )分別完成的是build和deploy的功能。
Build Phase
正在上傳…重新上傳取消
將Caffe model 轉(zhuǎn)換為 TensorRT object，首先使用其他深度學(xué)習(xí)框架訓(xùn)練好模型，然后丟進(jìn)tensorRT優(yōu)化器中進(jìn)行優(yōu)化，優(yōu)化后會(huì)產(chǎn)生一個(gè)文件，這個(gè)文件可以認(rèn)為是優(yōu)化后的模型，接著使用序列化方法將這個(gè)優(yōu)化好后的模型存儲(chǔ)在磁盤(pán)上，存儲(chǔ)到磁盤(pán)上的文件稱(chēng)為?plan file。
這個(gè)階段需要給tensorRT提供兩個(gè)文件，分別是
- 網(wǎng)絡(luò)模型文件（比如，caffe的deploy.prototxt）
- 訓(xùn)練好的權(quán)值文件（比如，caffe的net.caffemodel）
除此之外，還需要明確 batch size，并指明輸出層。
mnist例程中的caffe模型解析代碼：標(biāo)志是創(chuàng)建 IBuilder對(duì)象。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
// 解析caffemodel到tensorrt void caffeToTRTModel(const std::string& deployFile, // Path of Caffe prototxt file const std::string& modelFile, // Path of Caffe model file const std::vector<std::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 { // 1. Create builder //創(chuàng)建一個(gè) IBuilder，傳進(jìn)gLogger參數(shù)是為了方便打印信息。 //builder 這個(gè)地方感覺(jué)像是使用了建造者模式。 IBuilder* builder = createInferBuilder(gLogger); // Parse caffe model to populate network, then set the outputs const std::string deployFpath = locateFile(deployFile, directories); const std::string modelFpath = locateFile(modelFile, directories); std::cout << "Reading Caffe prototxt: " << deployFpath << "\n"; std::cout << "Reading Caffe model: " << modelFpath << "\n"; //創(chuàng)建一個(gè) network對(duì)象，但是這個(gè)network對(duì)象只是一個(gè)空架子，里面的屬性還沒(méi)有具體的數(shù)值。 INetworkDefinition* network = builder->createNetwork(); //創(chuàng)建一個(gè)caffe模型解析對(duì)象，parser,并調(diào)用解析函數(shù)，填充network對(duì)象， //將caffe模型中的blob解析為tensorRT中的tensor，賦給blob_name_to_tensor變量。 //此處使用了模型文件和權(quán)值文件。 ICaffeParser* parser = createCaffeParser(); const IBlobNameToTensor* blobNameToTensor = parser->parse(deployFpath.c_str(), modelFpath.c_str(), *network, DataType::kFLOAT); //標(biāo)記輸出blob. // Specify output tensors of network for (auto& s : outputs) network->markOutput(*blobNameToTensor->find(s.c_str())); // 設(shè)置batch size；設(shè)置工作空間 size。 builder->setMaxBatchSize(maxBatchSize); builder->setMaxWorkspaceSize(1 << 20); // 2.Build engine //使用network創(chuàng)建 CudaEngine，優(yōu)化方法在這里執(zhí)行。 //至此，caffe模型已轉(zhuǎn)換為tensorRT object。 ICudaEngine* engine = builder->buildCudaEngine(*network); assert(engine); //銷(xiāo)毀沒(méi)用的network對(duì)象和parser對(duì)象。 // Destroy parser and network network->destroy(); parser->destroy(); //將轉(zhuǎn)換好的tensorRT object序列化到內(nèi)存中，trtModelStream是一塊內(nèi)存空間。 //這里也可以序列化到磁盤(pán)中。 // Serialize engine and destroy it trtModelStream = engine->serialize(); //銷(xiāo)毀無(wú)用對(duì)象 engine->destroy(); builder->destroy(); //關(guān)閉protobuf庫(kù) shutdownProtobufLibrary(); } |
Deploy Phase
正在上傳…重新上傳取消
Deploy 階段需要文件如下：
- 標(biāo)簽文件（這個(gè)主要是將模型產(chǎn)生的數(shù)字標(biāo)號(hào)分類(lèi)，與真實(shí)的名稱(chēng)對(duì)應(yīng)起來(lái)），不過(guò)這個(gè)例子中就不需要了，因?yàn)镸NIST的真實(shí)分類(lèi)就是數(shù)字標(biāo)號(hào)。
Deploy 階段可以認(rèn)為從主函數(shù)中就已經(jīng)開(kāi)始了。標(biāo)志是創(chuàng)建 IRuntime 對(duì)象。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
int main(int argc, char** argv) { ……………… ……………… // Deserialize engine we serialized earlier // 創(chuàng)建運(yùn)行時(shí)環(huán)境 IRuntime對(duì)象，傳入 gLogger 用于打印信息 IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream->data(), trtModelStream->size(), nullptr); assert(engine != nullptr); trtModelStream->destroy(); //創(chuàng)建上下文環(huán)境，主要用于inference 函數(shù)中啟動(dòng)cuda核 IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); //2.deploy 階段：調(diào)用 inference 函數(shù)，進(jìn)行推理過(guò)程 // Run inference on input data float prob[OUTPUT_SIZE]; doInference(*context, data, prob, 1); ……………… ……………… } |
其中 doInference函數(shù)的詳細(xì)內(nèi)容如下：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) { //使用傳進(jìn)來(lái)的context恢復(fù)engine。 const ICudaEngine& engine = context.getEngine(); //engine.getNbBindings()是為了獲取與這個(gè)engine相關(guān)的輸入輸出tensor的數(shù)量。 //這個(gè)地方，輸入+輸出 總共就2個(gè)，所以做個(gè)驗(yàn)證。 // Pointers to input and output device buffers to pass to engine. // Engine requires exactly IEngine::getNbBindings() number of buffers. assert(engine.getNbBindings() == 2); //void* 型數(shù)組，主要用于下面GPU開(kāi)辟內(nèi)存。 void* buffers[2]; //獲取與這個(gè)engine相關(guān)的輸入輸出tensor的索引。 // In order to bind the buffers, we need to know the names of the input and output tensors. // Note that indices are guaranteed to be less than IEngine::getNbBindings() const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME); const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME); //為輸入輸出tensor開(kāi)辟顯存。 // Create GPU buffers on device CHECK(cudaMalloc(&buffers[inputIndex], batchSize * INPUT_H * INPUT_W * sizeof(float))); CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float))); //創(chuàng)建cuda流，用于管理數(shù)據(jù)復(fù)制，存取，和計(jì)算的并發(fā)操作 // Create stream cudaStream_t stream; CHECK(cudaStreamCreate(&stream)); //從內(nèi)存到顯存，從CPU到GPU，將輸入數(shù)據(jù)拷貝到顯存中 //input是讀入內(nèi)存中的數(shù)據(jù)；buffers[inputIndex]是顯存上的存儲(chǔ)區(qū)域，用于存放輸入數(shù)據(jù) // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); //啟動(dòng)cuda核，異步執(zhí)行推理計(jì)算 context.enqueue(batchSize, buffers, stream, nullptr); //從顯存到內(nèi)存，將計(jì)算結(jié)果拷貝回內(nèi)存中 //output是內(nèi)存中的存儲(chǔ)區(qū)域;buffers[outputIndex]是顯存中的存儲(chǔ)區(qū)域，存放模型輸出. CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); //這個(gè)是為了同步不同的cuda流。 cudaStreamSynchronize(stream); //銷(xiāo)毀流對(duì)象和釋放顯存 // Release stream and buffers cudaStreamDestroy(stream); CHECK(cudaFree(buffers[inputIndex])); CHECK(cudaFree(buffers[outputIndex])); } |
輔助函數(shù)
用到 common.cpp 文件中的輔助函數(shù)：locateFile( ) 和 readPGMFile( )
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
#include "common.h" // Locate path to file, given its filename or filepath suffix and possible dirs it might lie in // Function will also walk back MAX_DEPTH dirs from CWD to check for such a file path //查找文件 inline std::string locateFile(const std::string& filepathSuffix, const std::vector<std::string>& directories) { const int MAX_DEPTH{10}; bool found{false}; std::string filepath; for (auto& dir : directories) { filepath = dir + filepathSuffix; for (int i = 0; i < MAX_DEPTH && !found; i++) { std::ifstream checkFile(filepath); found = checkFile.is_open(); if (found) break; filepath = "../" + filepath; // Try again in parent dir } if (found) { break; } filepath.clear(); } if (filepath.empty()) { std::string directoryList = std::accumulate(directories.begin() + 1, directories.end(), directories.front(), [](const std::string& a, const std::string& b) { return a + "\n\t" + b; }); throw std::runtime_error("Could not find " + filepathSuffix + " in data directories:\n\t" + directoryList); } return filepath; } //讀取圖片 inline void readPGMFile(const std::string& fileName, uint8_t* buffer, int inH, int inW) { std::ifstream infile(fileName, std::ifstream::binary); assert(infile.is_open() && "Attempting to read from a file that is not open."); std::string magic, h, w, max; infile >> magic >> h >> w >> max; infile.seekg(1, infile.cur); infile.read(reinterpret_cast<char*>(buffer), inH * inW); } |
日志類(lèi)
common.h文件中有個(gè)日志類(lèi)：?class Logger : public nvinfer1::ILogger
這是一個(gè)日志類(lèi)，繼承自?nvinfer1::ILogger
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
// Logger for TensorRT info/warning/errors class Logger : public nvinfer1::ILogger { public: Logger(): Logger(Severity::kWARNING) {} Logger(Severity severity): reportableSeverity(severity) {} void log(Severity severity, const char* msg) override { // suppress messages with severity enum value greater than the reportable if (severity > reportableSeverity) return; switch (severity) { case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break; case Severity::kERROR: std::cerr << "ERROR: "; break; case Severity::kWARNING: std::cerr << "WARNING: "; break; case Severity::kINFO: std::cerr << "INFO: "; break; default: std::cerr << "UNKNOWN: "; break; } std::cerr << msg << std::endl; } Severity reportableSeverity{Severity::kWARNING}; }; |
nvinfer1::ILogger?這個(gè)類(lèi)位于 tensorRT頭文件?NvInfer.h?中，此文件路徑：?/usr/include/x86_64-linux-gnu/NvInfer.h
把 ILogger 類(lèi)摘出來(lái)：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
class ILogger { public: //! //! \enum Severity //! //! The severity corresponding to a log message. //! enum class Severity { kINTERNAL_ERROR = 0, //!< An internal error has occurred. Execution is unrecoverable. kERROR = 1, //!< An application error has occurred. kWARNING = 2, //!< An application error has been discovered, but TensorRT has recovered or fallen back to a default. kINFO = 3 //!< Informational messages. }; //! //! A callback implemented by the application to handle logging messages; //! //! \param severity The severity of the message. //! \param msg The log message, null terminated. //! virtual void log(Severity severity, const char* msg) = 0; protected: virtual ~ILogger() {} }; |
可見(jiàn)這個(gè)類(lèi) 是 builder, engine and runtime 的一個(gè)日志接口，這個(gè)類(lèi)應(yīng)該以單例模式使用。即當(dāng)有多個(gè)IRuntime 和/或 IBuilder 對(duì)象時(shí)，也只能使用同一個(gè)ILogger接口。
這個(gè)接口中有個(gè)枚舉類(lèi)?enum class Severity?定義了日志報(bào)告級(jí)別，分別為 kINTERNAL_ERROR，kERROR，kWARNING和kINFO；然后還有一個(gè)純虛函數(shù) log( ) ，用戶可以自定義這個(gè)函數(shù)，以實(shí)現(xiàn)不同效果的打印。
比如common.h 文件中Logger類(lèi)的 log()函數(shù)，就是根據(jù)不同的報(bào)告級(jí)別向標(biāo)準(zhǔn)錯(cuò)誤輸出流輸出帶有不同前綴的信息。這個(gè)地方是可以自己定義的，比如你可以設(shè)置為輸出信息到文件流然后把信息保存到txt文件中等。
以上就是使用tensorRT優(yōu)化MNIST的LeNet的一個(gè)簡(jiǎn)單的例子，其實(shí)對(duì)于mnist來(lái)說(shuō)，使用tensorRT加速的意義不大，因?yàn)檫@個(gè)模型本來(lái)就比較小，這里使用這個(gè)例子主要是為了學(xué)習(xí)tensorRT的用法。
參考
總結(jié)
以上是生活随笔為你收集整理的TensorRT(2)-基本使用：mnist手写体识别的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。
- 上一篇： 卡状态不符是啥意思
- 下一篇： TensorRT(4)-Profilin