loss.h

#ifndef SINGA_MODEL_LOSS_H_
#define SINGA_MODEL_LOSS_H_
#include <stack>
#include "singa/proto/model.pb.h"
#include "singa/core/tensor.h"
namespace singa {

// template <typename T = Tensor>
class Loss {
public:
  Loss() = default;
  void Setup(const string &conf) {
    LossConf loss;
    loss.ParseFromString(conf);
    Setup(loss);
  }
  virtual ~Loss() {};
  virtual void ToDevice(std::shared_ptr<Device> device) {}
  virtual void Setup(const LossConf &conf) {}

  virtual Tensor Forward(int flag, const Tensor &prediction,
                         const Tensor &target) = 0;

  float Evaluate(int flag, const Tensor &prediction, const Tensor &target) {
    Tensor loss = Forward(flag, prediction, target);
    return Sum<float>(loss) / (1.0f * loss.Size());
  }

  virtual Tensor Backward() = 0;
};

// ============= Mean Squared Error ===========================================
class MSE : public Loss {
 public:
  Tensor Forward(int flag, const Tensor& prediction,
      const Tensor& target) override;

  Tensor Backward() override;

 private:
  // to buffer intermediate data, i.e., prediction-target
  std::stack<Tensor> buf_;
};


// ===============Softamx Cross Entropy =======================================
class SoftmaxCrossEntropy : public Loss {
 public:
  Tensor Forward(int flag, const Tensor& prediction,
      const Tensor& target) override;

  Tensor Backward() override;

 private:
  // to buffer intermediate data, i.e., probability for each category and
  // the target (ground truth)
  std::stack<Tensor> buf_;
};

}  // namespace singa

#endif  // SINGA_MODEL_LOSS_H_