From ebd5b1a43e7deb4ae01d0ca3453c40b54afa02db Mon Sep 17 00:00:00 2001 From: walragatver Date: Tue, 11 Jun 2019 23:35:59 +0530 Subject: [PATCH] Add implementation of weight_norm layer. --- src/mlpack/methods/ann/layer/weight_norm.hpp | 30 +++- .../methods/ann/layer/weight_norm_impl.hpp | 134 ++++++++++++++++-- 2 files changed, 144 insertions(+), 20 deletions(-) diff --git a/src/mlpack/methods/ann/layer/weight_norm.hpp b/src/mlpack/methods/ann/layer/weight_norm.hpp index eeb7ba616b..eb9a2fcf43 100644 --- a/src/mlpack/methods/ann/layer/weight_norm.hpp +++ b/src/mlpack/methods/ann/layer/weight_norm.hpp @@ -14,9 +14,13 @@ #include #include "layer_types.hpp" + #include "../visitor/delete_visitor.hpp" #include "../visitor/delta_visitor.hpp" #include "../visitor/output_parameter_visitor.hpp" +#include "../visitor/reset_visitor.hpp" +#include "../visitor/weight_size_visitor.hpp" +#include "../visitor/weight_set_visitor.hpp" namespace mlpack { namespace ann /** Artificial Neural Network. */ { @@ -77,7 +81,7 @@ class WeightNorm * @param output Resulting output activations. */ template - void Forward(const arma::Mat&& input, arma::Mat&& output); + void Forward(arma::Mat&& input, arma::Mat&& output); /** * Backward pass through the layer. This function will call the @@ -101,7 +105,7 @@ class WeightNorm * @param gradient The calculated gradient. */ template - void Gradient(const arma::Mat&& input, + void Gradient(arma::Mat&& input, arma::Mat&& error, arma::Mat&& gradient); @@ -115,10 +119,10 @@ class WeightNorm //! Modify the gradient. OutputDataType& Gradient() { return gradient; } - //! Get the input parameter. + /*//! Get the input parameter. InputDataType const& InputParameter() const { return inputParameter; } //! Modify the input parameter. - InputDataType& InputParameter() { return inputParameter; } + InputDataType& InputParameter() { return inputParameter; }*/ //! Return the model modules. std::vector >& Model() @@ -147,14 +151,14 @@ class WeightNorm * @param args The layer parameter. */ template - void Add(Args... args) { network.push_back(new LayerType(args...)); } + void Add(Args... args); /* * Add a new module to the model. * * @param layer The Layer to be added to the model. */ - void Add(LayerTypes layer) { network.push_back(layer); } + void Add(LayerTypes layer); /** * Serialize the layer. @@ -187,20 +191,32 @@ class WeightNorm //! Locally-stored network modules. std::vector > network; + //! Locally stored number of elements in the weights of wrapped layer. + size_t networkWeightSize; + //! Locally-stored output parameter object. OutputDataType outputParameter; //! Locally-stored output parameter visitor module object. OutputParameterVisitor outputParameterVisitor; + // Reset the gradient for all modules that implement the Gradient function. + void ResetGradients(arma::mat& gradient); + + //! Locally-stored reset visitor. + ResetVisitor resetVisitor; + //! Locally-stored scalar parameter. - size_t scalarParameter; + OutputDataType scalarParameter; //! Locally-stored parameter vector. OutputDataType vectorParameter; //! Locally-stored parameters. OutputDataType weights; + + //! Locally-stored weight size visitor. + WeightSizeVisitor weightSizeVisitor; }; // class BatchNorm } // namespace ann diff --git a/src/mlpack/methods/ann/layer/weight_norm_impl.hpp b/src/mlpack/methods/ann/layer/weight_norm_impl.hpp index 7b7f5a24a1..9b9c1f4d10 100644 --- a/src/mlpack/methods/ann/layer/weight_norm_impl.hpp +++ b/src/mlpack/methods/ann/layer/weight_norm_impl.hpp @@ -16,14 +16,17 @@ // In case it is not included. #include "weight_norm.hpp" +#include "../visitor/forward_visitor.hpp" +#include "../visitor/backward_visitor.hpp" +#include "../visitor/gradient_visitor.hpp" + namespace mlpack { namespace ann { /** Artificial Neural Network. */ template WeightNorm::WeightNorm() : - model(false), - scalarParameter(1) + model(false) { // Nothing to do here. } @@ -40,38 +43,129 @@ template void WeightNorm::Reset() { - // It will call the reset function of weight norm layer. + size_t offset = boost::apply_visitor(WeightSetVisitor(std::move(weights), + 0), network[0]); + boost::apply_visitor(resetVisitor, network[0]); + + vectorParameter = arma::mat(weights.memptr() + offset, offset, 1, false, + false); + + scalarParameter = arma::mat(weights.memptr() + 2 * offset, 1, 1, false, + false); } template template void WeightNorm::Forward( - const arma::Mat&& input, arma::Mat&& output) + arma::Mat&& input, arma::Mat&& output) { - // It will call the Forward function of the wrapped layer. + weights.rows(0, networkWeightSize) = scalarParameter * vectorParameter / + std::sqrt(arma::accu(arma::square(vectorParameter))); + + boost::apply_visitor(ForwardVisitor(std::move(input), std::move( + boost::apply_visitor(outputParameterVisitor, network[0]))), + network[0]); + + output = boost::apply_visitor(outputParameterVisitor, network[0]); } template template void WeightNorm::Backward( - const arma::Mat&& input, arma::Mat&& gy, arma::Mat&& g) + const arma::Mat&& /* input */, arma::Mat&& gy, arma::Mat&& g) { - // It will directly call the Backward function of the wrapped layer. + boost::apply_visitor(BackwardVisitor(std::move(boost::apply_visitor( + outputParameterVisitor, network[0])), std::move(gy), std::move( + boost::apply_visitor(deltaVisitor, network[0]))), network[0]); + + g = boost::apply_visitor(deltaVisitor, network[0]); } template template void WeightNorm::Gradient( - const arma::Mat&& /* input */, + arma::Mat&& input, arma::Mat&& error, arma::Mat&& gradient) { - // First it will calculate the gradients of the wrapped layer. - // Then it will calculate gradients of the vector parameter v and scalar - // parameter g. + if (!model) + { + ResetGradients(gradient); + } + + // Calculate the gradients of the wrapped layer. + boost::apply_visitor(GradientVisitor(std::move(input), + std::move(error)), network[0]); + + // Store the norm of vectorParameter temporarily. + size_t normVectorParameter = std::sqrt(arma::accu(arma::square( + vectorParameter))); + + // Calculate gradients of the scalar parameter. + gradient[gradient.n_rows - 1] = arma::accu(gradient.rows(0, networkWeightSize) + % vectorParameter) / normVectorParameter; + + // Calculate gradients of the vector parameter. + gradient.rows(networkWeightSize, 2 * networkWeightSize) = scalarParameter / + normVectorParameter * (gradient.rows(0, networkWeightSize) - + gradient[gradient.n_rows - 1] / normVectorParameter * vectorParameter); +} + +template +template +void WeightNorm::Add( + Args... args) +{ + // Only one layer will be wrapped. + if(network.size() > 0) + { + std::for_each(network.begin(), network.end(), + boost::apply_visitor(deleteVisitor)); + network.clear(); + } + + network.push_back(new LayerType(args...)); + + // Now set the weights of the weight norm layer. + networkWeightSize = boost::apply_visitor(weightSizeVisitor, network[0]); + weights.set_size(2 * networkWeightSize + 1, 1); + + Reset(); +} + +template +void WeightNorm::Add( + LayerTypes layer) +{ + // Only one layer will be wrapped. + if(network.size() > 0) + { + std::for_each(network.begin(), network.end(), + boost::apply_visitor(deleteVisitor)); + network.clear(); + } + + network.push_back(layer); + + // Now set the weights of the weight norm layer. + networkWeightSize = boost::apply_visitor(weightSizeVisitor, network[0]); + weights.set_size(2 * networkWeightSize + 1, 1); + + Reset(); +} + +template +void WeightNorm::ResetGradients( + arma::mat& gradient) +{ + boost::apply_visitor(GradientSetVisitor(std::move(gradient), 0), + network[0]); } template void WeightNorm::serialize( Archive& ar, const unsigned int /* version */) { + if (Archive::is_loading::value) + { + std::for_each(network.begin(), network.end(), + boost::apply_visitor(deleteVisitor)); + network.clear(); + } + ar & BOOST_SERIALIZATION_NVP(network); ar & BOOST_SERIALIZATION_NVP(model); - ar & BOOST_SERIALIZATION_NVP(scalarParameter); - ar & BOOST_SERIALIZATION_NVP(vectorParameter); + ar & BOOST_SERIALIZATION_NVP(networkWeightSize); + + // If we are loading, we need to initialize the weights. + if (Archive::is_loading::value) + { + // The behavior in earlier versions was to always assume the weights needed + // to be reset. + weights.set_size(2 * networkWeightSize + 1, 1); + } } } // namespace ann