99 lines
2.9 KiB
C++
99 lines
2.9 KiB
C++
#include <thread>
|
|
#include <vector>
|
|
#include <memory>
|
|
#include "ibnnet.h"
|
|
#include "InferenceEngine.h"
|
|
|
|
// stuff we know about the network and the input/output blobs
|
|
static const int MAX_BATCH_SIZE = 4;
|
|
static const int INPUT_H = 224;
|
|
static const int INPUT_W = 224;
|
|
static const int OUTPUT_SIZE = 1000;
|
|
static const int DEVICE_ID = 0;
|
|
const char* INPUT_BLOB_NAME = "data";
|
|
const char* OUTPUT_BLOB_NAME = "prob";
|
|
extern Logger gLogger;
|
|
|
|
void run_infer(std::shared_ptr<trt::IBNNet> model) {
|
|
|
|
CHECK(cudaSetDevice(model->getDeviceID()));
|
|
|
|
if(!model->deserializeEngine()) {
|
|
std::cout << "DeserializeEngine Failed." << std::endl;
|
|
return;
|
|
}
|
|
|
|
/* support batch input data */
|
|
std::vector<cv::Mat> input;
|
|
input.emplace_back( cv::Mat(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(255,255,255)) ) ;
|
|
|
|
/* run inference */
|
|
model->inference(input);
|
|
|
|
/* get output data from cudaMalloc */
|
|
float* prob = model->getOutput();
|
|
|
|
/* print output */
|
|
std::cout << "\nOutput from thread_id: " << std::this_thread::get_id() << std::endl;
|
|
if( prob != nullptr ) {
|
|
for (size_t batch_idx = 0; batch_idx < input.size(); ++batch_idx) {
|
|
for (int p = 0; p < OUTPUT_SIZE; ++p) {
|
|
std::cout<< prob[batch_idx+p] << " ";
|
|
if ((p+1) % 10 == 0) {
|
|
std::cout << std::endl;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int main(int argc, char** argv) {
|
|
|
|
trt::EngineConfig engineCfg {
|
|
INPUT_BLOB_NAME,
|
|
OUTPUT_BLOB_NAME,
|
|
nullptr,
|
|
MAX_BATCH_SIZE,
|
|
INPUT_H,
|
|
INPUT_W,
|
|
OUTPUT_SIZE,
|
|
0,
|
|
DEVICE_ID};
|
|
|
|
if (argc == 2 && std::string(argv[1]) == "-s") {
|
|
std::cout << "Serializling Engine" << std::endl;
|
|
trt::IBNNet ibnnet{engineCfg, trt::IBN::A};
|
|
ibnnet.serializeEngine();
|
|
return 0;
|
|
} else if (argc == 2 && std::string(argv[1]) == "-d") {
|
|
|
|
/*
|
|
* Support multi thread inference (mthreads>1)
|
|
* Each thread holds their own CudaEngine
|
|
* They can run on different cuda device through trt::EngineConfig setting
|
|
*/
|
|
int mthreads = 1;
|
|
std::vector<std::thread> workers;
|
|
std::vector<std::shared_ptr<trt::IBNNet>> models;
|
|
|
|
for(int i = 0; i < mthreads; ++i) {
|
|
models.emplace_back( std::make_shared<trt::IBNNet>(engineCfg, trt::IBN::A) ); // For IBNB: trt::IBN::B
|
|
}
|
|
|
|
for(int i = 0; i < mthreads; ++i) {
|
|
workers.emplace_back( std::thread(run_infer, models[i]) );
|
|
}
|
|
|
|
for(auto & worker : workers) {
|
|
worker.join();
|
|
}
|
|
|
|
return 0;
|
|
} else {
|
|
std::cerr << "arguments not right!" << std::endl;
|
|
std::cerr << "./ibnnet -s // serialize model to plan file" << std::endl;
|
|
std::cerr << "./ibnnet -d // deserialize plan file and run inference" << std::endl;
|
|
return -1;
|
|
}
|
|
}
|