base_layer.h

#ifndef INCLUDE_BASE_LAYER_H_
#define INCLUDE_BASE_LAYER_H_

#include <vector>
#include <string>
#include <map>
#include <functional>
#include <utility>
#include <memory>
#include <chrono>
#include <algorithm>
#include <thread>

#include "proto/model.pb.h"
#include "utils/param.h"
#include "utils/common.h"
#include "utils/blob.h"

using std::vector;
using std::shared_ptr;
using std::make_shared;
using std::string;
using std::map;

namespace singa{

class Layer;
typedef shared_ptr<Layer> SLayer;
class Layer {
 public:
  Layer(){}
  virtual ~Layer(){}
  virtual void Init(const LayerProto &proto);
  void Init(const Layer& other, const vector<int>& shape);
  virtual void ToProto(LayerProto *layer_proto, bool copyData);
  virtual void Setup(const LayerProto& proto,
      const vector<SLayer>& srclayers)=0;
  virtual void Setup();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers)=0;
  virtual void SetupAfterPartition();
  virtual vector<shared_ptr<Param>> GetParams(){
    return vector<shared_ptr<Param>>();
  }
  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers)=0;
  virtual void ComputeFeature(bool training);
  virtual void ComputeGradient(const vector<SLayer>& srclayers)=0;
  virtual void ComputeGradient();
  virtual int partition_dimension() const {
    int ret=0;
    if(partition_type()==kLayerPartition)
      ret= 1;
    else if(partition_type()==kNone)
      ret= -1;
    return ret;
  }

  virtual ConnectionType connection_type(int k) const {
    CHECK_LT(k, srclayers_.size());
    return kOneToOne;
  }
  virtual PartitionType partition_type() const {
    return layer_proto_.partition_type();
  }
  virtual void set_partitionid(int id){
    layer_proto_.set_partitionid(id);
  }
  virtual int partitionid() const {
    return layer_proto_.partitionid();
  }
  virtual void set_name(string name){
    name_=name;
    layer_proto_.set_name(name);
  }
  virtual const string type() const {
    return layer_proto_.type();
  }
  const std::string &name() const {
    return layer_proto_.name();
  }
  const std::string &datablob() const {
    return layer_proto_.datablob();
  }
  const vector<int>& shape(const Layer* layer) const{
    return data(layer).shape();
  }

  virtual const Blob<float>& data(const Layer* from) const {
    return data_;
  }
  virtual Blob<float>* mutable_data(const Layer* from){
    return &data_;
  }

  virtual const Blob<float>& grad(const Layer* from) const {
    return grad_;
  }
  virtual Blob<float>* mutable_grad(const Layer* from) {
    return &grad_;
  }

  virtual const vector< SLayer> srclayers() const {
    return srclayers_;
  }
  virtual const vector<SLayer> dstlayers() const {
    return dstlayers_;
  }

  virtual const int srclayers_size() const {
    return srclayers_.size();
  }
  virtual const int dstlayers_size() const {
    return dstlayers_.size();
  }
  virtual void ClearDstLayers() {
    dstlayers_.clear();
  }
  virtual void ClearSrcLayers() {
    srclayers_.clear();
  }

  virtual void AddSrcLayer(SLayer src){
    srclayers_.push_back(src);
  }
  virtual void AddDstLayer(SLayer dst){
    dstlayers_.push_back(dst);
  }

  virtual bool is_datalayer() const {
    return false;
  }
  virtual bool is_parserlayer() const {
    return false;
  }
  virtual bool is_losslayer() const {
    return false;
  }
  virtual bool is_bridgesrclayer() const {
    return false;
  }
  virtual bool is_bridgedstlayer() const {
    return false;
  }
protected:
  string name_;
  Blob<float> data_, grad_;
  LayerProto layer_proto_;
  vector<SLayer> srclayers_, dstlayers_;
};

class BridgeSrcLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void SetupAfterPartition();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers);
  virtual void ComputeGradient(const vector<SLayer>& srclayers);
  virtual const Blob<float>& data(const Layer* from) const {
    return srclayers_[0]->data(this);
  }
  virtual Blob<float>* mutable_data(const Layer* from){
    return srclayers_[0]->mutable_data(this);
  }

  virtual const Blob<float>& grad(const Layer* from) const {
    return srclayers_[0]->grad(this);
  }
  virtual Blob<float>* mutable_grad(const Layer* from) {
    return srclayers_[0]->mutable_grad(this);
  }
  int dst_partition() const;
  virtual bool is_bridgesrclayer() const {
    return true;
  }
  virtual void set_ready(bool a) {
    ready_=a;
  }
  virtual bool ready() const {
    return ready_;
  }
 protected:
  bool ready_;
};
class BridgeDstLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void SetupAfterPartition();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers){
    ready_=false;
  }
  virtual void ComputeGradient(const vector<SLayer>& srclayers){}
  virtual bool is_bridgedstlayer() const {
    return true;
  }
  virtual void set_ready(bool a) {
    ready_=a;
  }
  virtual bool ready() const {
    return ready_;
  }
 protected:
  bool ready_;
};

class ConcateLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void SetupAfterPartition();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual void ComputeFeature(bool training, const vector<shared_ptr<Layer>>& srclayers);
  virtual void ComputeGradient(const vector<shared_ptr<Layer>>& srclayers);
};


class DataLayer: public Layer{
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers)=0;
  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers)=0;
  virtual bool is_datalayer() const {
    return true;
  }
  virtual void ComputeGradient(const vector<SLayer>& srclayers){};
  virtual const vector<Record>& records() const {
    return records_;
  }
  virtual void Setup(){
    vector<SLayer> dummy;
    Setup(layer_proto_,dummy);
    has_setup_=true;
  }
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual void SetupAfterPartition(){
    if(!has_setup_)
    Setup();
  }
  virtual PartitionType partition_type () const {
    return kNone;
  }

  virtual int batchsize() const {
    return layer_proto_.data_param().batchsize();
  }
  virtual const Record& sample() const {
    return sample_;
  }

  virtual Blob<float>* mutable_data(const Layer* layer) {
    return nullptr;
  }
  virtual Blob<float>* mutable_grad(const Layer* layer) {
    return nullptr;
  }
 protected:
  bool has_setup_;
  int random_skip_, batchsize_;
  Record sample_;
  vector<Record> records_;
};

class PrefetchLayer : public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;
  using Layer::SetupAfterPartition;

  virtual ~PrefetchLayer();
  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers);
  virtual void ComputeGradient(const vector<SLayer>& srclayers){};
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual const Blob<float>& data(const Layer* from) const ;
  virtual Blob<float>* mutable_data(const Layer* layer) ;

  virtual Blob<float>* mutable_grad(const Layer* layer){
    return nullptr;
  }
  virtual const Blob<float>& grad(const Layer* from) const {
    CHECK(false)<<"Loss layer has not gradient blob";
    return grad_;
  }
  virtual PartitionType partition_type () const {
    return kNone;
  }

  void Prefetch(bool training);
 protected:
  vector<shared_ptr<Layer>> sublayers_;
  map<string, Blob<float>> datablobs_;
  std::thread thread_;
};

class SliceLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void ComputeFeature(bool training, const vector<shared_ptr<Layer>>& srclayers);
  virtual void ComputeGradient(const vector<shared_ptr<Layer>>& srclayers);
  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void SetupAfterPartition();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual const Blob<float>& data(const Layer* layer) const;
  virtual const Blob<float>& grad(const Layer* layer) const;
  virtual Blob<float>* mutable_data(const Layer* layer);
  virtual Blob<float>* mutable_grad(const Layer* layer);
 protected:
  int SliceID(const Layer* layer) const;
  vector<Blob<float>> datavec_, gradvec_;
  int slice_dim_, slice_num_;
};

class SplitLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void Setup(const LayerProto& proto, const vector<SLayer>& srclayers);
  virtual void SetupAfterPartition();
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual void ComputeFeature(bool training, const vector<shared_ptr<Layer>>& srclayers);
  virtual void ComputeGradient(const vector<shared_ptr<Layer>>& srclayers);
};

class LossLayer: public Layer{
 public:
  using Layer::Setup;
  using Layer::SetupAfterPartition;

  virtual void Setup(const LayerProto& proto,
      const vector<SLayer>& srclayers)=0;
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers)=0;

  virtual Blob<float>* mutable_grad(const Layer* layer){
    return nullptr;
  }
  virtual const Blob<float>& grad(const Layer* from) const {
    CHECK(false)<<"Loss layer has not gradient blob";
    return grad_;
  }
  virtual bool is_losslayer() const {
    return true;
  }
  virtual const Blob<float>& metric() const {
    return metric_;
  }
 protected:
  Blob<float> metric_;
};

class ParserLayer: public Layer {
 public:
  using Layer::Setup;
  using Layer::SetupAfterPartition;
  using Layer::ComputeFeature;
  using Layer::ComputeGradient;

  virtual void Setup(const LayerProto& proto,
      const vector<SLayer>& srclayers)=0;
  virtual void ParseRecords(bool training, const vector<Record>& records,
      Blob<float>* blob)=0;

  virtual bool is_parserlayer() const {
    return true;
  }

  virtual void ComputeFeature(bool training, const vector<SLayer>& srclayers);
  virtual void ComputeGradient(const vector<SLayer>& srclayers){};
  virtual void Setup(){
    Setup(layer_proto_,srclayers_);
    has_setup_=true;
  }
  virtual void SetupAfterPartition(){
    if(!has_setup_)
      Setup();
  }
  virtual void SetupAfterPartition(const LayerProto& proto,
      const vector<int> &shape,
      const vector<SLayer>& srclayers){}

  virtual PartitionType partition_type () const{
    return kNone;
  }
  virtual Blob<float>* mutable_grad(const Layer* layer) {
    return nullptr;
  }
  virtual const Blob<float>& grad(const Layer* from) const {
    CHECK(false)<<"Parser layer has not gradient blob";
    return grad_;
  }

 private:
  bool has_setup_;
};
} // singa

#endif // INCLUDE_BASE_LAYER_H_