added support for both reorg and reorg3d layers of AlexeyAB/Darknet

include/tkDNN/Layer.h

@@ -534,13 +534,14 @@ public:
 class Reorg : public Layer {
 
 public:
-    Reorg(Network *net, int stride);
+    Reorg(Network *net, int stride, bool reorg3d = false);
     virtual ~Reorg();
     virtual layerType_t getLayerType() { return LAYER_REORG; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 
     int stride;
+    bool reorg3d;
 };
 
 /**

include/tkDNN/pluginsRT/ReorgRT.h

@@ -4,8 +4,9 @@
 class ReorgRT : public IPlugin {
 
 public:
-	ReorgRT(int stride) {
+	ReorgRT(int stride, bool reorg3d = false) {
 		this->stride = stride;
+		this->reorg3d = reorg3d;
 	}
 
 	~ReorgRT(){
@@ -21,9 +22,15 @@ public:
 	}
 
 	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
-		c = outputDims[0].d[0];
-		h = outputDims[0].d[1];
-		w = outputDims[0].d[2];
+		if (reorg3d) {
+			c = outputDims[0].d[0];
+			h = outputDims[0].d[1];
+			w = outputDims[0].d[2];
+		} else {
+			c = inputDims[0].d[0];
+			h = inputDims[0].d[1];
+			w = inputDims[0].d[2];
+		}
 	}
 
 	int initialize() override {
@@ -60,4 +67,5 @@ public:
 	}
 
 	int c, h, w, stride;
+	bool reorg3d;
 };

src/DarknetParser.cpp

@@ -153,7 +153,7 @@ namespace tk { namespace dnn {
             netLayers.push_back(new tk::dnn::Route(net, layers.data(), layers.size()));
 
         } else if(f.type == "reorg" || f.type == "reorg3d") {
-            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x));
+            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x, f.type == "reorg3d"));
 
         } else if(f.type == "region") {
             netLayers.push_back(new tk::dnn::Region(net, f.classes, f.coords, f.num));

src/NetworkRT.cpp

@@ -481,7 +481,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
     //std::cout<<"convert Reorg\n";
 
     //std::cout<<"New plugin REORG\n";
-    IPlugin *plugin = new ReorgRT(l->stride);
+    IPlugin *plugin = new ReorgRT(l->stride, l->reorg3d);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);
     return lRT;

src/Reorg.cpp

@@ -5,9 +5,10 @@
 
 namespace tk { namespace dnn {
 
-Reorg::Reorg(Network *net, int stride) : Layer(net) {
+Reorg::Reorg(Network *net, int stride, bool reorg3d) : Layer(net) {
 
     this->stride = stride;
+    this->reorg3d = reorg3d;
     
     output_dim.n = input_dim.n;
     output_dim.c = input_dim.c*stride*stride;
@@ -24,8 +25,10 @@ Reorg::~Reorg() {
 }
 
 dnnType* Reorg::infer(dataDim_t &dim, dnnType* srcData) {
-
-    reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride);
+    if (reorg3d)
+        reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride);
+    else
+        reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
 
     dim = output_dim;
     return dstData;