mnist tensorrt incomplete

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;
 }