Add implementation of weight_norm layer.

This commit is contained in:
walragatver
2019-08-23 00:14:27 +05:30
parent bf14b992c3
commit ebd5b1a43e
2 changed files with 144 additions and 20 deletions
+23 -7
View File
@@ -14,9 +14,13 @@
#include <mlpack/prereqs.hpp>
#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<typename eT>
void Forward(const arma::Mat<eT>&& input, arma::Mat<eT>&& output);
void Forward(arma::Mat<eT>&& input, arma::Mat<eT>&& output);
/**
* Backward pass through the layer. This function will call the
@@ -101,7 +105,7 @@ class WeightNorm
* @param gradient The calculated gradient.
*/
template<typename eT>
void Gradient(const arma::Mat<eT>&& input,
void Gradient(arma::Mat<eT>&& input,
arma::Mat<eT>&& error,
arma::Mat<eT>&& 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<LayerTypes<CustomLayers...> >& Model()
@@ -147,14 +151,14 @@ class WeightNorm
* @param args The layer parameter.
*/
template <class LayerType, class... Args>
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<CustomLayers...> layer) { network.push_back(layer); }
void Add(LayerTypes<CustomLayers...> layer);
/**
* Serialize the layer.
@@ -187,20 +191,32 @@ class WeightNorm
//! Locally-stored network modules.
std::vector<LayerTypes<CustomLayers...> > 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
+121 -13
View File
@@ -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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
WeightNorm<InputDataType, OutputDataType, CustomLayers...>::WeightNorm() :
model(false),
scalarParameter(1)
model(false)
{
// Nothing to do here.
}
@@ -40,38 +43,129 @@ template<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
template<typename eT>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::Forward(
const arma::Mat<eT>&& input, arma::Mat<eT>&& output)
arma::Mat<eT>&& input, arma::Mat<eT>&& 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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
template<typename eT>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::Backward(
const arma::Mat<eT>&& input, arma::Mat<eT>&& gy, arma::Mat<eT>&& g)
const arma::Mat<eT>&& /* input */, arma::Mat<eT>&& gy, arma::Mat<eT>&& 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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
template<typename eT>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::Gradient(
const arma::Mat<eT>&& /* input */,
arma::Mat<eT>&& input,
arma::Mat<eT>&& error,
arma::Mat<eT>&& 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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
template <class LayerType, class... Args>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::Add(
LayerTypes<CustomLayers...> 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<typename InputDataType, typename OutputDataType,
typename... CustomLayers>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::ResetGradients(
arma::mat& gradient)
{
boost::apply_visitor(GradientSetVisitor(std::move(gradient), 0),
network[0]);
}
template<typename InputDataType, typename OutputDataType,
@@ -80,10 +174,24 @@ template<typename Archive>
void WeightNorm<InputDataType, OutputDataType, CustomLayers...>::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