diff --git a/src/mlpack/methods/ann/loss_functions/reconstruction_loss.hpp b/src/mlpack/methods/ann/loss_functions/reconstruction_loss.hpp index 56f6f39cc3..1f0b606521 100644 --- a/src/mlpack/methods/ann/loss_functions/reconstruction_loss.hpp +++ b/src/mlpack/methods/ann/loss_functions/reconstruction_loss.hpp @@ -39,8 +39,14 @@ class ReconstructionLoss public: /** * Create the ReconstructionLoss object. + * + * @param reduction Specifies the reduction to apply to the output. If false, + * 'mean' reduction is used, where sum of the output will be + * divided by the number of elements in the output. If true, + * 'sum' reduction is used and the output will be summed. It + * is set to true by default. */ - ReconstructionLoss(); + ReconstructionLoss(const bool reduction = true); /** * Computes the reconstruction loss. @@ -71,6 +77,11 @@ class ReconstructionLoss //! Modify the output parameter. OutputDataType& OutputParameter() { return outputParameter; } + //! Get the type of reduction used. + bool Reduction() const { return reduction; } + //! Modify the type of reduction used. + bool& Reduction() { return reduction; } + /** * Serialize the layer */ @@ -83,6 +94,10 @@ class ReconstructionLoss //! Locally-stored output parameter object. OutputDataType outputParameter; + + //! Boolean value that tells if reduction + // is 'sum' or 'mean'. + bool reduction; }; // class ReconstructionLoss } // namespace ann diff --git a/src/mlpack/methods/ann/loss_functions/reconstruction_loss_impl.hpp b/src/mlpack/methods/ann/loss_functions/reconstruction_loss_impl.hpp index ca5c986f05..ccb70b90f0 100644 --- a/src/mlpack/methods/ann/loss_functions/reconstruction_loss_impl.hpp +++ b/src/mlpack/methods/ann/loss_functions/reconstruction_loss_impl.hpp @@ -23,7 +23,7 @@ ReconstructionLoss< InputDataType, OutputDataType, DistType ->::ReconstructionLoss() +>::ReconstructionLoss(const bool reduction) : reduction(reduction) { // Nothing to do here. } @@ -35,7 +35,12 @@ ReconstructionLoss::Forward( const PredictionType& prediction, const TargetType& target) { dist = DistType(prediction); - return -dist.LogProbability(target); + typename PredictionType::elem_type lossSum = -dist.LogProbability(target); + + if (reduction) + return lossSum; + + return lossSum / target.n_elem; } template @@ -47,15 +52,18 @@ void ReconstructionLoss::Backward( { dist.LogProbBackward(target, loss); loss *= -1; + + if (!reduction) + loss = loss / target.n_elem; } template template void ReconstructionLoss::serialize( - Archive& /* ar */, + Archive& ar, const uint32_t /* version */) { - // Nothing to do here. + ar(CEREAL_NVP(reduction)); } } // namespace ann