trainer.h

#ifndef INCLUDE_TRAINER_TRAINER_H_
#define INCLUDE_TRAINER_TRAINER_H_
#include "proto/cluster.pb.h"
#include "proto/model.pb.h"
#include "utils/updater.h"
#include "utils/param.h"
#include "utils/singleton.h"
#include "utils/factory.h"
#include "neuralnet/neuralnet.h"
#include "trainer/worker.h"
#include "trainer/server.h"
#include "communication/socket.h"

namespace singa {
class Trainer{
  public:
  class ParamInfo{
   public:
    ParamInfo(shared_ptr<Param> p,int local, int owner):
      num_update(0), next_version(0),num_local(local), num_total(1),
      owner_procs(owner){
        shares.push_back(p);
      }

    void AddParam(shared_ptr<Param> p, int local, int owner){
      num_local+=local;
      num_total+=1;
      if(owner>-1)
        owner_procs=owner;
      if(local>0){
        shares.push_back(p);
      }
    }
    int num_update, next_version; 

    int num_local; 
    int num_total; 
    int owner_procs; 
    vector<shared_ptr<Param>> shares;
  };

 typedef std::map<int, shared_ptr<ParamInfo>> ParamShard;

 public:
  void Start(const ModelProto& modelproto, const ClusterProto& clusterproto,
    int procs_id);

  // TODO add Resume() function to continue training from a previously stopped
  // point.

 protected:
  void Run(int nworkers, int nservers,
      const std::map<int, shared_ptr<ParamShard>>& shards);
  void RegisterDefaultClasses(const singa::ModelProto& proto);

  virtual int Sharding(int param_id);

    virtual Msg* HandleGet(shared_ptr<ParamInfo>counter, Msg** msg);
    virtual Msg* HandleGetResponse(shared_ptr<ParamInfo>counter, Msg** msg);

    virtual Msg* HandleUpdate(shared_ptr<ParamInfo>counter, Msg** msg);
  virtual int HandleUpdateResponse(shared_ptr<ParamInfo>counter, Msg** msg);

    virtual Msg* HandlePut(shared_ptr<ParamInfo>counter, Msg** msg);
    virtual Msg* HandleConnect(Msg** msg);

 protected:
  int procs_id_;
  shared_ptr<Router> router_;
};
} /* singa */
#endif // INCLUDE_TRAINER_TRAINER_H_