Add implementation of weight_norm layer.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user