yolo layer

include/Layer.h

@@ -20,6 +20,7 @@ enum layerType_t {
     LAYER_SHORTCUT,
     LAYER_UPSAMPLE,
     LAYER_REGION,
+    LAYER_YOLO
 };
 
 /** 
@@ -55,6 +56,7 @@ public:
             case LAYER_SHORTCUT:    return "Shortcut";
             case LAYER_UPSAMPLE:    return "Upsample";
             case LAYER_REGION:      return "Region";
+            case LAYER_YOLO:        return "Yolo";
             default:                return "unknown";
         }
     }
@@ -330,9 +332,23 @@ struct sortable_bbox {
     float **probs;
 };
 
+/**
+    Yolo3 layer
+*/
+class Yolo : public Layer {
+
+public:
+    Yolo(Network *net, int classes, int num);
+    virtual ~Yolo();
+    virtual layerType_t getLayerType() { return LAYER_YOLO; };
+
+    int classes, num;
+    
+    virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+};
+
 /**
     Region layer
-    Mantain same dimension but change C*H*W distribution
 */
 class Region : public Layer {
 

src/Yolo.cpp

@@ -0,0 +1,61 @@
+#include <iostream>
+
+#ifdef OPENCV
+    #include <opencv2/core/core.hpp>
+    #include <opencv2/highgui/highgui.hpp>
+    #include <opencv2/imgproc/imgproc.hpp>
+#endif
+
+#include "Layer.h"
+#include "kernels.h"
+
+namespace tk { namespace dnn {
+
+Yolo::Yolo(Network *net, int classes, int num) : 
+    Layer(net) {
+
+    this->classes = classes;
+    this->num = num;
+
+    // same
+    output_dim.n = input_dim.n;
+    output_dim.c = input_dim.c;
+    output_dim.h = input_dim.h;
+    output_dim.w = input_dim.w;
+    output_dim.l = input_dim.l;
+
+    checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(dnnType)) );
+}
+
+Yolo::~Yolo() {
+    checkCuda( cudaFree(dstData) );
+}
+
+int entry_index(int batch, int location, int entry, 
+            int classes, dataDim_t &input_dim, dataDim_t &output_dim) {
+    int n =   location / (input_dim.w*input_dim.h);
+    int loc = location % (input_dim.w*input_dim.h);
+    return batch*output_dim.tot() + n*input_dim.w*input_dim.h*(4+classes+1) +
+           entry*input_dim.w*input_dim.h + loc;
+}
+
+
+dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) {
+
+    checkCuda( cudaMemcpy(dstData, srcData, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+
+    for (int b = 0; b < dim.n; ++b){
+        for(int n = 0; n < num; ++n){
+            int index = entry_index(b, n*dim.w*dim.h, 0, classes, input_dim, output_dim);
+            activationLOGISTICForward(srcData + index, dstData + index, 2*dim.w*dim.h);
+            
+            index = entry_index(b, n*dim.w*dim.h, 4, classes, input_dim, output_dim);
+            activationLOGISTICForward(srcData + index, dstData + index, (1+classes)*dim.w*dim.h);
+        }
+    }
+
+    dim = output_dim;
+    return dstData;
+}
+
+}}

tests/yolo3_berkeley/yolo3_berkeley.cpp

@@ -61,7 +61,8 @@ const char *c78_bin     = "../tests/yolo3_berkeley/layers/c78.bin";
 const char *c79_bin     = "../tests/yolo3_berkeley/layers/c79.bin";
 const char *c80_bin     = "../tests/yolo3_berkeley/layers/c80.bin";
 const char *c81_bin     = "../tests/yolo3_berkeley/layers/c81.bin";
-const char *output_bin = "../tests/yolo3_berkeley/debug/layer81_out.bin";
+const char *c84_bin     = "../tests/yolo3_berkeley/layers/c84.bin";
+const char *output_bin = "../tests/yolo3_berkeley/debug/layer84_out.bin";
 
 int main() {
 
@@ -218,6 +219,13 @@ int main() {
     tk::dnn::Conv2d     c80  (&net,1024, 3, 3, 1, 1, 1, 1, c80_bin, true);
     tk::dnn::Activation a80  (&net, tk::dnn::ACTIVATION_LEAKY);
     tk::dnn::Conv2d     c81  (&net,  45, 1, 1, 1, 1, 0, 0, c81_bin, false);
+    tk::dnn::Yolo       g82  (&net,  10, 9);
+
+    tk::dnn::Layer *m83_layers[1] = { &a79 };
+    tk::dnn::Route      m83  (&net, m83_layers, 1);
+    tk::dnn::Conv2d     c84  (&net, 256, 1, 1, 1, 1, 0, 0, c84_bin, true);
+    tk::dnn::Activation a84  (&net, tk::dnn::ACTIVATION_LEAKY);
+
 
     // Load input
     dnnType *data;