From dee7f8132518b38327c7c1493190dbc7224c9177 Mon Sep 17 00:00:00 2001
From: tensor-tang <jian.j.tang@intel.com>
Date: Mon, 16 Oct 2017 13:54:08 +0800
Subject: [PATCH] add finish work of mkldnn

---
 paddle/gserver/gradientmachines/NeuralNetwork.cpp | 15 +++++++++++++++
 paddle/gserver/gradientmachines/NeuralNetwork.h   |  3 +++
 2 files changed, 18 insertions(+)

diff --git a/paddle/gserver/gradientmachines/NeuralNetwork.cpp b/paddle/gserver/gradientmachines/NeuralNetwork.cpp
index dcf0acb5a2c..dbadc352a4c 100644
--- a/paddle/gserver/gradientmachines/NeuralNetwork.cpp
+++ b/paddle/gserver/gradientmachines/NeuralNetwork.cpp
@@ -21,6 +21,10 @@ limitations under the License. */
 #include "paddle/utils/Logging.h"
 #include "paddle/utils/Stat.h"
 
+#ifdef PADDLE_USE_MKLDNN
+#include "paddle/gserver/layers/MKLDNNLayer.h"
+#endif
+
 #ifndef PADDLE_MOBILE_INFERENCE
 #include "MultiNetwork.h"
 #include "RecurrentGradientMachine.h"
@@ -300,6 +304,17 @@ void NeuralNetwork::backward(const UpdateCallback& callback) {
   }
 }
 
+void NeuralNetwork::finish() {
+#ifdef PADDLE_USE_MKLDNN
+  FOR_EACH_R(layer, layers_) {
+    MKLDNNLayerPtr dnnLayer = std::dynamic_pointer_cast<MKLDNNLayer>(*layer);
+    if (dnnLayer) {
+      dnnLayer->convertWeightsToPaddle();
+    }
+  }
+#endif
+}
+
 Argument NeuralNetwork::getLayerOutput(const std::string& layerName) {
   return getLayer(layerName)->getOutput();
 }
diff --git a/paddle/gserver/gradientmachines/NeuralNetwork.h b/paddle/gserver/gradientmachines/NeuralNetwork.h
index 56a1ec78460..16971883b42 100644
--- a/paddle/gserver/gradientmachines/NeuralNetwork.h
+++ b/paddle/gserver/gradientmachines/NeuralNetwork.h
@@ -134,6 +134,9 @@ public:
 
   const std::string& getName() const { return subModelName_; }
 
+  /// some finish work, like convert the weight format of MKLDNNLayers
+  void finish() override;
+
 protected:
   /**
    * The constructor of NeuralNetwork.
-- 
GitLab