neuralnet.h

#ifndef INCLUDE_NET_NET_H_
#define INCLUDE_NET_NET_H_

#include <glog/logging.h>
#include <vector>
#include <map>
#include <memory>

#include "proto/model.pb.h"
#include "neuralnet/layer.h"
#include "utils/factory.h"
#include "utils/graph.h"

using std::vector;
using std::string;
using std::map;
using std::shared_ptr;
namespace singa {
class NeuralNet {
 public:
  static void RegisterLayers();
  static shared_ptr<NeuralNet> SetupNeuralNet(const NetProto& np, Phase phase,
      int group_size);

 public:
  NeuralNet(NetProto net_proto, int group_size=1);
  std::string ToString();
  string DebugInfo();

  std::string ToAdjacency();
  void AddLayer(const LayerProto &layer_proto){};
  void AddLayer(const Layer* layer){};
  void ShareParams(shared_ptr<NeuralNet> other,int flag);
  void ToProto(NetProto *net_proto, bool copyData=false);
  const std::vector<shared_ptr<Layer>>& layers() {
    return layers_;
  }
  const std::vector<ParserLayer*>& parserlayers() {
    if(parserlayers_.size()==0){
      for(auto& layer: layers_)
        if(layer->is_parserlayer())
          parserlayers_.push_back(static_cast<ParserLayer*>(layer.get()));
    }
    return parserlayers_;
  }
  const std::vector<LossLayer*>& losslayers() {
    if(losslayers_.size()==0){
      for(auto& layer: layers_)
        if(layer->is_losslayer())
          losslayers_.push_back(static_cast<LossLayer*>(layer.get()));
    }
    return losslayers_;
  }
  const std::vector<DataLayer*>& datalayers() {
    if(datalayers_.size()==0){
      for(auto& layer: layers_)
        if(layer->is_datalayer())
          datalayers_.push_back(static_cast<DataLayer*>(layer.get()));
    }
    return datalayers_;
  }
  const std::vector<shared_ptr<Param>> &params()const {
    return params_;
  }
  shared_ptr<Layer> name2layer(string name){
    if (name2layer_.find(name)!=name2layer_.end())
      return name2layer_[name];
    else return nullptr;
  }

  shared_ptr<Param> paramid2param(int id) {
    if(paramid2param_.size()==0){
      for(auto& layer: layers_){
        for(shared_ptr<Param> p: layer->GetParams()){
          paramid2param_[p->id()]=p;
        }
      }
    }
    return paramid2param_[id];
  }

 protected:
  void ConstructNeuralNet(const NetProto &net_proto);
  void PartitionNeuralNet();
  map<string, shared_ptr<Layer>> GetNameToLayer(
    const vector<shared_ptr<Layer>>& layers);
  Graph CreatePartitonedGraph(const vector<shared_ptr<Layer>>& layers,
    const map<string, shared_ptr<Layer>>& name2layer);

  map<string, vector<shared_ptr<Layer>>> PartitionLayers(
      const vector<shared_ptr<Layer>>& layers);

 protected:
  vector<shared_ptr<Layer>> layers_;
  vector<ParserLayer*> parserlayers_;
  vector<LossLayer*> losslayers_;
  vector<DataLayer*> datalayers_;
  vector<shared_ptr<Param>> params_;
  map<string, shared_ptr<Layer>> name2layer_;
  map<int, shared_ptr<Param>> paramid2param_;

  map<string, LayerProto> name2layerproto_;
  int group_size_;
  Graph graph_;
};
}  // namespace singa
#endif  // INCLUDE_NET_NET_H_