diff --git a/include/Network.h b/include/Network.h index c7045dc..fb898a6 100644 --- a/include/Network.h +++ b/include/Network.h @@ -5,16 +5,31 @@ namespace tkDNN { +struct dataDim_t; +class Layer; +const int MAX_LAYERS = 256; + class Network { public: Network(); virtual ~Network(); + /** + Do inferece for every added layer + */ + value_type* infer(dataDim_t &dim, value_type* data); + + bool addLayer(Layer *l); + cudnnDataType_t dataType; cudnnTensorFormat_t tensorFormat; cudnnHandle_t cudnnHandle; cublasHandle_t cublasHandle; + +private: + Layer* layers[MAX_LAYERS]; //contains layers of the net + int num_layers; //current number of layers }; } diff --git a/src/Layer.cpp b/src/Layer.cpp index 1629812..7e689ee 100644 --- a/src/Layer.cpp +++ b/src/Layer.cpp @@ -12,6 +12,9 @@ Layer::Layer(Network *net, dataDim_t in_dim) { checkCUDNN( cudnnCreateTensorDescriptor(&srcTensorDesc) ); checkCUDNN( cudnnCreateTensorDescriptor(&dstTensorDesc) ); + + if(!net->addLayer(this)) + FatalError("Net reached max number of layers"); } Layer::~Layer() { diff --git a/src/Network.cpp b/src/Network.cpp index 88ccc58..3fafb60 100644 --- a/src/Network.cpp +++ b/src/Network.cpp @@ -1,6 +1,7 @@ #include #include "Network.h" +#include "Layer.h" namespace tkDNN { @@ -12,6 +13,8 @@ Network::Network() { checkCUDNN( cudnnCreate(&cudnnHandle) ); checkERROR( cublasCreate(&cublasHandle) ); + + num_layers = 0; } Network::~Network() { @@ -20,4 +23,21 @@ Network::~Network() { checkERROR( cublasDestroy(cublasHandle) ); } +value_type* Network::infer(dataDim_t &dim, value_type* data) { + + //do infer for every layer + for(int i=0; iinfer(dim, data); + + return data; +} + +bool Network::addLayer(Layer *l) { + if(num_layers == MAX_LAYERS) + return false; + + layers[num_layers++] = l; + return true; +} + } \ No newline at end of file diff --git a/tests/test.cpp b/tests/test.cpp index 3c24101..cf9cad5 100644 --- a/tests/test.cpp +++ b/tests/test.cpp @@ -29,11 +29,15 @@ int main() { TIMER_START // Inference + data = net.infer(dim, data); dim.print(); + + /* + //old Inference method data = c0.infer(dim, data); dim.print(); data = a0.infer(dim, data); dim.print(); data = c1.infer(dim, data); dim.print(); data = a1.infer(dim, data); dim.print(); - + */ TIMER_STOP // Print result printDeviceVector(dim.tot(), data);