Francesco Gatti
46 lines
1.2 KiB
C++
46 lines
1.2 KiB
C++
#include<iostream>
|
|
#include "tkdnn.h"
|
|
#include <stdlib.h> /* srand, rand */
|
|
|
|
|
|
int main(int argc, char *argv[]) {
|
|
|
|
if(argc < 2 || !fileExist(argv[1]))
|
|
FatalError("unable to read serialRT file");
|
|
|
|
int BATCH_SIZE = 1;
|
|
if(argc >2)
|
|
BATCH_SIZE = atoi(argv[2]);
|
|
|
|
//always same test
|
|
srand (0);
|
|
|
|
//convert network to tensorRT
|
|
tk::dnn::NetworkRT netRT(NULL, argv[1]);
|
|
|
|
tk::dnn::dataDim_t idim = netRT.input_dim;
|
|
tk::dnn::dataDim_t odim = netRT.output_dim;
|
|
idim.n = BATCH_SIZE;
|
|
odim.n = BATCH_SIZE;
|
|
dnnType *input = new float[idim.tot()];
|
|
dnnType *output = new float[odim.tot()];
|
|
dnnType *input_d;
|
|
checkCuda( cudaMalloc(&input_d, idim.tot()*sizeof(dnnType)));
|
|
|
|
std::cout<<"Testing with batchsize: "<<BATCH_SIZE<<"\n";
|
|
printCenteredTitle(" TENSORRT inference ", '=', 30);
|
|
for(int i=0; i<10; i++) {
|
|
for(int j=0; j<idim.tot(); j++) {
|
|
input[j] = ((float) rand() / (RAND_MAX));
|
|
}
|
|
checkCuda(cudaMemcpy(input_d, input, idim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice));
|
|
|
|
tk::dnn::dataDim_t dim = idim;
|
|
TIMER_START
|
|
netRT.infer(dim, input_d);
|
|
TIMER_STOP
|
|
}
|
|
|
|
return 0;
|
|
}
|