From a17e7800b9c969724e5ca1af03d3d2b2edd49964 Mon Sep 17 00:00:00 2001
From: Micaela Verucchi <micaelaverucchi@gmail.com>
Date: Mon, 23 Nov 2020 13:07:54 +0100
Subject: [PATCH] Add computation of #parameters, #MACC, and max feature map
 size in the tests

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
---
 include/tkDNN/Layer.h   |  4 ++++
 include/tkDNN/Network.h |  1 +
 src/Conv2d.cpp          |  5 +++++
 src/DeformConv2d.cpp    |  6 ++++++
 src/Layer.cpp           |  2 ++
 src/LayerWgs.cpp        |  5 ++++-
 src/Network.cpp         | 36 ++++++++++++++++++++++++++++++++++++
 7 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h
index 790a431..91a6af0 100644
--- a/include/tkDNN/Layer.h
+++ b/include/tkDNN/Layer.h
@@ -54,6 +54,10 @@ public:
 
     int id = 0;
     bool final;        //if the layer is the final one
+    uint n_params = 0;
+    uint feature_map_size = 0;
+    long unsigned MACC = 0;
+
 
     std::string getLayerName() {
         layerType_t type = getLayerType();
diff --git a/include/tkDNN/Network.h b/include/tkDNN/Network.h
index b78acff..6edf248 100644
--- a/include/tkDNN/Network.h
+++ b/include/tkDNN/Network.h
@@ -50,6 +50,7 @@ public:
     bool addLayer(Layer *l);
     void print();
     const char *getNetworkRTName(const char *network_name);
+    void adjustFeatureMapSizeWithShortcuts();
 
     cudnnDataType_t dataType;
     cudnnTensorFormat_t tensorFormat;
diff --git a/src/Conv2d.cpp b/src/Conv2d.cpp
index b57cf58..2901a70 100644
--- a/src/Conv2d.cpp
+++ b/src/Conv2d.cpp
@@ -166,6 +166,11 @@ Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
     }
     initCUDNN(deConv);
 
