Add Mobilenetv2 SSD Lite post and preprocessing, add mobilenet demo

Signed-off-by: xavier <micaelaverucchi@gmail.com>

CMakeLists.txt

@@ -121,6 +121,9 @@ target_link_libraries(yolo3_demo tkDNN)
 add_executable(centernet_demo demo/demo/demo_centernet.cpp)
 target_link_libraries(centernet_demo tkDNN)
 
+add_executable(mobilenet_demo demo/demo/demo_mobilenet.cpp)
+target_link_libraries(mobilenet_demo tkDNN)
+
 add_executable(map_demo demo/demo/map.cpp)
 target_link_libraries(map_demo tkDNN)
 

demo/demo/demo_mobilenet.cpp

@@ -0,0 +1,93 @@
+#include <iostream>
+#include <signal.h>
+#include <stdlib.h> /* srand, rand */
+#include <unistd.h>
+#include <mutex>
+#include "utils.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/videoio.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include "MobilenetDetection.h"
+
+bool gRun;
+bool SAVE_RESULT = false;
+
+void sig_handler(int signo)
+{
+    std::cout << "request gateway stop\n";
+    gRun = false;
+}
+
+int main(int argc, char *argv[])
+{
+
+    std::cout << "detection\n";
+    signal(SIGINT, sig_handler);
+
+    char *net = "mobilenetv2ssd.rt";
+    if (argc > 1)
+        net = argv[1];
+    char *input = "../demo/yolo_test.mp4";
+    if (argc > 2)
+        input = argv[2];
+
+    tk::dnn::MobilenetDetection mbnet;
+    mbnet.init(net);
+
+    gRun = true;
+
+    cv::VideoCapture cap(input);
+    if (!cap.isOpened())
+        gRun = false;
+    else
+        std::cout << "camera started\n";
+
+    cv::VideoWriter resultVideo;
+    if (SAVE_RESULT)
+    {
+        int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
+        int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
+        resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M', 'P', '4', 'V'), 30, cv::Size(w, h));
+    }
+
+    cv::Mat frame;
+    cv::Mat dnn_input;
+    cv::namedWindow("detection", cv::WINDOW_NORMAL);
+
+    while (gRun)
+    {
+        cap >> frame;
+        if (!frame.data)
+        {
+            break;
+        }
+
+        // this will be resized to the net format
+        dnn_input = frame.clone();
+        // TODO: async infer
+        mbnet.update(dnn_input);
+        // draw dets
+        frame = mbnet.draw();
+
+        cv::imshow("detection", frame);
+        cv::waitKey(1);
+        if (SAVE_RESULT)
+            resultVideo << frame;
+    }
+
+    std::cout << "detection end\n";
+
+    std::cout << COL_GREENB << "\n\nTime stats:\n";
+    std::cout << "Min: " << *std::min_element(mbnet.stats.begin(), mbnet.stats.end()) << " ms\n";
+    std::cout << "Max: " << *std::max_element(mbnet.stats.begin(), mbnet.stats.end()) << " ms\n";
+    double mean = 0;
+    for (int i = 0; i < mbnet.stats.size(); i++)
+        mean += mbnet.stats[i];
+    mean /= mbnet.stats.size();
+    std::cout << "Avg: " << mean << " ms\n"
+              << COL_END;
+    return 0;
+}

include/evaluation.h

@@ -1,5 +1,5 @@
 #ifndef EVALUATION_H
-#define EVALUATION_H_H
+#define EVALUATION_H
 
 #include <iostream>
 #include <vector>

include/tkDNN/Layer.h

@@ -17,6 +17,7 @@ enum layerType_t {
     LAYER_ACTIVATION_CRELU,
     LAYER_ACTIVATION_LEAKY,
     LAYER_FLATTEN,
+    LAYER_RESHAPE,
     LAYER_MULADD,
     LAYER_POOLING,
     LAYER_SOFTMAX,
@@ -62,6 +63,7 @@ public:
             case LAYER_ACTIVATION_CRELU:    return "ActivationCReLU";
             case LAYER_ACTIVATION_LEAKY:    return "ActivationLeaky";
             case LAYER_FLATTEN:             return "Flatten";
+            case LAYER_RESHAPE:             return "Reshape";
             case LAYER_MULADD:              return "MulAdd";
             case LAYER_POOLING:             return "Pooling";
             case LAYER_SOFTMAX:             return "Softmax";
@@ -265,6 +267,20 @@ public:
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 };
 
+/**
+    Reshape layer
+*/
+class Reshape : public Layer {
+
+public:
+    Reshape(Network *net, dataDim_t new_dim, bool final=false); 
+    virtual ~Reshape();
+    virtual layerType_t getLayerType() { return LAYER_RESHAPE; };
+
+    virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+
+};
+
 
 /**
     MulAdd layer
@@ -329,11 +345,13 @@ protected:
 class Softmax : public Layer {
 
 public:
-    Softmax(Network *net); 
+    Softmax(Network *net, const tk::dnn::dataDim_t* dim=nullptr, bool final=false, const cudnnSoftmaxMode_t mode=CUDNN_SOFTMAX_MODE_CHANNEL); 
     virtual ~Softmax();
     virtual layerType_t getLayerType() { return LAYER_SOFTMAX; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+    dataDim_t dim;
+    cudnnSoftmaxMode_t mode;
 };
 
 /**
@@ -343,7 +361,7 @@ public:
 class Route : public Layer {
 
 public:
-    Route(Network *net, Layer **layers, int layers_n); 
+    Route(Network *net, Layer **layers, int layers_n, bool final=false); 
     virtual ~Route();
     virtual layerType_t getLayerType() { return LAYER_ROUTE; };
 
@@ -410,6 +428,11 @@ struct box {
     int cl;
     float x, y, w, h;
     float prob;
+
+    void print() 
+    {
+        std::cout<<"x: "<<x<<"\ty: "<<y<<"\tw: "<<w<<"\th: "<<h<<"\tcl: "<<cl<<"\tprob: "<<prob<<std::endl;
+    }
 };
 struct sortable_bbox {
     int index;

include/tkDNN/MobilenetDetection.h

@@ -0,0 +1,112 @@
+#ifndef MOBILENETDETECTION_H
+#define MOBILENETDETECTION_H
+
+#include <iostream>
+#include "tkdnn.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/videoio.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#define N_COORDS 4
+
+
+struct SSDSpec
+{
+    int feature_size = 0;
+    int shrinkage = 0;
+    int box_width = 0;
+    int box_height = 0;
+    int ratio1 = 0;
+    int ratio2 = 0;
+
+    SSDSpec() {}
+
+    SSDSpec(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2) : feature_size(feature_size), shrinkage(shrinkage), box_width(box_width), box_height(box_height),
+                                                                                                      ratio1(ratio1), ratio2(ratio2) {}
+
+    void setAll(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2)
+    {
+        this->feature_size = feature_size;
+        this->shrinkage = shrinkage;
+        this->box_width = box_width;
+        this->box_height = box_height;
+        this->ratio1 = ratio1;
+        this->ratio2 = ratio2;
+    }
+
+    void print()
+    {
+        std::cout << "fsize: " << feature_size << "\tshrinkage: " << shrinkage << "\t box W:" << box_width << "\tbox H: " << box_height << "\t x ratio:" << ratio1 << "\t y ratio:" << ratio2 << std::endl;
+    }
+};
+
+
+namespace tk
+{
+namespace dnn
+{
+class MobilenetDetection
+{
+
+private:
+    tk::dnn::NetworkRT *netRT = nullptr;
+
+    int classes = 21;
+    float iou_threshold = 0.45;
+    float center_variance = 0.1;
+    float size_variance = 0.2;
+    float conf_thresh = 0.4;
+    int input_h = 300;
+    int input_w = 300;
+    int image_size = 300;
+
+    float *priors = nullptr;
+    int n_priors = 0;
+
+    cv::Mat origImg;
+    cv::Mat bgr[3];
+
+    float *input, *input_d;
+    float *locations_h, *confidences_h;
+
+    tk::dnn::dataDim_t dim;
+
+    dnnType *conf;
+    dnnType *loc;
+
+    
+
+    float __colors[6][3] = {{1, 0, 1}, {0, 0, 1}, {0, 1, 1}, {0, 1, 0}, {1, 1, 0}, {1, 0, 0}};
+    int baseline = 0;
+    float fontScale = 0.5;
+    int thickness = 2;
+
+    void generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp = true);
+    void convert_locatios_to_boxes_and_center(float *priors, const int n_priors, float *locations, const float center_variance, const float size_variance);
+    float iou(const tk::dnn::box &a, const tk::dnn::box &b);
+    std::vector<tk::dnn::box> postprocess(float *locations, float *confidences, const int n_values, const float threshold, const int n_classes, const float iou_thresh, const int width, const int height);
+    float get_color2(int c, int x, int max);
+
+    cv::Scalar colors[256];
+    std::vector<std::string> voc_class_name;
+
+public:
+    // keep track of inference times (ms)
+    std::vector<double> stats;
+    std::vector<tk::dnn::box> detected;
+
+    MobilenetDetection() {}
+    ~MobilenetDetection() {}
+
+    void init(std::string tensor_path);
+    cv::Mat draw();
+    void update(cv::Mat &img);
+};
+
+} // namespace dnn
+} // namespace tk
+
+
+#endif /*MOBILENETDETECTION_H*/

