mnist tensorrt incomplete

include/Layer.h

@@ -69,7 +69,6 @@ public:
              const char* fname_weights, bool batchnorm = false); 
     virtual ~LayerWgs();
 
-protected:
     int inputs, outputs;
     std::string weights_path;
 
@@ -137,9 +136,9 @@ public:
 
     virtual value_type* infer(dataDim_t &dim, value_type* srcData);
 
-protected:
     int kernelH, kernelW, strideH, strideW;
 
+protected:
     cudnnFilterDescriptor_t filterDesc;
     cudnnConvolutionDescriptor_t convDesc;
     cudnnConvolutionFwdAlgo_t algo;

include/utils.h

@@ -63,8 +63,8 @@
 }
 
 void readBinaryFile(const char* fname, int size, value_type** data_h, value_type** data_d, int seek = 0);
-int  checkResult(int size, value_type *data_d, value_type *correct_d);
-void printDeviceVector(int size, value_type* vec_d);
+int  checkResult(int size, value_type *data_d, value_type *correct_d, bool device = true);
+void printDeviceVector(int size, value_type* vec_d, bool device = true);
 void resize(int size, value_type **data);
 
 void matrixTranspose(cublasHandle_t handle, value_type* srcData, value_type* dstData, int rows, int cols);

src/utils.cpp

@@ -28,28 +28,42 @@ void readBinaryFile(const char* fname, int size, value_type** data_h, value_type
                                 cudaMemcpyHostToDevice) );
 }
 
