Caffe2 - C++ API
A deep learning, cross platform ML framework
blobs_queue.h
1 
17 #pragma once
18 
19 #include <atomic>
20 #include <condition_variable>
21 #include <memory>
22 #include <mutex>
23 #include <queue>
24 
25 #include "caffe2/core/blob_stats.h"
26 #include "caffe2/core/logging.h"
27 #include "caffe2/core/stats.h"
28 #include "caffe2/core/tensor.h"
29 #include "caffe2/core/workspace.h"
30 
31 namespace caffe2 {
32 
33 // A thread-safe, bounded, blocking queue.
34 // Modelled as a circular buffer.
35 
36 // Containing blobs are owned by the workspace.
37 // On read, we swap out the underlying data for the blob passed in for blobs
38 
39 class BlobsQueue : public std::enable_shared_from_this<BlobsQueue> {
40  public:
41  BlobsQueue(
42  Workspace* ws,
43  const std::string& queueName,
44  size_t capacity,
45  size_t numBlobs,
46  bool enforceUniqueName,
47  const std::vector<std::string>& fieldNames = {});
48 
49  ~BlobsQueue() {
50  close();
51  }
52 
53  bool blockingRead(
54  const std::vector<Blob*>& inputs,
55  float timeout_secs = 0.0f);
56  bool tryWrite(const std::vector<Blob*>& inputs);
57  bool blockingWrite(const std::vector<Blob*>& inputs);
58  void close();
59  size_t getNumBlobs() const {
60  return numBlobs_;
61  }
62 
63  private:
64  bool canWrite();
65  void doWrite(const std::vector<Blob*>& inputs);
66 
67  std::atomic<bool> closing_{false};
68 
69  size_t numBlobs_;
70  std::mutex mutex_; // protects all variables in the class.
71  std::condition_variable cv_;
72  int64_t reader_{0};
73  int64_t writer_{0};
74  std::vector<std::vector<Blob*>> queue_;
75  const std::string name_;
76 
77  struct QueueStats {
78  CAFFE_STAT_CTOR(QueueStats);
79  CAFFE_EXPORTED_STAT(queue_balance);
80  CAFFE_EXPORTED_STAT(queue_dequeued_records);
81  CAFFE_DETAILED_EXPORTED_STAT(queue_dequeued_bytes);
82  } stats_;
83 };
84 } // namespace caffe2
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.