include/tkDNN/NetworkRT.h

@@ -34,6 +34,8 @@ using namespace nvinfer1;
 #include "pluginsRT/ResizeLayerRT.h"
 //#include "pluginsRT/Int8Calibrator.h"
 #include "pluginsRT/DeformableConvRT.h"
+#include "pluginsRT/FlattenConcatRT.h"
+#include "pluginsRT/ReshapeRT.h"
 
 class PluginFactory : IPluginFactory
 {
@@ -83,6 +85,8 @@ public:
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Softmax *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Route *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Flatten *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reshape *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);

include/tkDNN/pluginsRT/FlattenConcatRT.h

@@ -0,0 +1,76 @@
+#include<cassert>
+
+class FlattenConcatRT : public IPlugin {
+
+public:
+	FlattenConcatRT() {
+		stat = cublasCreate(&handle);
+		if (stat != CUBLAS_STATUS_SUCCESS) {
+			printf ("CUBLAS initialization failed\n");
+			return;
+  		}
+	}
+
+	~FlattenConcatRT(){
+
+	}
+
+	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], 1, 1};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+		assert(nbOutputs == 1 && nbInputs ==1);
+		rows = inputDims[0].d[0];
+		cols = inputDims[0].d[1] * inputDims[0].d[2];
+		c = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2];
+		h = 1;
+		w = 1;
+	}
+
+	int initialize() override {
+		return 0;
+	}
+
+	virtual void terminate() override {
+		checkERROR(cublasDestroy(handle));
+	}
+
+	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( cudaMemcpy(dstData, srcData, rows*cols*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+	
+		float const alpha(1.0);
+		float const beta(0.0);
+		checkERROR( cublasSgeam( handle, CUBLAS_OP_T, CUBLAS_OP_N, rows, cols, &alpha, srcData, cols, &beta, srcData, rows, dstData, rows ));
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 5*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);
+		tk::dnn::writeBUF(buf, rows);
+		tk::dnn::writeBUF(buf, cols);
+	}
+
+	int c, h, w;
+	int rows, cols;
+	cublasStatus_t stat; 
+	cublasHandle_t handle; 
+};

include/tkDNN/pluginsRT/ReshapeRT.h

@@ -0,0 +1,61 @@
+#include<cassert>
+
+class ReshapeRT : public IPlugin {
+
+public:
+	ReshapeRT(dataDim_t new_dim) {
+		n = new_dim.n;
+		c = new_dim.c;
+		h = new_dim.h;
+		w = new_dim.w;
+	}
+
+	~ReshapeRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{ c,h,w};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+	}
+
+	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( cudaMemcpy(dstData, srcData, c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 4*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, n);
+		tk::dnn::writeBUF(buf, c);
+		tk::dnn::writeBUF(buf, h);
+		tk::dnn::writeBUF(buf, w);
+	}
+
+	int n, c, h, w;
+};

include/tkDNN/pluginsRT/SoftmaxRT.h

@@ -0,0 +1,64 @@
+#include<cassert>
+
+class SoftmaxRT : public IPlugin {
+
+public:
+	SoftmaxRT(const tk::dnn::dataDim_t* dim) {
+		assert(dim != nullptr);
+		this->dim.n = dim->n;
+		this->dim.c = dim->c;
+		this->dim.h = dim->h;
+		this->dim.w = dim->w;
+		this->dim.l = dim->l;
+	}
+
+	~SoftmaxRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsNCHW{this->dim.n,this->dim.c,this->dim.h,this->dim.w };
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {		
+	}
+
+	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]);
+		
+		
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 5*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, this->dim.n);
+		tk::dnn::writeBUF(buf, this->dim.c);
+		tk::dnn::writeBUF(buf, this->dim.h);
+		tk::dnn::writeBUF(buf, this->dim.w);
+		tk::dnn::writeBUF(buf, this->dim.l);
+	}
+
+	dataDim_t dim;
+};

src/CenternetDetection.cpp