-void printDeviceVector(int size, value_type* vec_d)
+void printDeviceVector(int size, value_type* vec_d, bool device)
 {
     value_type *vec;
-    vec = new value_type[size];
-    cudaDeviceSynchronize();
-    cudaMemcpy(vec, vec_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
+    if(device) {
+        vec = new value_type[size];
+        cudaDeviceSynchronize();
+        cudaMemcpy(vec, vec_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
+    } else {
+        vec = vec_d;
+    }
+    
     for (int i = 0; i < size; i++)
     {
         std::cout << vec[i] << " ";
     }
     std::cout << std::endl;
-    delete [] vec;
+
+    if(device)
+        delete [] vec;
 }
 
-int checkResult(int size, value_type *data_d, value_type *correct_d) {
+int checkResult(int size, value_type *data_d, value_type *correct_d, bool device) {
 
     value_type *data_h, *correct_h;
-    data_h = new value_type[size];
-    correct_h = new value_type[size];
-    cudaDeviceSynchronize();
-    cudaMemcpy(data_h, data_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
-    cudaMemcpy(correct_h, correct_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
+
+    if(device) {
+        data_h = new value_type[size];
+        correct_h = new value_type[size];
+        cudaDeviceSynchronize();
+        cudaMemcpy(data_h, data_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
+        cudaMemcpy(correct_h, correct_d, size*sizeof(value_type), cudaMemcpyDeviceToHost);
+    
+    } else {
+        data_h = data_d;
+        correct_h = correct_d;
+    }
 
     int diffs = 0;
     for(int i=0; i<size; i++) {
@@ -59,8 +73,10 @@ int checkResult(int size, value_type *data_d, value_type *correct_d) {
         }
     }
 
-    delete [] data_h;
-    delete [] correct_h;
+    if(device) {
+        delete [] data_h;
+        delete [] correct_h;
+    }
 
     return diffs;
 }

tests/mnist/test_mnistRT.cpp

@@ -1,4 +1,5 @@
 #include<iostream>
+#include<cassert>
 #include "tkdnn.h"
 #include "NvInfer.h"
 
@@ -9,8 +10,10 @@ const char *d2_bin      = "../tests/mnist/layers/d2.bin";
 const char *d3_bin      = "../tests/mnist/layers/d3.bin";
 const char *output_bin   = "../tests/mnist/output.bin";
 
+using namespace nvinfer1;
+
 // Logger for info/warning/errors
-class Logger : public nvinfer1::ILogger			
+class Logger : public ILogger			
 {
 	void log(Severity severity, const char* msg) override
 	{
@@ -59,8 +62,105 @@ int main() {
 
     std::cout<<"\n==== TensorRT ====\n";
 	// create the builder
-	nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
-	nvinfer1::INetworkDefinition* network = builder->createNetwork();
+	IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
+	INetworkDefinition* network = builder->createNetwork();
+
+    DataType dt = DataType::kFLOAT;
+	//  Create input of shape { 1, 1, 28, 28 } with name referenced by "data"
+	auto input = network->addInput("data", dt, DimsCHW{ 1, 28, 28});
+	assert(input != nullptr);
+
+    tkDNN::Conv2d *c0 = (tkDNN::Conv2d*) (net.layers[0]); 
+    Weights w { dt, c0->data_h, c0->inputs*c0->outputs*c0->kernelH*c0->kernelW};
+    Weights b { dt, c0->bias_h, c0->outputs};
+	// Add a convolution layer with 20 outputs and a 5x5 filter.
+	auto conv1 = network->addConvolution(*input, 20, DimsHW{5, 5}, w, b);
+	assert(conv1 != nullptr);
+	conv1->setStride(DimsHW{1, 1});
+    conv1->getOutput(0)->setName("out");
+
+/*
+	// Add a max pooling layer with stride of 2x2 and kernel size of 2x2.
+	auto pool1 = network->addPooling(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
+	assert(pool1 != nullptr);
+	pool1->setStride(DimsHW{2, 2});
+
+	// Add a second convolution layer with 50 outputs and a 5x5 filter.
+	auto conv2 = network->addConvolution(*pool1->getOutput(0), 50, DimsHW{5, 5}, weightMap["conv2filter"], weightMap["conv2bias"]);
+	assert(conv2 != nullptr);
+	conv2->setStride(DimsHW{1, 1});
+
+	// Add a second max pooling layer with stride of 2x2 and kernel size of 2x3>
+	auto pool2 = network->addPooling(*conv2->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
+	assert(pool2 != nullptr);
+	pool2->setStride(DimsHW{2, 2});
+
+	// Add a fully connected layer with 500 outputs.
+	auto ip1 = network->addFullyConnected(*pool2->getOutput(0), 500, weightMap["ip1filter"], weightMap["ip1bias"]);
+	assert(ip1 != nullptr);
+
+	// Add an activation layer using the ReLU algorithm.
+	auto relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);
+	assert(relu1 != nullptr);
+
+	// Add a second fully connected layer with 20 outputs.
+	auto ip2 = network->addFullyConnected(*relu1->getOutput(0), OUTPUT_SIZE, weightMap["ip2filter"], weightMap["ip2bias"]);
+	assert(ip2 != nullptr);
+
+	// Add a softmax layer to determine the probability.
+	auto prob = network->addSoftMax(*ip2->getOutput(0));
+	assert(prob != nullptr);
+	prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);
+*/
+	network->markOutput(*conv1->getOutput(0));
+
+	// Build the engine
+	builder->setMaxBatchSize(1);
+	builder->setMaxWorkspaceSize(1 << 20);
+
+	auto engine = builder->buildCudaEngine(*network);
+	// we don't need the network any more
+	network->destroy();
+
+	IExecutionContext *context = engine->createExecutionContext();
+
+	// run inference
+	// input and output buffer pointers that we pass to the engine - the engine requires exactly IEngine::getNbBindings(),
+	// of these, but in this case we know that there is exactly one input and one output.
+	assert(engine->getNbBindings() == 2);
+	void* buffers[2];
+
+	// 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()
+	int inputIndex = engine->getBindingIndex("data"); 
+    int outputIndex = engine->getBindingIndex("out");
+
+    float output[5*5*20];
+	// create GPU buffers and a stream
+	checkCuda(cudaMalloc(&buffers[inputIndex], 28*28*sizeof(float)));
+	checkCuda(cudaMalloc(&buffers[outputIndex], 5*5*20*sizeof(float)));
+
+	cudaStream_t stream;
+	checkCuda(cudaStreamCreate(&stream));
+
+	// DMA the input to the GPU,  execute the batch asynchronously, and DMA it back:
+	checkCuda(cudaMemcpyAsync(buffers[inputIndex], input_h, 1 * 28*28* sizeof(float), cudaMemcpyHostToDevice, stream));
+	context->enqueue(1, buffers, stream, nullptr);
+	checkCuda(cudaMemcpyAsync(output, buffers[outputIndex],5*5*20*sizeof(float), cudaMemcpyDeviceToHost, stream));
+	cudaStreamSynchronize(stream);
+
+
+    std::cout<<"\n==== CHECK CUDNN RESULT ====\n";
+    std::cout<<"Diff: "<<checkResult(dim.tot(), (float*)buffers[outputIndex], c0->dstData)<<"\n";
+
+	// release the stream and the buffers
+	cudaStreamDestroy(stream);
+	checkCuda(cudaFree(buffers[inputIndex]));
+	checkCuda(cudaFree(buffers[outputIndex]));
+
+	// destroy the engine
+	context->destroy();
+	engine->destroy();
 
     return 0;
 }