worker.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
* 
*   http://www.apache.org/licenses/LICENSE-2.0
* 
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*
*************************************************************/

#ifndef SINGA_TRAINER_WORKER_H_
#define SINGA_TRAINER_WORKER_H_
#include "neuralnet/neuralnet.h"
#include "proto/job.pb.h"
#include "communication/socket.h"

namespace singa {
const int kCollectSleepTime=5;
class Worker {
 public:
  static Worker* Create(const JobProto& proto);
  virtual void Init(int thread_id, int grp_id, int id);
  virtual ~Worker();
  void Setup(const JobProto& job, shared_ptr<NeuralNet> train_net,
      shared_ptr<NeuralNet> valid_net, shared_ptr<NeuralNet> test_net);
  void Run();
  void InitLocalParams();

  void Checkpoint(int step, shared_ptr<NeuralNet> net);
  void Test(int nsteps, Phase phase, shared_ptr<NeuralNet> net);
  virtual void TrainOneBatch(int step, Metric* perf)=0;
  virtual void TestOneBatch(int step, Phase phase, shared_ptr<NeuralNet> net,
      Metric* perf)=0;
  void Report(const string& prefix, const Metric & perf);

  int Put(Param* param, int step);
  int Get(Param* param, int step);
  int Update(Param* param, int step);
  int Collect(Param* param, int step);
  int CollectAll(shared_ptr<NeuralNet> net, int step);
  void ReceiveBlobs(
    bool data, bool grad, BridgeLayer* layer, shared_ptr<NeuralNet> net);
  void SendBlobs(
    bool data, bool grad, BridgeLayer* layer, shared_ptr<NeuralNet> net);

  inline bool DisplayNow(int step) const;
  inline bool DisplayDebugInfo(int step) const;
  inline bool StopNow(int step) const;
  inline bool CheckpointNow(int step) const;
  inline bool TestNow(int step) const;
  inline bool ValidateNow(int step) const;

  int grp_id() const { return grp_id_;}

  int id() const { return id_;}

 protected:
  int thread_id_, grp_id_, id_;
  int step_;
  JobProto job_conf_;
  shared_ptr<NeuralNet> train_net_, test_net_, validation_net_;
  Dealer* layer_dealer_, *dealer_;
};

class BPWorker: public Worker{
 public:
  ~BPWorker(){}
  void Init(int thread_id, int grp_id, int id) override;
  void TrainOneBatch(int step, Metric* perf) override;
  void TestOneBatch(int step, Phase phase, shared_ptr<NeuralNet> net,
      Metric* perf) override;

  void Forward(int step, Phase phase, shared_ptr<NeuralNet> net, Metric* perf);
  void Backward(int step, shared_ptr<NeuralNet> net);
};

class CDWorker: public Worker{
 public:
  void TrainOneBatch(int step, Metric* perf) override;
  void TestOneBatch(int step, Phase phase, shared_ptr<NeuralNet> net,
      Metric* perf) override;
};

inline int BlobTrgt(int grp, int layer) {
  return (grp << 16) | layer;
}

inline int BlobGrp(int blob_trgt) {
  return blob_trgt >> 16;
}

inline int BlobLayer(int blob_trgt) {
  static int mask = (1 << 16) -1;
  return blob_trgt & mask;
}
}  // namespace singa

#endif  // SINGA_TRAINER_WORKER_H_