@@ -226,7 +226,7 @@ void CenternetDetection::update(cv::Mat &imageORIG) {
     sz_old = sz;
     cv::cuda::GpuMat im_Orig; 
     im_Orig = cv::cuda::GpuMat(imageORIG);
-    cv::cuda::resize (im_Orig, imageF1_d, cv::Size(new_width, new_height)); 
+    // cv::cuda::resize (im_Orig, imageF1_d, cv::Size(new_width, new_height)); 
     checkCuda( cudaDeviceSynchronize() );
     
     sz = imageF1_d.size();
@@ -235,7 +235,7 @@ void CenternetDetection::update(cv::Mat &imageORIG) {
     std::cout << " TIME resize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
     step_t = end_t;
     
-    cv::cuda::warpAffine(imageF1_d, imageF2_d, trans, cv::Size(inp_width, inp_height), cv::INTER_LINEAR );
+    // cv::cuda::warpAffine(imageF1_d, imageF2_d, trans, cv::Size(inp_width, inp_height), cv::INTER_LINEAR );
     checkCuda( cudaDeviceSynchronize() );
     end_t = std::chrono::steady_clock::now();
     std::cout << " TIME warpAffine: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
@@ -248,7 +248,7 @@ void CenternetDetection::update(cv::Mat &imageORIG) {
     step_t = end_t;
     
     dim2 = dim;
-    cv::cuda::split(imageF1_d,bgr);//split source
+    // cv::cuda::split(imageF1_d,bgr);//split source
     end_t = std::chrono::steady_clock::now();
     std::cout << " TIME split: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
     step_t = end_t;

src/MobilenetDetection.cpp

@@ -0,0 +1,330 @@
+#include "MobilenetDetection.h"
+
+bool boxProbCmp(const tk::dnn::box &a, const tk::dnn::box &b)
+{
+    return (a.prob > b.prob);
+}
+
+namespace tk
+{
+namespace dnn
+{
+
+void MobilenetDetection::generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp)
+{
+    n_priors = 0;
+    for (int i = 0; i < n_specs; i++)
+    {
+        n_priors += specs[i].feature_size * specs[i].feature_size * 6;
+    }
+
+    // std::cout<<"n priors: "<<n_priors<<std::endl;
+    // std::cout<<"n priors: "<<n_specs<<std::endl;
+
+    priors = (float *)malloc(N_COORDS * n_priors * sizeof(float));
+
+    int i_prio = 0;
+    float scale, x_center, y_center, h, w, size, ratio;
+    int min, max;
+    for (int i = 0; i < n_specs; i++)
+    {
+        scale = (float)image_size / (float)specs[i].shrinkage;
+        min = specs[i].box_height > specs[i].box_width ? specs[i].box_width : specs[i].box_height;
+        max = specs[i].box_height < specs[i].box_width ? specs[i].box_width : specs[i].box_height;
+        for (int j = 0; j < specs[i].feature_size; j++)
+        {
+            for (int k = 0; k < specs[i].feature_size; k++)
+            {
+                //small sized square box
+                size = min;
+                x_center = (k + 0.5f) / scale;
+                y_center = (j + 0.5f) / scale;
+                h = w = (float)size / (float)image_size;
+
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w;
+                priors[i_prio * N_COORDS + 3] = h;
+                ++i_prio;
+
+                //big sized square box
+                size = sqrt(max * min);
+                h = w = (float)size / (float)image_size;
+
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w;
+                priors[i_prio * N_COORDS + 3] = h;
+                ++i_prio;
+
+                //change h/w ratio of the small sized box
+                size = min;
+                h = w = size / (float)image_size;
+                ratio = sqrt(specs[i].ratio1);
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w * ratio;
+                priors[i_prio * N_COORDS + 3] = h / ratio;
+                ++i_prio;
+
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w / ratio;
+                priors[i_prio * N_COORDS + 3] = h * ratio;
+                ++i_prio;
+
+                ratio = sqrt(specs[i].ratio2);
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w * ratio;
+                priors[i_prio * N_COORDS + 3] = h / ratio;
+                ++i_prio;
+
+                priors[i_prio * N_COORDS + 0] = x_center;
+                priors[i_prio * N_COORDS + 1] = y_center;
+                priors[i_prio * N_COORDS + 2] = w / ratio;
+                priors[i_prio * N_COORDS + 3] = h * ratio;
+                ++i_prio;
+            }
+        }
+    }
+
+    if (clamp)
+    {
+        for (int i = 0; i < n_priors * N_COORDS; i++)
+        {
+            priors[i] = priors[i] > 1.0f ? 1.0f : priors[i];
+            priors[i] = priors[i] < 0.0f ? 0.0f : priors[i];
+
+            // std::cout<<priors[i]<<" ";
+            // if((i+1)%4 == 0)
+            //     std::cout<< i/4 <<" " <<std::endl;
+        }
+    }
+}
+
+void MobilenetDetection::convert_locatios_to_boxes_and_center(float *priors, const int n_priors, float *locations, const float center_variance, const float size_variance)
+{
+    float cur_x, cur_y;
+    for (int i = 0; i < n_priors; i++)
+    {
+        locations[i * N_COORDS + 0] = locations[i * N_COORDS + 0] * center_variance * priors[i * N_COORDS + 2] + priors[i * N_COORDS + 0];
+        locations[i * N_COORDS + 1] = locations[i * N_COORDS + 1] * center_variance * priors[i * N_COORDS + 3] + priors[i * N_COORDS + 1];
+        locations[i * N_COORDS + 2] = exp(locations[i * N_COORDS + 2] * size_variance) * priors[i * N_COORDS + 2];
+        locations[i * N_COORDS + 3] = exp(locations[i * N_COORDS + 3] * size_variance) * priors[i * N_COORDS + 3];
+
+        cur_x = locations[i * N_COORDS + 0];
+        cur_y = locations[i * N_COORDS + 1];
+
+        locations[i * N_COORDS + 0] = cur_x - locations[i * N_COORDS + 2] / 2;
+        locations[i * N_COORDS + 1] = cur_y - locations[i * N_COORDS + 3] / 2;
+        locations[i * N_COORDS + 2] = cur_x + locations[i * N_COORDS + 2] / 2;
+        locations[i * N_COORDS + 3] = cur_y + locations[i * N_COORDS + 3] / 2;
+
+        // std::cout<<locations[i*N_COORDS + 0]<<" "<<locations[i*N_COORDS + 1]<<" "<<locations[i*N_COORDS + 2]<<" "<<locations[i*N_COORDS + 3]<<" "<<std::endl;
+    }
+}
+
+float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b)
+{
+    float max_x = a.x > b.x ? a.x : b.x;
+    float max_y = a.y > b.y ? a.y : b.y;
+    float min_w = a.w < b.w ? a.w : b.w;
+    float min_h = a.h < b.h ? a.h : b.h;
+
+    float ao_w = min_w - max_x > 0 ? min_w - max_x : 0;
+    float ao_h = min_h - max_y > 0 ? min_h - max_y : 0;
+
+    // std::cout<<" ao w: "<<ao_w<<" ao h: "<<ao_h<<std::endl;
+
+    float area_overlap = ao_w * ao_h;
+    float area_0_w = a.w - a.x > 0 ? a.w - a.x : 0;
+    float area_0_h = a.h - a.y > 0 ? a.h - a.y : 0;
+
+    float area_1_w = b.w - b.x > 0 ? b.w - b.x : 0;
+    float area_1_h = b.h - b.y > 0 ? b.h - b.y : 0;
+
+    float area_0 = area_0_h * area_0_w;
+    float area_1 = area_1_h * area_1_w;
+
+    // std::cout<<" area_overlap : "<<area_overlap<<" area_0: "<<area_0<<" area_1: "<<area_1<<std::endl;
+
+    float iou = area_overlap / (area_0 + area_1 - area_overlap + 1e-5);
+    return iou;
+}
+
+std::vector<tk::dnn::box> MobilenetDetection::postprocess(float *locations, float *confidences, const int n_values, const float threshold, const int n_classes, const float iou_thresh, const int width, const int height)
+{
+    float *conf_per_class;
+    std::vector<tk::dnn::box> detections;
+    for (int i = 1; i < n_classes; i++)
+    {
+        conf_per_class = &confidences[i * n_values];
+        std::vector<tk::dnn::box> boxes;
+        for (int j = 0; j < n_values; j++)
+        {
+
+            if (conf_per_class[j] > threshold)
+            {
+                tk::dnn::box b;
+                b.cl = i;
+                b.prob = conf_per_class[j];
+                b.x = locations[j * N_COORDS + 0];
+                b.y = locations[j * N_COORDS + 1];
+                b.w = locations[j * N_COORDS + 2];
+                b.h = locations[j * N_COORDS + 3];
+
+                boxes.push_back(b);
+            }
+        }
+        std::sort(boxes.begin(), boxes.end(), boxProbCmp);
+        // for(auto b:boxes)
+        //     b.print();
+        std::vector<tk::dnn::box> remaining;
+        while (boxes.size() > 0)
+        {
+            remaining.clear();
+
+            tk::dnn::box b;
+            b.cl = boxes[0].cl;
+            b.prob = boxes[0].prob;
+            b.x = boxes[0].x * width;
+            b.y = boxes[0].y * height;
+            b.w = boxes[0].w * width;
+            b.h = boxes[0].h * height;
+            detections.push_back(b);
+            for (size_t j = 1; j < boxes.size(); j++)
+            {
+                if (iou(boxes[0], boxes[j]) <= iou_thresh)
+                {
+                    remaining.push_back(boxes[j]);
+                }
+            }
+            boxes = remaining;
+        }
+    }
+    // std::cout<<"picked"<<std::endl;
+    // for(auto b:detections)
+    //     b.print();
+
+    return detections;
+}
+
+float MobilenetDetection::get_color2(int c, int x, int max)
+{
+    float ratio = ((float)x / max) * 5;
+    int i = floor(ratio);
+    int j = ceil(ratio);
+    ratio -= i;
+    float r = (1 - ratio) * __colors[i % 6][c % 3] + ratio * __colors[j % 6][c % 3];
+    //printf("%f\n", r);
+    return r;
+}
+
+void MobilenetDetection::init(std::string tensor_path)
+{
+    const int n_SSDSpec = 6;
+    SSDSpec specs[6];
+    specs[0].setAll(19, 16, 60, 105, 2, 3);
+    specs[1].setAll(10, 32, 105, 150, 2, 3);
+    specs[2].setAll(5, 64, 150, 195, 2, 3);
+    specs[3].setAll(3, 100, 195, 240, 2, 3);
+    specs[4].setAll(2, 150, 240, 285, 2, 3);
+    specs[5].setAll(1, 300, 285, 330, 2, 3);
+
+    generate_ssd_priors(specs, n_SSDSpec);
+
+    netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
+
+    checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot()));
+    checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot()));
+
+    locations_h = (float *)malloc(N_COORDS * n_priors * sizeof(float));
+    confidences_h = (float *)malloc(n_priors * classes * sizeof(float));
+
+    dim = tk::dnn::dataDim_t(1, 3, input_w, input_h, 1);
+
+    for (int c = 0; c < classes; c++)
+    {
+        int offset = c * 123457 % classes;
+        float r = get_color2(2, offset, classes);
+        float g = get_color2(1, offset, classes);
+        float b = get_color2(0, offset, classes);
+        colors[c] = cv::Scalar(int(255.0 * b), int(255.0 * g), int(255.0 * r));
+    }
+
+    const char *voc_class_name_[] = {
+        "BACKGROUND", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
+        "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike",
+        "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"};
+    voc_class_name = std::vector<std::string>(voc_class_name_, std::end(voc_class_name_));
+}
+
+cv::Mat MobilenetDetection::draw()
+{
+    tk::dnn::box b;
+    for (size_t i = 0; i < detected.size(); i++)
+    {
+        b = detected[i];
+        std::string det_class = voc_class_name[b.cl];
+        cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point(b.w, b.h), colors[b.cl], 2);
+        // draw label
+        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
+        cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point((b.x + textSize.width - 2), (b.y - textSize.height - 2)), colors[b.cl], -1);
+        cv::putText(origImg, det_class, cv::Point(b.x, (b.y - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
+    }
+    return origImg;
+}
+
+void MobilenetDetection::update(cv::Mat &img)
+{
+    TIMER_START
+    detected.clear();
+
+    //save origin image
+    origImg = img;
+    cv::Size sz = origImg.size();
+
+    //resize image, remove mean, divide by std
+    cv::Mat frame_resize, frame_nomean, frame_scaled;
+    resize(origImg, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
+    frame_resize.convertTo(frame_nomean, CV_32FC3, 1, -127);
+    frame_nomean.convertTo(frame_scaled, CV_32FC3, 1 / 128.0, 0);
+
+    //copy image into tensor and copy it into GPU
+    cv::split(frame_scaled, bgr);
+    for (int i = 0; i < netRT->input_dim.c; i++)
+    {
+        int idx = i * frame_scaled.rows * frame_scaled.cols;
+        memcpy((void *)&input[idx], (void *)bgr[i].data, frame_scaled.rows * frame_scaled.cols * sizeof(dnnType));
+    }
+    checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
+
+    //do inference
+    tk::dnn::dataDim_t dim2 = dim;
+    printCenteredTitle(" TENSORRT inference ", '=', 30);
+    {
+        dim2.print();
+        TIMER_START
+        netRT->infer(dim2, input_d);
+        TIMER_STOP
+        dim2.print();
+    }
+
+    //get confidences and locations
+    conf = (dnnType *)netRT->buffersRT[3];
+    loc = (dnnType *)netRT->buffersRT[4];
+
+    checkCuda(cudaMemcpy(confidences_h, conf, n_priors * classes * sizeof(float), cudaMemcpyDeviceToHost));
+    checkCuda(cudaMemcpy(locations_h, loc, N_COORDS * n_priors * sizeof(float), cudaMemcpyDeviceToHost));
+
+    //postprocess
+    convert_locatios_to_boxes_and_center(priors, n_priors, locations_h, center_variance, size_variance);
+    detected = postprocess(locations_h, confidences_h, n_priors, conf_thresh, classes, iou_threshold, sz.width, sz.height);
+
+    TIMER_STOP
+    stats.push_back(t_ns);
+}
+
+} // namespace dnn
+} // namespace tk

