TensorRT version

This commit is contained in:
Francesco Gatti
2017-08-01 17:51:49 +02:00
parent 1cfe70365f
commit ed5e5d58b5
6 changed files with 16 additions and 11 deletions
+5 -2
View File
@@ -1,4 +1,5 @@
#include <iostream>
#include "NvInfer.h"
#include "tkdnn.h"
#include "Network.h"
@@ -10,8 +11,10 @@ Network::Network() {
float tk_ver = float(tkDNN::getVersion())/1000;
float cu_ver = float(cudnnGetVersion())/1000;
float rt_ver = float(NV_TENSORRT_MAJOR) + float(NV_TENSORRT_MINOR)/10 + float(NV_TENSORRT_PATCH)/100;
std::cout<<"New NETWORK (tkDNN v"<<tk_ver<<", CUDNN v"<<cu_ver<<")\n";
std::cout<<"New NETWORK (tkDNN v"<<tk_ver
<<", CUDNN v"<<cu_ver<<", TensorRT v"<<rt_ver<<")\n";
dataType = CUDNN_DATA_FLOAT;
tensorFormat = CUDNN_TENSOR_NCHW;
@@ -32,7 +35,7 @@ value_type* Network::infer(dataDim_t &dim, value_type* data) {
//do infer for every layer
for(int i=0; i<num_layers; i++)
data = layers[i]->infer(dim, data);
return data;
}