workspace.h

#ifndef CAFFE2_CORE_WORKSPACE_H_
#define CAFFE2_CORE_WORKSPACE_H_

#include "caffe2/core/common.h"

#ifndef CAFFE2_MOBILE
#error "mobile build state not defined"
#endif

#include <climits>
#include <cstddef>
#include <mutex>
#include <typeinfo>
#include <vector>

#include "caffe2/core/blob.h"
#include "caffe2/core/registry.h"
#include "caffe2/core/net.h"
#include "caffe2/proto/caffe2.pb.h"
#include "caffe2/utils/signal_handler.h"
#if CAFFE2_MOBILE
#include "caffe2/utils/threadpool/ThreadPool.h"
#endif // CAFFE2_MOBILE

CAFFE2_DECLARE_bool(caffe2_print_blob_sizes_at_exit);

namespace caffe2 {

class NetBase;
struct CompiledExecutionStep;

struct StopOnSignal {
  StopOnSignal()
      : handler_(std::make_shared<SignalHandler>(
            SignalHandler::Action::STOP,
            SignalHandler::Action::STOP)) {}

  StopOnSignal(const StopOnSignal& other) : handler_(other.handler_) {}

  bool operator()(int iter) {
    return handler_->CheckForSignals() != SignalHandler::Action::STOP;
  }

  std::shared_ptr<SignalHandler> handler_;
};


class Workspace {
 public:
  typedef std::function<bool(int)> ShouldContinue;
  typedef CaffeMap<string, unique_ptr<Blob> > BlobMap;
  typedef CaffeMap<string, unique_ptr<NetBase> > NetMap;
  Workspace() {
  }
  explicit Workspace(const string& root_folder)
      : root_folder_(root_folder) {}
  explicit Workspace(Workspace* const shared)
      : shared_(shared) {}
  Workspace(const string& root_folder, Workspace* shared)
      : root_folder_(root_folder), shared_(shared) {}
  ~Workspace() {
    if (FLAGS_caffe2_print_blob_sizes_at_exit) {
      PrintBlobSizes();
    }
  }

  void SetParentWorkspace(Workspace* shared) {
    shared_ = shared;
  }

  vector<string> LocalBlobs() const;

  vector<string> Blobs() const;

  const string& RootFolder() { return root_folder_; }
  inline bool HasBlob(const string& name) const {
    return (blob_map_.count(name) || (shared_ && shared_->HasBlob(name)));
  }

  void PrintBlobSizes();

  Blob* CreateBlob(const string& name);
  bool RemoveBlob(const string& name);
  const Blob* GetBlob(const string& name) const;
  Blob* GetBlob(const string& name);

  NetBase* CreateNet(const NetDef& net_def, bool overwrite = false);

  NetBase* GetNet(const string& net_name);
  void DeleteNet(const string& net_name);
  bool RunNet(const string& net_name);

  vector<string> Nets() const {
    vector<string> names;
    for (auto& entry : net_map_) {
      names.push_back(entry.first);
    }
    return names;
  }

  bool RunPlan(const PlanDef& plan_def,
               ShouldContinue should_continue = StopOnSignal{});

#if CAFFE2_MOBILE
  /*
   * Returns a CPU threadpool instace for parallel execution of
   * work. The threadpool is created lazily; if no operators use it,
   * then no threadpool will be created.
   */
  ThreadPool* GetThreadPool();
#endif

  // RunOperatorOnce and RunNetOnce runs an operator or net once. The difference
  // between RunNet and RunNetOnce lies in the fact that RunNet allows you to
  // have a persistent net object, while RunNetOnce creates a net and discards
  // it on the fly - this may make things like database read and random number
  // generators repeat the same thing over multiple calls.
  bool RunOperatorOnce(const OperatorDef& op_def);
  bool RunNetOnce(const NetDef& net_def);

 protected:
  bool ExecuteStepRecursive(CompiledExecutionStep& execution);

 private:
  BlobMap blob_map_;
  NetMap net_map_;
  string root_folder_ = ".";
  Workspace* shared_ = nullptr;
#if CAFFE2_MOBILE
  std::unique_ptr<ThreadPool> thread_pool_;
  std::mutex thread_pool_creation_mutex_;
#endif // CAFFE2_MOBILE

  DISABLE_COPY_AND_ASSIGN(Workspace);
};

}  // namespace caffe2

#endif  // CAFFE2_CORE_WORKSPACE_H_