mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

inference fix

Browse Source
This commit is contained in:
Francesco Gatti
2017-08-14 11:36:48 +02:00
parent b3a369dc29
commit 6110fffbb5
1 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_rtinference/rtinference.cpp
+4 -5
View File
@@ -3,23 +3,22 @@
int main(int argc, char *argv[]) { int main(int argc, char *argv[]) {
// Network layout
tkDNN::dataDim_t dim(1, 3, 608, 608, 1);
tkDNN::Network net(dim);
if(argc < 2 || !fileExist(argv[1])) if(argc < 2 || !fileExist(argv[1]))
FatalError("unable to read serialRT file"); FatalError("unable to read serialRT file");
//convert network to tensorRT //convert network to tensorRT
tkDNN::NetworkRT netRT(&net, argv[1]); tkDNN::NetworkRT netRT(NULL, argv[1]);
tkDNN::dataDim_t dim = netRT.input_dim;
dnnType *data; dnnType *data;
checkCuda(cudaMalloc(&data, dim.tot()*sizeof(dnnType))); checkCuda(cudaMalloc(&data, dim.tot()*sizeof(dnnType)));
printCenteredTitle(" TENSORRT inference ", '=', 30); { printCenteredTitle(" TENSORRT inference ", '=', 30); {
dim.print();
TIMER_START TIMER_START
data = netRT.infer(dim, data); data = netRT.infer(dim, data);
TIMER_STOP TIMER_STOP
dim.print();
} }
return 0; return 0;