+    if(this->groups != 1)
+        MACC = kernelH*kernelW*output_dim.c*output_dim.w*output_dim.h;
+    else
+        MACC = input_dim.c*kernelH*kernelW*output_dim.c*output_dim.w*output_dim.h;
+
     // allocate warkspace
     if (ws_sizeInBytes!=0) {
         checkCuda( cudaMalloc(&workSpace, ws_sizeInBytes) );
diff --git a/src/DeformConv2d.cpp b/src/DeformConv2d.cpp
index dbb71e1..826cbb0 100644
--- a/src/DeformConv2d.cpp
+++ b/src/DeformConv2d.cpp
@@ -73,6 +73,12 @@ DeformConv2d::DeformConv2d( Network *net, int out_ch, int deformable_group, int
     
     output_dim.c = out_ch;
     initCUDNN();
+
+    if(this->deformableGroup != 1)
+        MACC = kernelH*kernelW*output_dim.c*output_dim.w*output_dim.h;
+    else
+        MACC = input_dim.c*kernelH*kernelW*output_dim.c*output_dim.w*output_dim.h;
+        
     //allocate data for infer result
     checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
 }
diff --git a/src/Layer.cpp b/src/Layer.cpp
index a355b90..7e13a48 100644
--- a/src/Layer.cpp
+++ b/src/Layer.cpp
@@ -18,6 +18,8 @@ Layer::Layer(Network *net) {
         if(!net->addLayer(this))
             FatalError("Net reached max number of layers");    
     }
+
+    feature_map_size = input_dim.tot() + output_dim.tot();
 }
 
 Layer::~Layer() {
diff --git a/src/LayerWgs.cpp b/src/LayerWgs.cpp
index a761327..2f875a4 100644
--- a/src/LayerWgs.cpp
+++ b/src/LayerWgs.cpp
@@ -19,6 +19,8 @@ LayerWgs::LayerWgs(Network *net, int inputs, int outputs,
     int seek = 0;
     readBinaryFile(weights_path.c_str(), inputs*outputs*kh*kw*kl, &data_h, &data_d, seek);
     seek += inputs*outputs*kh*kw*kl;
+    n_params = seek;
+    
     this->additional_bias = additional_bias;
     if(additional_bias) {
         readBinaryFile(weights_path.c_str(), outputs, &bias2_h, &bias2_d, seek);
@@ -26,15 +28,16 @@ LayerWgs::LayerWgs(Network *net, int inputs, int outputs,
     }
     
     readBinaryFile(weights_path.c_str(), outputs, &bias_h, &bias_d, seek);
+    seek += outputs;
 
     this->batchnorm = batchnorm;
     if(batchnorm) {
-        seek += outputs;
         readBinaryFile(weights_path.c_str(), outputs, &scales_h, &scales_d, seek);
         seek += outputs;
         readBinaryFile(weights_path.c_str(), outputs, &mean_h, &mean_d, seek);
         seek += outputs;
         readBinaryFile(weights_path.c_str(), outputs, &variance_h, &variance_d, seek);
+        seek += outputs;
 
         float eps = TKDNN_BN_MIN_EPSILON;
 
diff --git a/src/Network.cpp b/src/Network.cpp
index 7fa291f..2c3c8c7 100644
--- a/src/Network.cpp
+++ b/src/Network.cpp
@@ -96,6 +96,28 @@ dataDim_t Network::getOutputDim() {
             return layers[num_layers-1]->output_dim;
 }
 
+void Network::adjustFeatureMapSizeWithShortcuts(){
+    layerType_t layer_type;
+    int shortcutted_idx;
+
+    for(int i=0; i<num_layers; i++) {
+        layer_type = layers[i]->getLayerType();
+        if(layer_type == LAYER_SHORTCUT){
+            shortcutted_idx = -1;
+            for(int j=0; j<num_layers; j++) {
+                if(static_cast<tk::dnn::Shortcut*>(layers[i])->backLayer == layers[j]){
+                    shortcutted_idx = j;
+                    break;
+                }
+            }
+            if(shortcutted_idx == -1)
+                FatalError("Problem when computing featuer_map_size with shortcuts");
+            for(int j=shortcutted_idx+1; j<i; ++j)
+                layers[j]->feature_map_size += layers[shortcutted_idx]->output_dim.tot();
+        }
+    }
+}
+
 void Network::print() {
 
     printCenteredTitle(" NETWORK MODEL ", '=', 60);
@@ -106,10 +128,21 @@ void Network::print() {
     std::cout.width(16); std::cout<<std::left<<"output (H*W,CH)";
     std::cout<<"\n";
 
+    adjustFeatureMapSizeWithShortcuts();
+    
+    long long unsigned int tot_params = 0;
+    long long unsigned int max_feature_map_size = 0;
+    long long unsigned int tot_MACC = 0;
+    
     for(int i=0; i<num_layers; i++) {
         dataDim_t in = layers[i]->input_dim;
         dataDim_t out = layers[i]->output_dim;
 
+        tot_params += layers[i]->n_params;
+        tot_MACC += layers[i]->MACC;
+        if(layers[i]->feature_map_size> max_feature_map_size)
+            max_feature_map_size = layers[i]->feature_map_size;
+
         std::cout.width(3); std::cout<<std::right<<i;
         std::cout<<" ";
         std::cout.width(16); std::cout<<std::left<<layers[i]->getLayerName();
@@ -128,6 +161,9 @@ void Network::print() {
     }
     printCenteredTitle("", '=', 60);
     std::cout<<"\n";
+    std::cout<<"N params: "<<tot_params<<std::endl;
+    std::cout<<"Max feature map size: "<<max_feature_map_size<<std::endl;
+    std::cout<<"N MACC: "<<tot_MACC<<std::endl<<std::endl;
     printCudaMemUsage();
 }
 const char *Network::getNetworkRTName(const char *network_name){