weight_norm.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_ANN_LAYER_WEIGHTNORM_HPP
#define MLPACK_METHODS_ANN_LAYER_WEIGHTNORM_HPP

#include <mlpack/prereqs.hpp>
#include "layer.hpp"

namespace mlpack {
namespace ann  {

template <
  typename InputType = arma::mat,
  typename OutputType = arma::mat
>
class WeightNormType : public Layer<InputType, OutputType>
{
 public:
  WeightNormType();

  WeightNormType(Layer<InputType, OutputType>* layer);

  ~WeightNormType();

  WeightNormType(const WeightNormType& other);
  WeightNormType(WeightNormType&& other);
  WeightNormType& operator=(const WeightNormType& other);
  WeightNormType& operator=(WeightNormType&& other);

  WeightNormType* Clone() const { return new WeightNormType(*this); }

  void SetWeights(typename OutputType::elem_type* weightsPtr);

  void Forward(const InputType& input, OutputType& output);

  void Backward(const InputType& input,
                const OutputType& gy,
                OutputType& g);

  void Gradient(const InputType& input,
                const OutputType& error,
                OutputType& gradient);

  OutputType const& Parameters() const { return weights; }
  OutputType& Parameters() { return weights; }

  Layer<InputType, OutputType>* const& WrappedLayer() { return wrappedLayer; }

  const size_t WeightSize() const { return wrappedLayer->WeightSize(); }

  const std::vector<size_t> OutputDimensions() const
  {
    wrappedLayer->InputDimensions() = inputDimensions;
    return wrappedLayer->OutputDimensions();
  }

  template<typename Archive>
  void serialize(Archive& ar, const uint32_t /* version */);

 private:
  size_t biasWeightSize;

  Layer<InputType, OutputType>* wrappedLayer;

  size_t layerWeightSize;

  void ResetGradients(OutputType& gradient);

  OutputType scalarParameter;

  OutputType vectorParameter;

  OutputType weights;

  OutputType layerGradients;

  OutputType layerWeights;
}; // class WeightNormType.

// Standard WeightNorm layer.
typedef WeightNormType<arma::mat, arma::mat> WeightNorm;

} // namespace ann
} // namespace mlpack

// Include the implementation.
#include "weight_norm_impl.hpp"

#endif