worker.h

#ifndef INCLUDE_TRAINER_WORKER_H_
#define INCLUDE_TRAINER_WORKER_H_
#include <map>
#include <exception>
#include "neuralnet/neuralnet.h"
#include "proto/model.pb.h"
#include "utils/cluster.h"
#include "communication/socket.h"
#include "communication/msg.h"

namespace singa {
const int kCollectSleepTime=5;//milliseconds;
class Worker {
 public:
  Worker(int thread_id, int group_id, int worker_id);
  ~Worker(){}
  void Setup(const ModelProto& model, shared_ptr<NeuralNet> train_net);
  void set_test_net(shared_ptr<NeuralNet> test_net){
    test_net_=test_net;
  }
  void set_validation_net(shared_ptr<NeuralNet> val_net){
    validation_net_=val_net;
  }


  void Stop();
  int Put(shared_ptr<Param> param, int step);
  int Get(shared_ptr<Param> param, int step);
  int Update(shared_ptr<Param> param, int step);
  int Collect(shared_ptr<Param> param, int step);
  int CollectAll(shared_ptr<NeuralNet> net, int step);
  void RunOneBatch(int step, Metric* perf=nullptr);
  virtual void TrainOneBatch(int step)=0;
  virtual void TestOneBatch(shared_ptr<NeuralNet> net, int step, Phase phase)=0;
  void Test(shared_ptr<NeuralNet> net, int nsteps, const string &prefix);

  virtual void Run();


  const bool DisplayNow(const int step) const {
    return (modelproto_.display_frequency() > 0
        && step >= modelproto_.display_after_steps()
        && ((step - modelproto_.display_after_steps())
          % modelproto_.display_frequency() == 0));
  }

  const bool DisplayDebugInfo(const int step) const {
    return DisplayNow(step)&&modelproto_.debug()&&group_id_==0;
  }
  const void DisplayPerformance(const Metric & perf, const string& prefix);

  const bool StopNow(const int step) const{
    return (step >= modelproto_.train_steps());
  }
  const bool CheckpointNow(const int step) const{
    return (group_id_==0
        && modelproto_.checkpoint_frequency() > 0
        && step >= modelproto_.checkpoint_after_steps()
        && ((step - modelproto_.checkpoint_after_steps())
          % modelproto_.checkpoint_frequency() == 0));
  }
  const bool TestNow(const int step) const{
    return (group_id_==0
        && modelproto_.test_frequency() > 0
        && modelproto_.test_steps() > 0
        && step >= modelproto_.test_after_steps()
        && ((step - modelproto_.test_after_steps())
          % modelproto_.test_frequency() == 0));
  }
  const bool ValidateNow(const int step) {
    return (group_id_==0
        && modelproto_.validation_frequency() > 0
        && modelproto_.validation_steps() > 0
        && step >= modelproto_.validation_after_steps()
        && ((step - modelproto_.validation_after_steps())
          % modelproto_.validation_frequency() == 0));
  }


  void ReceiveBlobs(shared_ptr<NeuralNet> net);
  void SendBlob();
 protected:
  int thread_id_, group_id_, worker_id_;
  int step_;
  ModelProto modelproto_;
  shared_ptr<NeuralNet> train_net_, test_net_, validation_net_;
  shared_ptr<Dealer> layer_dealer_, param_dealer_;
  Poller layer_poller_, param_poller_;
};

class BPWorker: public Worker{
 public:
  ~BPWorker(){}
  BPWorker(int thread_id, int group_id, int worker_id):Worker(thread_id, group_id, worker_id){}
  virtual void TrainOneBatch(int step);
  virtual void TestOneBatch(shared_ptr<NeuralNet> net, int step, Phase phase);
  void Forward(shared_ptr<NeuralNet> net, int step, bool training);
  void Backward(shared_ptr<NeuralNet> net, int step);
};
}  // namespace singa

#endif  // INCLUDE_TRAINER_WORKER_H_