better network model

src/Activation.cpp

@@ -5,8 +5,8 @@
 
 namespace tkDNN {
 
-Activation::Activation(Network *net, dataDim_t input_dim, int act_mode) : 
-    Layer(net, input_dim) {
+Activation::Activation(Network *net, int act_mode) : 
+    Layer(net) {
 
     this->act_mode = act_mode;
     checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(value_type)) );

src/Conv2d.cpp

@@ -4,12 +4,11 @@
 
 namespace tkDNN {
 
-Conv2d::Conv2d( Network *net, dataDim_t in_dim, int out_ch, 
-                int kernelH, int kernelW, int strideH, int strideW,
-                int paddingH, int paddingW,
+Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW, 
+                int strideH, int strideW, int paddingH, int paddingW,
                 const char* fname_weights, bool batchnorm) : 
     
-    LayerWgs(net, in_dim, in_dim.c, out_ch, kernelH, kernelW, 1, 
+    LayerWgs(net, net->getOutputDim().c, out_ch, kernelH, kernelW, 1, 
              fname_weights, batchnorm) {
 
     this->kernelH = kernelH;

src/Dense.cpp

@@ -4,9 +4,8 @@
 
 namespace tkDNN {
 
-Dense::Dense(Network *net, dataDim_t in_dim, 
-    int out_ch, const char* fname_weights) : 
-    LayerWgs(net, in_dim, in_dim.tot(), out_ch, 1, 1, 1, fname_weights) {
+Dense::Dense(Network *net, int out_ch, const char* fname_weights) : 
+    LayerWgs(net, net->getOutputDim().tot(), out_ch, 1, 1, 1, fname_weights) {
 
     output_dim.n = 1;
     output_dim.c = out_ch;

src/Flatten.cpp

@@ -5,8 +5,7 @@
 
 namespace tkDNN {
 
-Flatten::Flatten(Network *net, dataDim_t input_dim) : 
-    Layer(net, input_dim) {
+Flatten::Flatten(Network *net) : Layer(net) {
 
     checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(value_type)) );
 

src/Layer.cpp

@@ -4,11 +4,11 @@
 
 namespace tkDNN {
 
-Layer::Layer(Network *net, dataDim_t in_dim) {
+Layer::Layer(Network *net) {
 
     this->net = net;
-    this->input_dim = in_dim;
-    this->output_dim = in_dim;
+    this->input_dim = net->getOutputDim();
+    this->output_dim = input_dim;
     
     checkCUDNN( cudnnCreateTensorDescriptor(&srcTensorDesc) );
     checkCUDNN( cudnnCreateTensorDescriptor(&dstTensorDesc) );

src/LayerWgs.cpp

@@ -4,9 +4,9 @@
 
 namespace tkDNN {
 
-LayerWgs::LayerWgs(Network *net, dataDim_t in_dim, 
-                int inputs, int outputs, int kh, int kw, int kl, 
-                const char* fname_weights, bool batchnorm) : Layer(net, in_dim) {
+LayerWgs::LayerWgs(Network *net, int inputs, int outputs, 
+                   int kh, int kw, int kl, 
+                   const char* fname_weights, bool batchnorm) : Layer(net) {
 
     this->inputs  = inputs;
     this->outputs = outputs;    

src/MulAdd.cpp

@@ -5,8 +5,7 @@
 
 namespace tkDNN {
 
-MulAdd::MulAdd(Network *net, dataDim_t input_dim, value_type mul, value_type add) : 
-    Layer(net, input_dim) {
+MulAdd::MulAdd(Network *net, value_type mul, value_type add) : Layer(net) {
 
     this->mul = mul;
     this->add = add;

src/Network.cpp

@@ -7,7 +7,8 @@
 
 namespace tkDNN {
 
-Network::Network() {
+Network::Network(dataDim_t input_dim) {
+    this->input_dim = input_dim;
 
     float tk_ver = float(tkDNN::getVersion())/1000;
     float cu_ver = float(cudnnGetVersion())/1000;
@@ -47,4 +48,12 @@ bool Network::addLayer(Layer *l) {
     return true;
 }
 
+dataDim_t Network::getOutputDim() {
+
+        if(num_layers == 0)
+            return input_dim;
+        else
+            return layers[num_layers-1]->output_dim;
+}
+
 }

src/Pooling.cpp

@@ -5,9 +5,9 @@
 
 namespace tkDNN {
 
-Pooling::Pooling(   Network *net, dataDim_t input_dim, 
-                    int winH, int winW, int strideH, int strideW, tkdnnPoolingMode_t pool_mode) : 
-    Layer(net, input_dim) {
+Pooling::Pooling(   Network *net, int winH, int winW, 
+                    int strideH, int strideW, tkdnnPoolingMode_t pool_mode) : 
+    Layer(net) {
 
 
     if(winH != strideH || winW != strideW)

src/Region.cpp

@@ -5,9 +5,8 @@
 
 namespace tkDNN {
 
-Region::Region(Network *net, dataDim_t input_dim, 
-               int classes, int coords, int num, float thresh) : 
-    Layer(net, input_dim) {
+Region::Region(Network *net, int classes, int coords, int num, float thresh) : 
+    Layer(net) {
 
     this->classes = classes;
     this->coords = coords;

src/Reorg.cpp

@@ -5,8 +5,7 @@
 
 namespace tkDNN {
 
-Reorg::Reorg(Network *net, dataDim_t input_dim, int stride) : 
-    Layer(net, input_dim) {
+Reorg::Reorg(Network *net, int stride) : Layer(net) {
 
     this->stride = stride;
     

src/Route.cpp

@@ -5,17 +5,11 @@
 
 namespace tkDNN {
 
-Route::Route(Network *net, int *layers_id, int layers_n) : 
-    Layer(net, dataDim_t()) {
+Route::Route(Network *net, Layer **layers, int layers_n) : Layer(net) {
 
+    this->layers = layers;
     this->layers_n = layers_n;
 
-    //get layers
-    layers = new Layer*[layers_n];
-    for(int i=0; i<layers_n; i++) 
-        layers[i] = net->layers[layers_id[i]]; 
-        
-
     //get dims
     output_dim.l = 1;
     output_dim.c = 0;

src/Softmax.cpp

@@ -5,8 +5,7 @@
 
 namespace tkDNN {
 
-Softmax::Softmax(Network *net, dataDim_t input_dim) : 
-    Layer(net, input_dim) {
+Softmax::Softmax(Network *net) : Layer(net) {
 
     checkCuda( cudaMalloc(&dstData, input_dim.tot()*sizeof(value_type)) );