Caffe2 - C++ API
A deep learning, cross platform ML framework
rebatching_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/logging.h"
26 #include "caffe2/core/operator.h"
27 #include "caffe2/core/stats.h"
28 #include "caffe2/core/tensor.h"
29 
30 namespace caffe2 {
31 
32 // TODO: This is a very naive implementation with a single mutex. We can do the
33 // atomic index + circular queue optimizations or pull something more
34 // heavy-weight later
35 
37  public:
38  RebatchingQueue(size_t capacity, size_t numBlobs);
39 
40  ~RebatchingQueue();
41 
42  bool enqueueOne(
43  CPUContext& context,
44  const std::vector<const TensorCPU*>& inputs);
45 
46  bool enqueueMany(
47  CPUContext& context,
48  const std::vector<const TensorCPU*>& inputs);
49 
50  bool dequeue(
51  CPUContext& context,
52  size_t numElements,
53  const std::vector<TensorCPU*>& outputs);
54 
55  size_t capacity() const;
56 
57  size_t numBlobs() const;
58 
59  bool isClosed() const;
60 
61  void close();
62 
63  private:
64  bool enqueue(std::vector<std::vector<TensorCPU>> splittedInputs);
65 
66  bool canWrite() const;
67  bool canRead() const;
68 
69  const size_t capacity_;
70  const size_t numBlobs_;
71 
72  mutable std::mutex mutex_;
73 
74  bool isClosed_{false};
75 
76  uint64_t head_{0};
77  uint64_t tail_{0};
78 
79  std::condition_variable cvEmpty_;
80  std::condition_variable cvOverflow_;
81 
82  std::vector<std::vector<TensorCPU>> queue_;
83 };
84 } // caffe2
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:82
Copyright (c) 2016-present, Facebook, Inc.