src/NetworkRT.cpp

@@ -179,6 +179,10 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
         return convert_layer(input, (Softmax*) l);
     if(type == LAYER_ROUTE)
         return convert_layer(input, (Route*) l);
+    if(type == LAYER_FLATTEN)
+        return convert_layer(input, (Flatten*) l);
+    if(type == LAYER_RESHAPE)
+        return convert_layer(input, (Reshape*) l);
     if(type == LAYER_REORG)
         return convert_layer(input, (Reorg*) l);
     if(type == LAYER_REGION)
@@ -389,6 +393,24 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
     return lRT;
 }
 
+ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) {
+
+    IPlugin *plugin = new FlattenConcatRT();
+    IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+    checkNULL(lRT);
+    return lRT;
+}
+
+ILayer* NetworkRT::convert_layer(ITensor *input, Reshape *l) {
+    // std::cout<<"convert Reshape\n";
+
+    l->output_dim.print();
+    IPlugin *plugin = new ReshapeRT(l->output_dim);
+    IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+    checkNULL(lRT);
+    return lRT;
+}
+
 ILayer* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
     //std::cout<<"convert Reorg\n";
 
@@ -597,6 +619,24 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
         return r;
     } 
 
+    if(name.find("Flatten") == 0) {
+        FlattenConcatRT *r = new FlattenConcatRT(); 
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        r->rows = readBUF<int>(buf);
+        r->cols = readBUF<int>(buf);
+        return r;
+    } 
+
+    if(name.find("Reshape") == 0) {
+
+        dataDim_t new_dim(readBUF<int>(buf), readBUF<int>(buf),readBUF<int>(buf), readBUF<int>(buf));
+        ReshapeRT *r = new ReshapeRT(new_dim); 
+        
+        return r;
+    } 
+
     if(name.find("Yolo") == 0) {
         YoloRT *r = new YoloRT(readBUF<int>(buf),    //classes
                                 readBUF<int>(buf),   //num

src/Reshape.cpp

@@ -0,0 +1,37 @@
+#include <iostream>
+
+#include "Layer.h"
+#include "kernels.h"
+
+namespace tk { namespace dnn {
+
+Reshape::Reshape(Network *net, dataDim_t new_dim, bool final) : Layer(net, final) {
+
+    checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(dnnType)) );
+
+    output_dim.n = new_dim.n;
+    output_dim.c = new_dim.c;
+    output_dim.h = new_dim.h;
+    output_dim.w = new_dim.w;
+    output_dim.l = new_dim.l;
+
+}
+
+Reshape::~Reshape() {
+
+    checkCuda( cudaFree(dstData) );
+}
+
+dnnType* Reshape::infer(dataDim_t &dim, dnnType* srcData) {
+
+    //transpose per channel
+
+    checkCuda( cudaMemcpy(dstData, srcData, dim.n*dim.c*dim.h*dim.w*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+
+    //update data dimensions    
+    dim = output_dim;
+
+    return dstData;
+}
+
+}}

src/Route.cpp

@@ -5,7 +5,7 @@
 
 namespace tk { namespace dnn {
 
-Route::Route(Network *net, Layer **layers, int layers_n) : Layer(net) {
+Route::Route(Network *net, Layer **layers, int layers_n, bool final) : Layer(net, final) {
 
     this->layers_n = layers_n;
     if(layers_n > MAX_INPUT_LAYERS)

src/Softmax.cpp

@@ -5,22 +5,40 @@
 
 namespace tk { namespace dnn {
 
-Softmax::Softmax(Network *net) : Layer(net) {
+Softmax::Softmax(Network *net, const tk::dnn::dataDim_t* dim, bool final, const cudnnSoftmaxMode_t mode) : Layer(net, final) {
 
     checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(dnnType)) );
 
+    this->mode = mode;
+    if(dim == nullptr)
+    {
+        this->dim.n= input_dim.n;
+        this->dim.c= input_dim.c;
+        this->dim.h= input_dim.h;
+        this->dim.w= input_dim.w;
+        this->dim.l= input_dim.l;
+    }
+    else
+    {
+        this->dim.n= dim->n;
+        this->dim.c= dim->c;
+        this->dim.h= dim->h;
+        this->dim.w= dim->w;
+        this->dim.l= dim->l;
+    }
+
     checkCUDNN( cudnnSetTensor4dDescriptor(srcTensorDesc,
                                         net->tensorFormat,
                                         net->dataType,
-                                        input_dim.n*input_dim.l, 
+                                        this->dim.n*this->dim.l, 
-                                        input_dim.c,
+                                        this->dim.c,
-                                        input_dim.h, input_dim.w) );
+                                        this->dim.h, this->dim.w) );
     checkCUDNN( cudnnSetTensor4dDescriptor(dstTensorDesc,
                                             net->tensorFormat,
                                             net->dataType,
-                                            input_dim.n*input_dim.l, 
+                                            this->dim.n*this->dim.l, 
-                                            input_dim.c,
+                                            this->dim.c,
-                                            input_dim.h, input_dim.w) );
+                                            this->dim.h, this->dim.w) );
 }
 
 Softmax::~Softmax() {
@@ -34,7 +52,7 @@ dnnType* Softmax::infer(dataDim_t &dim, dnnType* srcData) {
     dnnType beta  = dnnType(0);
     checkCUDNN( cudnnSoftmaxForward(net->cudnnHandle,
                                     CUDNN_SOFTMAX_ACCURATE ,
-                                    CUDNN_SOFTMAX_MODE_CHANNEL,
+                                    this->mode,
                                     &alpha,
                                     srcTensorDesc,
                                     srcData,

tests/mobilenetv2ssd/mobilenetv2ssd.cpp

@@ -1,139 +1,143 @@
 #include <iostream>
 #include "tkdnn.h"
 
-const char *output_bin = "../tests/mobilenetv2ssd/debug/regression_headers-5.bin";
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/videoio.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+const char *output_bin1 = "../tests/mobilenetv2ssd/debug/classification_headers-5.bin";
+const char *output_bin2 = "../tests/mobilenetv2ssd/debug/regression_headers-5.bin";
 const char *input_bin = "../tests/mobilenetv2ssd/debug/input.bin";
 
 const char *conv0_bin = "../tests/mobilenetv2ssd/layers/base_net-0-0.bin";
-const char *inverted_residual1[]={
+const char *inverted_residual1[] = {
-"../tests/mobilenetv2ssd/layers/base_net-1-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-1-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-1-conv-3.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-1-conv-3.bin"};
-const char *inverted_residual2[]={
+const char *inverted_residual2[] = {
-"../tests/mobilenetv2ssd/layers/base_net-2-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-2-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-2-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-2-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-2-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-2-conv-6.bin"};
-const char *inverted_residual3[]={
+const char *inverted_residual3[] = {
-"../tests/mobilenetv2ssd/layers/base_net-3-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-3-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-3-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-3-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-3-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-3-conv-6.bin"};
-const char *inverted_residual4[]={
+const char *inverted_residual4[] = {
-"../tests/mobilenetv2ssd/layers/base_net-4-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-4-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-4-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-4-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-4-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-4-conv-6.bin"};
-const char *inverted_residual5[]={
+const char *inverted_residual5[] = {
-"../tests/mobilenetv2ssd/layers/base_net-5-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-5-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-5-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-5-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-5-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-5-conv-6.bin"};
-const char *inverted_residual6[]={
+const char *inverted_residual6[] = {
-"../tests/mobilenetv2ssd/layers/base_net-6-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-6-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-6-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-6-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-6-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-6-conv-6.bin"};
-const char *inverted_residual7[]={
+const char *inverted_residual7[] = {
-"../tests/mobilenetv2ssd/layers/base_net-7-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-7-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-7-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-7-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-7-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-7-conv-6.bin"};
-const char *inverted_residual8[]={
+const char *inverted_residual8[] = {
-"../tests/mobilenetv2ssd/layers/base_net-8-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-8-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-8-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-8-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-8-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-8-conv-6.bin"};
-const char *inverted_residual9[]={
+const char *inverted_residual9[] = {
-"../tests/mobilenetv2ssd/layers/base_net-9-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-9-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-9-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-9-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-9-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-9-conv-6.bin"};
-const char *inverted_residual10[]={
+const char *inverted_residual10[] = {
-"../tests/mobilenetv2ssd/layers/base_net-10-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-10-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-10-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-10-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-10-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-10-conv-6.bin"};
-const char *inverted_residual11[]={
+const char *inverted_residual11[] = {
-"../tests/mobilenetv2ssd/layers/base_net-11-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-11-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-11-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-11-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-11-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-11-conv-6.bin"};
-const char *inverted_residual12[]={
+const char *inverted_residual12[] = {
-"../tests/mobilenetv2ssd/layers/base_net-12-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-12-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-12-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-12-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-12-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-12-conv-6.bin"};
-const char *inverted_residual13[]={
+const char *inverted_residual13[] = {
-"../tests/mobilenetv2ssd/layers/base_net-13-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-13-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-13-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-13-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-13-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-13-conv-6.bin"};
-const char *inverted_residual14[]={
+const char *inverted_residual14[] = {
-"../tests/mobilenetv2ssd/layers/base_net-14-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-14-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-14-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-14-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-14-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-14-conv-6.bin"};
-const char *inverted_residual15[]={
+const char *inverted_residual15[] = {
-"../tests/mobilenetv2ssd/layers/base_net-15-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-15-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-15-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-15-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-15-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-15-conv-6.bin"};
-const char *inverted_residual16[]={
+const char *inverted_residual16[] = {
-"../tests/mobilenetv2ssd/layers/base_net-16-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-16-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-16-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-16-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-16-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-16-conv-6.bin"};
-const char *inverted_residual17[]={
+const char *inverted_residual17[] = {
-"../tests/mobilenetv2ssd/layers/base_net-17-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-17-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/base_net-17-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/base_net-17-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/base_net-17-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/base_net-17-conv-6.bin"};
 
 const char *conv18 = "../tests/mobilenetv2ssd/layers/base_net-18-0.bin";
 
-const char *extras0[]={
+const char *extras0[] = {
-"../tests/mobilenetv2ssd/layers/extras-0-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/extras-0-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/extras-0-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/extras-0-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/extras-0-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/extras-0-conv-6.bin"};
-const char *extras1[]={
+const char *extras1[] = {
-"../tests/mobilenetv2ssd/layers/extras-1-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/extras-1-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/extras-1-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/extras-1-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/extras-1-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/extras-1-conv-6.bin"};
-const char *extras2[]={
+const char *extras2[] = {
-"../tests/mobilenetv2ssd/layers/extras-2-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/extras-2-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/extras-2-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/extras-2-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/extras-2-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/extras-2-conv-6.bin"};
-const char *extras3[]={
+const char *extras3[] = {
-"../tests/mobilenetv2ssd/layers/extras-3-conv-0.bin",
+    "../tests/mobilenetv2ssd/layers/extras-3-conv-0.bin",
-"../tests/mobilenetv2ssd/layers/extras-3-conv-3.bin",
+    "../tests/mobilenetv2ssd/layers/extras-3-conv-3.bin",
-"../tests/mobilenetv2ssd/layers/extras-3-conv-6.bin"};
+    "../tests/mobilenetv2ssd/layers/extras-3-conv-6.bin"};
 
-const char *classification_header0[]={
+const char *classification_header0[] = {
-"../tests/mobilenetv2ssd/layers/classification_headers-0-0.bin",
+    "../tests/mobilenetv2ssd/layers/classification_headers-0-0.bin",
-"../tests/mobilenetv2ssd/layers/classification_headers-0-3.bin"};
+    "../tests/mobilenetv2ssd/layers/classification_headers-0-3.bin"};
-const char *classification_header1[]={
+const char *classification_header1[] = {
-"../tests/mobilenetv2ssd/layers/classification_headers-1-0.bin",
+    "../tests/mobilenetv2ssd/layers/classification_headers-1-0.bin",
-"../tests/mobilenetv2ssd/layers/classification_headers-1-3.bin"};
+    "../tests/mobilenetv2ssd/layers/classification_headers-1-3.bin"};
-const char *classification_header2[]={
+const char *classification_header2[] = {
-"../tests/mobilenetv2ssd/layers/classification_headers-2-0.bin",
+    "../tests/mobilenetv2ssd/layers/classification_headers-2-0.bin",
-"../tests/mobilenetv2ssd/layers/classification_headers-2-3.bin"};
+    "../tests/mobilenetv2ssd/layers/classification_headers-2-3.bin"};
-const char *classification_header3[]={
+const char *classification_header3[] = {
-"../tests/mobilenetv2ssd/layers/classification_headers-3-0.bin",
+    "../tests/mobilenetv2ssd/layers/classification_headers-3-0.bin",
-"../tests/mobilenetv2ssd/layers/classification_headers-3-3.bin"};
+    "../tests/mobilenetv2ssd/layers/classification_headers-3-3.bin"};
-const char *classification_header4[]={
+const char *classification_header4[] = {
-"../tests/mobilenetv2ssd/layers/classification_headers-4-0.bin",
+    "../tests/mobilenetv2ssd/layers/classification_headers-4-0.bin",
-"../tests/mobilenetv2ssd/layers/classification_headers-4-3.bin"};
+    "../tests/mobilenetv2ssd/layers/classification_headers-4-3.bin"};
 
 const char *classification_header5 = "../tests/mobilenetv2ssd/layers/classification_headers-5.bin";
 
-const char *regression_header0[]={
+const char *regression_header0[] = {
-"../tests/mobilenetv2ssd/layers/regression_headers-0-0.bin",
+    "../tests/mobilenetv2ssd/layers/regression_headers-0-0.bin",
-"../tests/mobilenetv2ssd/layers/regression_headers-0-3.bin"};
+    "../tests/mobilenetv2ssd/layers/regression_headers-0-3.bin"};
-const char *regression_header1[]={
+const char *regression_header1[] = {
-"../tests/mobilenetv2ssd/layers/regression_headers-1-0.bin",
+    "../tests/mobilenetv2ssd/layers/regression_headers-1-0.bin",
-"../tests/mobilenetv2ssd/layers/regression_headers-1-3.bin"};
+    "../tests/mobilenetv2ssd/layers/regression_headers-1-3.bin"};
-const char *regression_header2[]={
+const char *regression_header2[] = {
-"../tests/mobilenetv2ssd/layers/regression_headers-2-0.bin",
+    "../tests/mobilenetv2ssd/layers/regression_headers-2-0.bin",
-"../tests/mobilenetv2ssd/layers/regression_headers-2-3.bin"};
+    "../tests/mobilenetv2ssd/layers/regression_headers-2-3.bin"};
-const char *regression_header3[]={
+const char *regression_header3[] = {
-"../tests/mobilenetv2ssd/layers/regression_headers-3-0.bin",
+    "../tests/mobilenetv2ssd/layers/regression_headers-3-0.bin",
-"../tests/mobilenetv2ssd/layers/regression_headers-3-3.bin"};
+    "../tests/mobilenetv2ssd/layers/regression_headers-3-3.bin"};
-const char *regression_header4[]={
+const char *regression_header4[] = {
-"../tests/mobilenetv2ssd/layers/regression_headers-4-0.bin",
+    "../tests/mobilenetv2ssd/layers/regression_headers-4-0.bin",
-"../tests/mobilenetv2ssd/layers/regression_headers-4-3.bin"};
+    "../tests/mobilenetv2ssd/layers/regression_headers-4-3.bin"};
 
 const char *regression_header5 = "../tests/mobilenetv2ssd/layers/regression_headers-5.bin";
 
 
-
-
-
 int main()
 {
+    int classes = 21;
 
     // Network layout
     tk::dnn::dataDim_t dim(1, 3, 300, 300, 1);
@@ -143,11 +147,10 @@ int main()
     tk::dnn::Activation relu3(&net, CUDNN_ACTIVATION_RELU);
 
     //Inverted Residual 1
-    
+
-    tk::dnn::Conv2d conv2(&net, 32, 3, 3, 1, 1, 1, 1, inverted_residual1[0], true,false, false,32);
+    tk::dnn::Conv2d conv2(&net, 32, 3, 3, 1, 1, 1, 1, inverted_residual1[0], true, false, false, 32);
     tk::dnn::Activation relu5(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d conv3(&net, 16, 1, 1, 1, 1, 0, 0, inverted_residual1[1], true);
-    
 
     //Inverted Residual 2
     tk::dnn::Conv2d ir_2_conv1(&net, 96, 1, 1, 1, 1, 0, 0, inverted_residual2[0], true);
@@ -157,14 +160,14 @@ int main()
     tk::dnn::Conv2d ir_2_conv3(&net, 24, 1, 1, 1, 1, 0, 0, inverted_residual2[2], true);
 
     //Inverted Residual 3
-    tk::dnn::Layer    *last = &ir_2_conv3;
+    tk::dnn::Layer *last = &ir_2_conv3;
     tk::dnn::Conv2d ir_3_conv1(&net, 144, 1, 1, 1, 1, 0, 0, inverted_residual3[0], true);
     tk::dnn::Activation relu_3_1(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_3_conv2(&net, 144, 3, 3, 1, 1, 1, 1, inverted_residual3[1], true, false, false, 144);
     tk::dnn::Activation relu_3_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_3_conv3(&net, 24, 1, 1, 1, 1, 0, 0, inverted_residual3[2], true);
-    
+
-    tk::dnn::Shortcut s3_0 (&net, last);
+    tk::dnn::Shortcut s3_0(&net, last);
     // //Inverted Residual 4
     tk::dnn::Conv2d ir_4_conv1(&net, 144, 1, 1, 1, 1, 0, 0, inverted_residual4[0], true);
     tk::dnn::Activation relu_4_1(&net, CUDNN_ACTIVATION_RELU);
@@ -179,8 +182,8 @@ int main()
     tk::dnn::Conv2d ir_5_conv2(&net, 192, 3, 3, 1, 1, 1, 1, inverted_residual5[1], true, false, false, 192);
     tk::dnn::Activation relu_5_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_5_conv3(&net, 32, 1, 1, 1, 1, 0, 0, inverted_residual5[2], true);
-    
+
-    tk::dnn::Shortcut s5_0 (&net, last);
+    tk::dnn::Shortcut s5_0(&net, last);
     // // // //Inverted Residual 6
     last = &s5_0;
     tk::dnn::Conv2d ir_6_conv1(&net, 192, 1, 1, 1, 1, 0, 0, inverted_residual6[0], true);
@@ -189,7 +192,7 @@ int main()
     tk::dnn::Activation relu_6_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_6_conv3(&net, 32, 1, 1, 1, 1, 0, 0, inverted_residual6[2], true);
 
-    tk::dnn::Shortcut s6_0 (&net, last);
+    tk::dnn::Shortcut s6_0(&net, last);
     //Inverted Residual 7
     tk::dnn::Conv2d ir_7_conv1(&net, 192, 1, 1, 1, 1, 0, 0, inverted_residual7[0], true);
     tk::dnn::Activation relu_7_1(&net, CUDNN_ACTIVATION_RELU);
@@ -204,8 +207,8 @@ int main()
     tk::dnn::Conv2d ir_8_conv2(&net, 384, 3, 3, 1, 1, 1, 1, inverted_residual8[1], true, false, false, 384);
     tk::dnn::Activation relu_8_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_8_conv3(&net, 64, 1, 1, 1, 1, 0, 0, inverted_residual8[2], true);
-    
+
-    tk::dnn::Shortcut s8_0 (&net, last);
+    tk::dnn::Shortcut s8_0(&net, last);
     //Inverted Residual 9
     last = &s8_0;
     tk::dnn::Conv2d ir_9_conv1(&net, 384, 1, 1, 1, 1, 0, 0, inverted_residual9[0], true);
@@ -214,7 +217,7 @@ int main()
     tk::dnn::Activation relu_9_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_9_conv3(&net, 64, 1, 1, 1, 1, 0, 0, inverted_residual9[2], true);
 
-    tk::dnn::Shortcut s9_0 (&net, last);
+    tk::dnn::Shortcut s9_0(&net, last);
     //Inverted Residual 10
     last = &s9_0;
     tk::dnn::Conv2d ir_10_conv1(&net, 384, 1, 1, 1, 1, 0, 0, inverted_residual10[0], true);
@@ -223,7 +226,7 @@ int main()
     tk::dnn::Activation relu_10_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_10_conv3(&net, 64, 1, 1, 1, 1, 0, 0, inverted_residual10[2], true);
 
-    tk::dnn::Shortcut s10_0 (&net, last);
+    tk::dnn::Shortcut s10_0(&net, last);
     //Inverted Residual 11
     tk::dnn::Conv2d ir_11_conv1(&net, 384, 1, 1, 1, 1, 0, 0, inverted_residual11[0], true);
     tk::dnn::Activation relu_11_1(&net, CUDNN_ACTIVATION_RELU);
@@ -239,7 +242,7 @@ int main()
     tk::dnn::Activation relu_12_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_12_conv3(&net, 96, 1, 1, 1, 1, 0, 0, inverted_residual12[2], true);
 
-    tk::dnn::Shortcut s12_0 (&net, last);
+    tk::dnn::Shortcut s12_0(&net, last);
     last = &s12_0;
     //Inverted Residual 13
     tk::dnn::Conv2d ir_13_conv1(&net, 576, 1, 1, 1, 1, 0, 0, inverted_residual13[0], true);
@@ -248,10 +251,10 @@ int main()
     tk::dnn::Activation relu_13_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_13_conv3(&net, 96, 1, 1, 1, 1, 0, 0, inverted_residual13[2], true);
 
-    tk::dnn::Shortcut s13_0 (&net, last);
+    tk::dnn::Shortcut s13_0(&net, last);
     // //Inverted Residual 14
     tk::dnn::Conv2d ir_14_conv1(&net, 576, 1, 1, 1, 1, 0, 0, inverted_residual14[0], true);
-    tk::dnn::Activation relu_14_1(&net, CUDNN_ACTIVATION_RELU);    
+    tk::dnn::Activation relu_14_1(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_14_conv2(&net, 576, 3, 3, 2, 2, 1, 1, inverted_residual14[1], true, false, false, 576);
     tk::dnn::Activation relu_14_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_14_conv3(&net, 160, 1, 1, 1, 1, 0, 0, inverted_residual14[2], true);
@@ -264,7 +267,7 @@ int main()
     tk::dnn::Activation relu_15_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_15_conv3(&net, 160, 1, 1, 1, 1, 0, 0, inverted_residual15[2], true);
 
-    tk::dnn::Shortcut s15_0 (&net, last);
+    tk::dnn::Shortcut s15_0(&net, last);
     //Inverted Residual 16
     last = &s15_0;
     tk::dnn::Conv2d ir_16_conv1(&net, 960, 1, 1, 1, 1, 0, 0, inverted_residual16[0], true);
@@ -273,7 +276,7 @@ int main()
     tk::dnn::Activation relu_16_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d ir_16_conv3(&net, 160, 1, 1, 1, 1, 0, 0, inverted_residual16[2], true);
 
-    tk::dnn::Shortcut s16_0 (&net, last);
+    tk::dnn::Shortcut s16_0(&net, last);
     //Inverted Residual 17
     tk::dnn::Conv2d ir_17_conv1(&net, 960, 1, 1, 1, 1, 0, 0, inverted_residual17[0], true);
     tk::dnn::Activation relu_17_1(&net, CUDNN_ACTIVATION_RELU);
@@ -284,7 +287,7 @@ int main()
     //Conv 18
     tk::dnn::Conv2d ir_18_conv1(&net, 1280, 1, 1, 1, 1, 0, 0, conv18, true);
     tk::dnn::Activation relu_18_1(&net, CUDNN_ACTIVATION_RELU);
-    tk::dnn::Layer * header_1[1] = {&relu_18_1};
+    tk::dnn::Layer *header_1[1] = {&relu_18_1};
 
     // //extras Inverted Residual 0
     tk::dnn::Conv2d e_0_conv1(&net, 256, 1, 1, 1, 1, 0, 0, extras0[0], true);
@@ -292,7 +295,7 @@ int main()
     tk::dnn::Conv2d e_0_conv2(&net, 256, 3, 3, 2, 2, 1, 1, extras0[1], true, false, false, 256);
     tk::dnn::Activation e_relu_0_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d e_0_conv3(&net, 512, 1, 1, 1, 1, 0, 0, extras0[2], true);
-    tk::dnn::Layer * header_2[1] = {&e_0_conv3};
+    tk::dnn::Layer *header_2[1] = {&e_0_conv3};
 
     // //extras Inverted Residual 1
     tk::dnn::Conv2d e_1_conv1(&net, 128, 1, 1, 1, 1, 0, 0, extras1[0], true);
@@ -300,15 +303,15 @@ int main()
     tk::dnn::Conv2d e_1_conv2(&net, 128, 3, 3, 2, 2, 1, 1, extras1[1], true, false, false, 128);
     tk::dnn::Activation e_relu_1_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d e_1_conv3(&net, 256, 1, 1, 1, 1, 0, 0, extras1[2], true);
-    tk::dnn::Layer * header_3[1] = {&e_1_conv3};
+    tk::dnn::Layer *header_3[1] = {&e_1_conv3};
-    
+
     //extras Inverted Residual 2
     tk::dnn::Conv2d e_2_conv1(&net, 128, 1, 1, 1, 1, 0, 0, extras2[0], true);
     tk::dnn::Activation e_relu_2_1(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d e_2_conv2(&net, 128, 3, 3, 2, 2, 1, 1, extras2[1], true, false, false, 128);
     tk::dnn::Activation e_relu_2_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d e_2_conv3(&net, 256, 1, 1, 1, 1, 0, 0, extras2[2], true);
-    tk::dnn::Layer * header_4[1] = {&e_2_conv3};
+    tk::dnn::Layer *header_4[1] = {&e_2_conv3};
 
     //extras Inverted Residual 3
     tk::dnn::Conv2d e_3_conv1(&net, 64, 1, 1, 1, 1, 0, 0, extras3[0], true);
@@ -316,79 +319,144 @@ int main()
     tk::dnn::Conv2d e_3_conv2(&net, 64, 3, 3, 2, 2, 1, 1, extras3[1], true, false, false, 64);
     tk::dnn::Activation e_relu_3_2(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d e_3_conv3(&net, 64, 1, 1, 1, 1, 0, 0, extras3[2], true);
-    tk::dnn::Layer * header_5[1] = {&e_3_conv3};
+    tk::dnn::Layer *header_5[1] = {&e_3_conv3};
 
     // classification header 0
-    tk::dnn::Layer * header_0[1] = {&relu_14_1};
+    tk::dnn::Layer *header_0[1] = {&relu_14_1};
-    tk::dnn::Route  rout_ch_0(&net, header_0, 1);
+    tk::dnn::Route rout_ch_0(&net, header_0, 1);
-    tk::dnn::Conv2d ch_0_conv1(&net, 576, 3, 3, 1, 1, 1, 1, classification_header0[0], true, false, false, 576,true);
+    tk::dnn::Conv2d ch_0_conv1(&net, 576, 3, 3, 1, 1, 1, 1, classification_header0[0], true, false, false, 576, true);
-    tk::dnn::Activation ch_relu_0_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation ch_relu_0_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d ch_0_conv2(&net, 126, 1, 1, 1, 1, 0, 0, classification_header0[1], false);
+    tk::dnn::Layer *conf0[1] = {&ch_0_conv2};
 
     // // classification header 1
-    tk::dnn::Route  rout_ch_1(&net, header_1, 1);
+    tk::dnn::Route rout_ch_1(&net, header_1, 1);
-    tk::dnn::Conv2d ch_1_conv1(&net, 1280, 3, 3, 1, 1, 1, 1, classification_header1[0], true, false, false, 1280,true);
+    tk::dnn::Conv2d ch_1_conv1(&net, 1280, 3, 3, 1, 1, 1, 1, classification_header1[0], true, false, false, 1280, true);
-    tk::dnn::Activation ch_relu_1_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation ch_relu_1_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d ch_1_conv2(&net, 126, 1, 1, 1, 1, 0, 0, classification_header1[1], false);
+    tk::dnn::Layer *conf1[1] = {&ch_1_conv2};
 
     // //classification header 2
-    tk::dnn::Route  rout_ch_2(&net, header_2, 1);
+    tk::dnn::Route rout_ch_2(&net, header_2, 1);
     tk::dnn::Conv2d ch_2_conv1(&net, 512, 3, 3, 1, 1, 1, 1, classification_header2[0], true, false, false, 512, true);
-    tk::dnn::Activation ch_relu_2_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation ch_relu_2_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d ch_2_conv2(&net, 126, 1, 1, 1, 1, 0, 0, classification_header2[1], false);
+    tk::dnn::Layer *conf2[1] = {&ch_2_conv2};
 
     // //classification header 3
-    tk::dnn::Route  rout_ch_3(&net, header_3, 1);
+    tk::dnn::Route rout_ch_3(&net, header_3, 1);
     tk::dnn::Conv2d ch_3_conv1(&net, 256, 3, 3, 1, 1, 1, 1, classification_header3[0], true, false, false, 256, true);
-    tk::dnn::Activation ch_relu_3_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation ch_relu_3_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d ch_3_conv2(&net, 126, 1, 1, 1, 1, 0, 0, classification_header3[1], false);
+    tk::dnn::Layer *conf3[1] = {&ch_3_conv2};
 
     // //classification header 4
-    tk::dnn::Route  rout_ch_4(&net, header_4, 1);
+    tk::dnn::Route rout_ch_4(&net, header_4, 1);
     tk::dnn::Conv2d ch_4_conv1(&net, 256, 3, 3, 1, 1, 1, 1, classification_header4[0], true, false, false, 256, true);
-    tk::dnn::Activation ch_relu_4_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation ch_relu_4_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d ch_4_conv2(&net, 126, 1, 1, 1, 1, 0, 0, classification_header4[1], false);
+    tk::dnn::Layer *conf4[1] = {&ch_4_conv2};
 
     // //classification header 5
-    tk::dnn::Route  rout_ch_5(&net, header_5, 1);
+    tk::dnn::Route rout_ch_5(&net, header_5, 1);
-    tk::dnn::Conv2d ch_5_conv(&net, 126, 1, 1, 1, 1, 0, 0, classification_header5, false);
+    tk::dnn::Conv2d ch_5_conv(&net, 126, 1, 1, 1, 1, 0, 0, classification_header5, false, false, true);
+    tk::dnn::Layer *conf5[1] = {&ch_5_conv};
 
     //regression header 0
-    tk::dnn::Route  rout_rh_0(&net, header_0, 1);
+    tk::dnn::Route rout_rh_0(&net, header_0, 1);
     tk::dnn::Conv2d rh_0_conv1(&net, 576, 3, 3, 1, 1, 1, 1, regression_header0[0], true, false, false, 576, true);
-    tk::dnn::Activation rh_relu_0_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation rh_relu_0_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d rh_0_conv2(&net, 24, 1, 1, 1, 1, 0, 0, regression_header0[1], false);
+    tk::dnn::Layer *loc0[1] = {&rh_0_conv2};
 
     // //regression header 1
-    tk::dnn::Route  rout_rh_1(&net, header_1, 1);
+    tk::dnn::Route rout_rh_1(&net, header_1, 1);
     tk::dnn::Conv2d rh_1_conv1(&net, 1280, 3, 3, 1, 1, 1, 1, regression_header1[0], true, false, false, 1280, true);
-    tk::dnn::Activation rh_relu_1_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation rh_relu_1_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d rh_1_conv2(&net, 24, 1, 1, 1, 1, 0, 0, regression_header1[1], false);
+    tk::dnn::Layer *loc1[1] = {&rh_1_conv2};
 
     //regression header 2
-    tk::dnn::Route  rout_rh_2(&net, header_2, 1);
+    tk::dnn::Route rout_rh_2(&net, header_2, 1);
     tk::dnn::Conv2d rh_2_conv1(&net, 512, 3, 3, 1, 1, 1, 1, regression_header2[0], true, false, false, 512, true);
-    tk::dnn::Activation rh_relu_2_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation rh_relu_2_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d rh_2_conv2(&net, 24, 1, 1, 1, 1, 0, 0, regression_header2[1], false);
+    tk::dnn::Layer *loc2[1] = {&rh_2_conv2};
 
     //regression header 3
-    tk::dnn::Route  rout_rh_3(&net, header_3, 1);
+    tk::dnn::Route rout_rh_3(&net, header_3, 1);
     tk::dnn::Conv2d rh_3_conv1(&net, 256, 3, 3, 1, 1, 1, 1, regression_header3[0], true, false, false, 256, true);
-    tk::dnn::Activation rh_relu_3_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation rh_relu_3_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d rh_3_conv2(&net, 24, 1, 1, 1, 1, 0, 0, regression_header3[1], false);
+    tk::dnn::Layer *loc3[1] = {&rh_3_conv2};
 
     //regression header 4
 
-    tk::dnn::Route  rout_rh_4(&net, header_4, 1);
+    tk::dnn::Route rout_rh_4(&net, header_4, 1);
     tk::dnn::Conv2d rh_4_conv1(&net, 256, 3, 3, 1, 1, 1, 1, regression_header4[0], true, false, false, 256, true);
-    tk::dnn::Activation rh_relu_4_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU,6);
+    tk::dnn::Activation rh_relu_4_1(&net, CUDNN_ACTIVATION_CLIPPED_RELU, 6);
     tk::dnn::Conv2d rh_4_conv2(&net, 24, 1, 1, 1, 1, 0, 0, regression_header4[1], false);
+    tk::dnn::Layer *loc4[1] = {&rh_4_conv2};
 
     //regression header 5
-    tk::dnn::Route  rout_rh_5(&net, header_5, 1);
+    tk::dnn::Route rout_rh_5(&net, header_5, 1);
-    tk::dnn::Conv2d rh_5_conv(&net, 24, 1, 1, 1, 1, 0, 0, regression_header5, false);
+    tk::dnn::Conv2d rh_5_conv(&net, 24, 1, 1, 1, 1, 0, 0, regression_header5, false, false, true);
+    tk::dnn::Layer *loc5[1] = {&rh_5_conv};
 
-    //flatten confidence and flatten locations
+    last = &rh_5_conv;
+
+    //flatten all confidence
+    tk::dnn::Route r_conf_0(&net, conf0, 1);
+    tk::dnn::Flatten fl_c_0(&net);
+    tk::dnn::Route r_conf_1(&net, conf1, 1);
+    tk::dnn::Flatten fl_c_1(&net);
+    tk::dnn::Route r_conf_2(&net, conf2, 1);
+    tk::dnn::Flatten fl_c_2(&net);
+    tk::dnn::Route r_conf_3(&net, conf3, 1);
+    tk::dnn::Flatten fl_c_3(&net);
+    tk::dnn::Route r_conf_4(&net, conf4, 1);
+    tk::dnn::Flatten fl_c_4(&net);
+    tk::dnn::Route r_conf_5(&net, conf5, 1);
+    tk::dnn::Flatten fl_c_5(&net);
+
+    // //flatten all locations
+    tk::dnn::Route r_loc_0(&net, loc0, 1);
+    tk::dnn::Flatten fl_l_0(&net);
+    tk::dnn::Route r_loc_1(&net, loc1, 1);
+    tk::dnn::Flatten fl_l_1(&net);
+    tk::dnn::Route r_loc_2(&net, loc2, 1);
+    tk::dnn::Flatten fl_l_2(&net);
+    tk::dnn::Route r_loc_3(&net, loc3, 1);
+    tk::dnn::Flatten fl_l_3(&net);
+    tk::dnn::Route r_loc_4(&net, loc4, 1);
+    tk::dnn::Flatten fl_l_4(&net);
+    tk::dnn::Route r_loc_5(&net, loc5, 1);
+    tk::dnn::Flatten fl_l_5(&net);
+
+    // //concat confidence + softmax
+    tk::dnn::Layer *confidences[6] = {&fl_c_0, &fl_c_1, &fl_c_2, &fl_c_3, &fl_c_4, &fl_c_5};
+    tk::dnn::Route rout_conf(&net, confidences, 6);
+    tk::dnn::dataDim_t olddim_c = net.layers[net.num_layers - 1]->output_dim;
+    tk::dnn::dataDim_t dim_resh(1, olddim_c.c * olddim_c.h * olddim_c.w / classes, classes, 1, 1);
+
+    tk::dnn::Reshape reshape_conf1(&net, dim_resh);
+    tk::dnn::Flatten fl_l_6(&net);
+    tk::dnn::dataDim_t newdim_c(1, classes, olddim_c.c * olddim_c.h * olddim_c.w / classes, 1, 1);
+
+    tk::dnn::Reshape reshape_conf2(&net, newdim_c);
+
+    tk::dnn::Softmax sm_1(&net, &newdim_c, true);
+    // tk::dnn::Flatten fl_l_7(&net);
+    // tk::dnn::Reshape reshape_conf3(&net,dim_resh, true);
+    tk::dnn::Layer *conf = &sm_1;
+
+    //concat locations
+    tk::dnn::Layer *locations[6] = {&fl_l_0, &fl_l_1, &fl_l_2, &fl_l_3, &fl_l_4, &fl_l_5};
+    tk::dnn::Route rout_loc(&net, locations, 6);
+    tk::dnn::dataDim_t olddim_l = net.layers[net.num_layers - 1]->output_dim;
+    tk::dnn::dataDim_t newdim_l(1, olddim_l.c * olddim_l.h * olddim_l.w / 4, 1, 4, 1);
+    tk::dnn::Reshape reshape_loc(&net, newdim_l, true);
+    tk::dnn::Layer *loc = &reshape_loc;
 
     // Load input
     dnnType *data;
@@ -399,14 +467,9 @@ int main()
     //print network model
     net.print();
 
-    //convert network to tensorRT
+    // convert network to tensorRT
     tk::dnn::NetworkRT netRT(&net, "mobilenetv2ssd.rt");
 
-    
-    tk::dnn::dataDim_t out_dim;
-    out_dim = net.layers[net.num_layers-1]->output_dim;
-    dnnType *cudnn_out, *rt_out;
-
     tk::dnn::dataDim_t dim1 = dim; //input dim
     printCenteredTitle(" CUDNN inference ", '=', 30);
     {
@@ -416,9 +479,11 @@ int main()
         TIMER_STOP
         dim1.print();
     }
-    cudnn_out = net.layers[net.num_layers-1]->dstData;
 
-    printDeviceVector(64, cudnn_out, true);
+    dnnType *cudnn_out1 = conf5[0]->dstData;
+    tk::dnn::dataDim_t out_dim1 = conf5[0]->output_dim;
+    dnnType *cudnn_out2 = loc5[0]->dstData;
+    tk::dnn::dataDim_t out_dim2 = loc5[0]->output_dim;
 
     tk::dnn::dataDim_t dim2 = dim;
     printCenteredTitle(" TENSORRT inference ", '=', 30);
@@ -429,20 +494,48 @@ int main()
         TIMER_STOP
         dim2.print();
     }
-    rt_out = (dnnType *)netRT.buffersRT[1];
 
+    dnnType *rt_out1 = (dnnType *)netRT.buffersRT[1];
+    dnnType *rt_out2 = (dnnType *)netRT.buffersRT[2];
+    dnnType *rt_out3 = (dnnType *)netRT.buffersRT[3];
+    dnnType *rt_out4 = (dnnType *)netRT.buffersRT[4];
 
     printCenteredTitle(std::string(" RESNET CHECK RESULTS ").c_str(), '=', 30);
-    dnnType *out, *out_h;
+    dnnType *out1, *out1_h;
-    int odim = out_dim.tot();
+    int odim1 = out_dim1.tot();
-    readBinaryFile(output_bin, odim, &out_h, &out);
+    readBinaryFile(output_bin1, odim1, &out1_h, &out1);
-    std::cout << "CUDNN vs correct";
-    checkResult(odim, cudnn_out, out);
 
-    std::cout << "TRT   vs correct";
+    dnnType *out2, *out2_h;
-    checkResult(odim, rt_out, out);
+    int odim2 = out_dim2.tot();
-    std::cout << "CUDNN vs TRT    ";
+    readBinaryFile(output_bin2, odim2, &out2_h, &out2);
-    checkResult(odim, cudnn_out, rt_out);
+    std::cout << "CUDNN vs correct" << std::endl;
+    checkResult(odim1, cudnn_out1, out1);
+    checkResult(odim2, cudnn_out2, out2);
+
+    std::cout << "TRT   vs correct" << std::endl;
+    checkResult(odim1, rt_out1, out1);
+    checkResult(odim2, rt_out2, out2);
+
+    std::cout << "CUDNN vs TRT    " << std::endl;
+    checkResult(odim1, cudnn_out1, rt_out1);
+    checkResult(odim2, cudnn_out2, rt_out2);
+
+    std::cout << "---------------------------------------------------" << std::endl;
+    std::cout << "Confidence CUDNN" << std::endl;
+    printDeviceVector(64, conf->dstData, true);
+    std::cout << "Locations CUDNN" << std::endl;
+    printDeviceVector(64, loc->dstData, true);
+    std::cout << "---------------------------------------------------" << std::endl;
+
+    std::cout << "Confidence tensorRT" << std::endl;
+    printDeviceVector(64, rt_out3, true);
+    std::cout << "Locations tensorRT" << std::endl;
+    printDeviceVector(64, rt_out4, true);
+    std::cout << "---------------------------------------------------" << std::endl;
+
+    std::cout << "CUDNN vs TRT    " << std::endl;
+    checkResult(conf->output_dim.tot(), conf->dstData, rt_out3);
+    checkResult(loc->output_dim.tot(), loc->dstData, rt_out4);
 
     return 0;
 }