diff --git a/include/Layer.h b/include/Layer.h index 5b9ace5..369551d 100644 --- a/include/Layer.h +++ b/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 { diff --git a/src/Yolo.cpp b/src/Yolo.cpp new file mode 100644 index 0000000..81236aa --- /dev/null +++ b/src/Yolo.cpp @@ -0,0 +1,61 @@ +#include + +#ifdef OPENCV + #include + #include + #include +#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; +} + +}} diff --git a/tests/yolo3_berkeley/yolo3_berkeley.cpp b/tests/yolo3_berkeley/yolo3_berkeley.cpp index c0b3689..441ec48 100644 --- a/tests/yolo3_berkeley/yolo3_berkeley.cpp +++ b/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;