shortcut rt test

include/NetworkRT.h

@@ -43,6 +43,7 @@ public:
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Route *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reorg *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Region *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l);
 
     bool serialize(const char *filename);
     bool deserialize(const char *filename);

src/NetworkRT.cpp

@@ -14,6 +14,7 @@ using namespace nvinfer1;
 #include "pluginsRT/ActivationLeakyRT.cpp"
 #include "pluginsRT/ReorgRT.cpp"
 #include "pluginsRT/RegionRT.cpp"
+#include "pluginsRT/ShortcutRT.cpp"
 #include "pluginsRT/Int8Calibrator.cpp"
 
 // Logger for info/warning/errors
@@ -167,6 +168,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
         return convert_layer(input, (Reorg*) l);
     if(type == LAYER_REGION)
         return convert_layer(input, (Region*) l);
+    if(type == LAYER_SHORTCUT)
+        return convert_layer(input, (Shortcut*) l);
 
     FatalError("Layer not implemented in tensorRT");
     return NULL;
@@ -321,6 +324,18 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Region *l) {
     return lRT;
 }
 
+ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
+    //std::cout<<"convert Shortcut\n";
+
+    //std::cout<<"New plugin Shortcut\n";
+    ITensor *tens = tensors[l->backLayer];
+    IPlugin *plugin = new ShortcutRT(tens);
+    IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+    checkNULL(lRT);
+    return lRT;
+}
+
+
 bool NetworkRT::serialize(const char *filename) {
 
     std::ofstream p(filename);

src/pluginsRT/ShortcutRT.cpp

@@ -0,0 +1,64 @@
+#include<cassert>
+#include "kernels.h"
+
+class ShortcutRT : public IPlugin {
+
+public:
+	ShortcutRT(ITensor *tens) {
+		this->back_layer = tens;
+	}
+
+	~ShortcutRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+		c = inputDims[0].d[0];
+		h = inputDims[0].d[1];
+		w = inputDims[0].d[2];
+	}
+
+	int initialize() override {
+
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+
+		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
+
+		checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 3*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, c);
+		tk::dnn::writeBUF(buf, h);
+		tk::dnn::writeBUF(buf, w);
+	}
+
+	int c, h, w;
+	ITensor *back_layer;
+};

tests/yolo3_berkeley/yolo3_berkeley.cpp

@@ -77,7 +77,7 @@ const char *c102_bin  = "../tests/yolo3_berkeley/layers/c102.bin";
 const char *c103_bin  = "../tests/yolo3_berkeley/layers/c103.bin";
 const char *c104_bin  = "../tests/yolo3_berkeley/layers/c104.bin";
 const char *c105_bin  = "../tests/yolo3_berkeley/layers/c105.bin";
-const char *output_bin = "../tests/yolo3_berkeley/debug/layer106_out.bin";
+const char *output_bin = "../tests/yolo3_berkeley/debug/layer11_out.bin";
 
 int main() {
 
@@ -93,6 +93,7 @@ int main() {
     tk::dnn::Activation a2   (&net, tk::dnn::ACTIVATION_LEAKY);
     tk::dnn::Conv2d     c3   (&net,  64, 3, 3, 1, 1, 1, 1,  c3_bin, true);
     tk::dnn::Activation a3   (&net, tk::dnn::ACTIVATION_LEAKY);
+/*
     tk::dnn::Shortcut   s4   (&net, &a1);
     tk::dnn::Conv2d     c5   (&net, 128, 3, 3, 2, 2, 1, 1,  c5_bin, true);
     tk::dnn::Activation a5   (&net, tk::dnn::ACTIVATION_LEAKY);
@@ -283,7 +284,7 @@ int main() {
     tk::dnn::Activation a104 (&net, tk::dnn::ACTIVATION_LEAKY);
     tk::dnn::Conv2d     c105 (&net,  45, 1, 1, 1, 1, 0, 0, c105_bin, false);
     tk::dnn::Yolo       g106 (&net,  10, 3);
-
+*/
     // Load input
     dnnType *data;
     dnnType *input_h;
@@ -292,8 +293,10 @@ int main() {
     //print network model
     net.print();
 
+    //convert network to tensorRT
+    tk::dnn::NetworkRT netRT(&net, "yolo3_berkeley.rt");
 
-    dnnType *out_data;  // cudnn output
+    dnnType *out_data, *out_data2; // cudnn output, tensorRT output
 
     tk::dnn::dataDim_t dim1 = dim; //input dim
     printCenteredTitle(" CUDNN inference ", '=', 30); {
@@ -303,11 +306,22 @@ int main() {
         TIMER_STOP
         dim1.print();   
     }
+    
+    tk::dnn::dataDim_t dim2 = dim;
+    printCenteredTitle(" TENSORRT inference ", '=', 30); {
+        dim2.print();
+        TIMER_START
+        out_data2 = netRT.infer(dim2, data);
+        TIMER_STOP
+        dim2.print();
+    }
 
     printCenteredTitle(" CHECK RESULTS ", '=', 30);
     dnnType *out, *out_h;
     int out_dim = net.getOutputDim().tot();
     readBinaryFile(output_bin, out_dim, &out_h, &out);
     std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out);
+    std::cout<<"TRT   vs correct"; checkResult(out_dim, out_data2, out);
+    std::cout<<"CUDNN vs TRT    "; checkResult(out_dim, out_data, out_data2);
     return 0;
 }