diff --git a/paddle/gserver/gradientmachines/NeuralNetwork.cpp b/paddle/gserver/gradientmachines/NeuralNetwork.cpp index dcf0acb5a2cc7698625a2e254d4c12a32bc9631d..dbadc352a4ccd7483bf67e1025c212f514e32a24 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(*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 56a1ec78460731554c9b47cf3f517f7654dc314f..16971883b42786e2cb48faafd2a7e95d45075ac8 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.