/src/tesseract/src/lstm/networkscratch.h
Line | Count | Source (jump to first uncovered line) |
1 | | /////////////////////////////////////////////////////////////////////// |
2 | | // File: networkscratch.h |
3 | | // Description: Scratch space for Network layers that hides distinction |
4 | | // between float/int implementations. |
5 | | // Author: Ray Smith |
6 | | // |
7 | | // (C) Copyright 2014, Google Inc. |
8 | | // Licensed under the Apache License, Version 2.0 (the "License"); |
9 | | // you may not use this file except in compliance with the License. |
10 | | // You may obtain a copy of the License at |
11 | | // http://www.apache.org/licenses/LICENSE-2.0 |
12 | | // Unless required by applicable law or agreed to in writing, software |
13 | | // distributed under the License is distributed on an "AS IS" BASIS, |
14 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
15 | | // See the License for the specific language governing permissions and |
16 | | // limitations under the License. |
17 | | /////////////////////////////////////////////////////////////////////// |
18 | | |
19 | | #ifndef TESSERACT_LSTM_NETWORKSCRATCH_H_ |
20 | | #define TESSERACT_LSTM_NETWORKSCRATCH_H_ |
21 | | |
22 | | #include <mutex> |
23 | | #include "matrix.h" |
24 | | #include "networkio.h" |
25 | | |
26 | | namespace tesseract { |
27 | | |
28 | | // Generic scratch space for network layers. Provides NetworkIO that can store |
29 | | // a complete set (over time) of intermediates, and vector<float> |
30 | | // scratch space that auto-frees after use. The aim here is to provide a set |
31 | | // of temporary buffers to network layers that can be reused between layers |
32 | | // and don't have to be reallocated on each call. |
33 | | class NetworkScratch { |
34 | | public: |
35 | 2 | NetworkScratch() : int_mode_(false) {} |
36 | 0 | ~NetworkScratch() = default; |
37 | | |
38 | | // Sets the network representation. If the representation is integer, then |
39 | | // default (integer) NetworkIOs are separated from the always-float variety. |
40 | | // This saves memory by having separate int-specific and float-specific |
41 | | // stacks. If the network representation is float, then all NetworkIOs go |
42 | | // to the float stack. |
43 | 0 | void set_int_mode(bool int_mode) { |
44 | 0 | int_mode_ = int_mode; |
45 | 0 | } |
46 | | |
47 | | // Class that acts like a NetworkIO (by having an implicit cast operator), |
48 | | // yet actually holds a pointer to NetworkIOs in the source NetworkScratch, |
49 | | // and knows how to unstack the borrowed pointers on destruction. |
50 | | class IO { |
51 | | public: |
52 | | // The NetworkIO should be sized after construction. |
53 | | IO(const NetworkIO &src, NetworkScratch *scratch) |
54 | 690k | : int_mode_(scratch->int_mode_ && src.int_mode()), scratch_space_(scratch) { |
55 | 690k | network_io_ = |
56 | 690k | int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow(); |
57 | 690k | } |
58 | | // Default constructor for arrays. Use one of the Resize functions |
59 | | // below to initialize and size. |
60 | 345k | IO() : int_mode_(false), network_io_(nullptr), scratch_space_(nullptr) {} |
61 | | |
62 | 1.03M | ~IO() { |
63 | 1.03M | if (scratch_space_ == nullptr) { |
64 | 345k | ASSERT_HOST(network_io_ == nullptr); |
65 | 690k | } else if (int_mode_) { |
66 | 0 | scratch_space_->int_stack_.Return(network_io_); |
67 | 690k | } else { |
68 | 690k | scratch_space_->float_stack_.Return(network_io_); |
69 | 690k | } |
70 | 1.03M | } |
71 | | // Resizes the array (and stride), avoiding realloc if possible, to the |
72 | | // size from various size specs: |
73 | | // Same time size, given number of features. |
74 | 0 | void Resize(const NetworkIO &src, int num_features, NetworkScratch *scratch) { |
75 | 0 | if (scratch_space_ == nullptr) { |
76 | 0 | int_mode_ = scratch->int_mode_ && src.int_mode(); |
77 | 0 | scratch_space_ = scratch; |
78 | 0 | network_io_ = |
79 | 0 | int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow(); |
80 | 0 | } |
81 | 0 | network_io_->Resize(src, num_features); |
82 | 0 | } |
83 | | // Resizes to a specific size as a temp buffer. No batches, no y-dim. |
84 | 0 | void Resize2d(bool int_mode, int width, int num_features, NetworkScratch *scratch) { |
85 | 0 | if (scratch_space_ == nullptr) { |
86 | 0 | int_mode_ = scratch->int_mode_ && int_mode; |
87 | 0 | scratch_space_ = scratch; |
88 | 0 | network_io_ = |
89 | 0 | int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow(); |
90 | 0 | } |
91 | 0 | network_io_->Resize2d(int_mode, width, num_features); |
92 | 0 | } |
93 | | // Resize forcing a float representation with the width of src and the given |
94 | | // number of features. |
95 | 0 | void ResizeFloat(const NetworkIO &src, int num_features, NetworkScratch *scratch) { |
96 | 0 | if (scratch_space_ == nullptr) { |
97 | 0 | int_mode_ = false; |
98 | 0 | scratch_space_ = scratch; |
99 | 0 | network_io_ = scratch_space_->float_stack_.Borrow(); |
100 | 0 | } |
101 | 0 | network_io_->ResizeFloat(src, num_features); |
102 | 0 | } |
103 | | |
104 | | // Returns a ref to a NetworkIO that enables *this to be treated as if |
105 | | // it were just a NetworkIO*. |
106 | 1.03M | NetworkIO &operator*() { |
107 | 1.03M | return *network_io_; |
108 | 1.03M | } |
109 | 0 | NetworkIO *operator->() { |
110 | 0 | return network_io_; |
111 | 0 | } |
112 | 1.03M | operator NetworkIO *() { |
113 | 1.03M | return network_io_; |
114 | 1.03M | } |
115 | | |
116 | | private: |
117 | | // True if this is from the always-float stack, otherwise the default stack. |
118 | | bool int_mode_; |
119 | | // The NetworkIO that we have borrowed from the scratch_space_. |
120 | | NetworkIO *network_io_; |
121 | | // The source scratch_space_. Borrowed pointer, used to free the |
122 | | // NetworkIO. Don't delete! |
123 | | NetworkScratch *scratch_space_; |
124 | | }; // class IO. |
125 | | |
126 | | // Class that acts like a fixed array of float, yet actually uses space |
127 | | // from a vector<float> in the source NetworkScratch, and knows how |
128 | | // to unstack the borrowed vector on destruction. |
129 | | class FloatVec { |
130 | | public: |
131 | | // The array will have size elements in it, uninitialized. |
132 | 0 | FloatVec(int size, NetworkScratch *scratch) : vec_(nullptr), scratch_space_(scratch) { |
133 | 0 | Init(size, scratch); |
134 | 0 | } |
135 | | // Default constructor is for arrays. Use Init to setup. |
136 | 3.45M | FloatVec() : vec_(nullptr), data_(nullptr), scratch_space_(nullptr) {} |
137 | 3.45M | ~FloatVec() { |
138 | 3.45M | if (scratch_space_ != nullptr) { |
139 | 3.10M | scratch_space_->vec_stack_.Return(vec_); |
140 | 3.10M | } |
141 | 3.45M | } |
142 | | |
143 | 3.10M | void Init(int /*size*/, int reserve, NetworkScratch *scratch) { |
144 | 3.10M | if (scratch_space_ != nullptr && vec_ != nullptr) { |
145 | 0 | scratch_space_->vec_stack_.Return(vec_); |
146 | 0 | } |
147 | 3.10M | scratch_space_ = scratch; |
148 | 3.10M | vec_ = scratch_space_->vec_stack_.Borrow(); |
149 | | // TODO: optimize. |
150 | 3.10M | vec_->resize(reserve); |
151 | 3.10M | data_ = &(*vec_)[0]; |
152 | 3.10M | } |
153 | | |
154 | 1.38M | void Init(int size, NetworkScratch *scratch) { |
155 | 1.38M | Init(size, size, scratch); |
156 | 1.38M | } |
157 | | |
158 | | // Use the cast operator instead of operator[] so the FloatVec can be used |
159 | | // as a TFloat* argument to a function call. |
160 | 295M | operator TFloat *() const { |
161 | 295M | return data_; |
162 | 295M | } |
163 | 0 | TFloat *get() { |
164 | 0 | return data_; |
165 | 0 | } |
166 | | |
167 | | private: |
168 | | // Vector borrowed from the scratch space. Use Return to free it. |
169 | | std::vector<TFloat> *vec_; |
170 | | // Short-cut pointer to the underlying array. |
171 | | TFloat *data_; |
172 | | // The source scratch_space_. Borrowed pointer, used to free the |
173 | | // vector. Don't delete! |
174 | | NetworkScratch *scratch_space_; |
175 | | }; // class FloatVec |
176 | | |
177 | | // Class that acts like a 2-D array of TFloat, yet actually uses space |
178 | | // from the source NetworkScratch, and knows how to unstack the borrowed |
179 | | // array on destruction. |
180 | | class GradientStore { |
181 | | public: |
182 | | // Default constructor is for arrays. Use Init to setup. |
183 | 0 | GradientStore() : array_(nullptr), scratch_space_(nullptr) {} |
184 | 0 | ~GradientStore() { |
185 | 0 | if (scratch_space_ != nullptr) { |
186 | 0 | scratch_space_->array_stack_.Return(array_); |
187 | 0 | } |
188 | 0 | } |
189 | | |
190 | 0 | void Init(int size1, int size2, NetworkScratch *scratch) { |
191 | 0 | if (scratch_space_ != nullptr && array_ != nullptr) { |
192 | 0 | scratch_space_->array_stack_.Return(array_); |
193 | 0 | } |
194 | 0 | scratch_space_ = scratch; |
195 | 0 | array_ = scratch_space_->array_stack_.Borrow(); |
196 | 0 | array_->Resize(size1, size2, 0.0); |
197 | 0 | } |
198 | | |
199 | | // Accessors to get to the underlying TransposedArray. |
200 | 0 | TransposedArray *get() const { |
201 | 0 | return array_; |
202 | 0 | } |
203 | 0 | const TransposedArray &operator*() const { |
204 | 0 | return *array_; |
205 | 0 | } |
206 | | |
207 | | private: |
208 | | // Array borrowed from the scratch space. Use Return to free it. |
209 | | TransposedArray *array_; |
210 | | // The source scratch_space_. Borrowed pointer, used to free the |
211 | | // vector. Don't delete! |
212 | | NetworkScratch *scratch_space_; |
213 | | }; // class GradientStore |
214 | | |
215 | | // Class that does the work of holding a stack of objects, a stack pointer |
216 | | // and a vector of in-use flags, so objects can be returned out of order. |
217 | | // It is safe to attempt to Borrow/Return in multiple threads. |
218 | | template <typename T> |
219 | | class Stack { |
220 | | public: |
221 | 8 | Stack() = default; tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Stack() Line | Count | Source | 221 | 4 | Stack() = default; |
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Stack() Line | Count | Source | 221 | 2 | Stack() = default; |
tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Stack() Line | Count | Source | 221 | 2 | Stack() = default; |
|
222 | | |
223 | 0 | ~Stack() { |
224 | 0 | for (auto data : stack_) { |
225 | 0 | delete data; |
226 | 0 | } |
227 | 0 | } Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::~Stack() Unexecuted instantiation: tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::~Stack() Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::~Stack() |
228 | | |
229 | | // Lends out the next free item, creating one if none available, sets |
230 | | // the used flags and increments the stack top. |
231 | 3.79M | T *Borrow() { |
232 | 3.79M | std::lock_guard<std::mutex> lock(mutex_); |
233 | 3.79M | if (stack_top_ == stack_.size()) { |
234 | 12 | stack_.push_back(new T); |
235 | 12 | flags_.push_back(false); |
236 | 12 | } |
237 | 3.79M | flags_[stack_top_] = true; |
238 | 3.79M | return stack_[stack_top_++]; |
239 | 3.79M | } tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Borrow() Line | Count | Source | 231 | 690k | T *Borrow() { | 232 | 690k | std::lock_guard<std::mutex> lock(mutex_); | 233 | 690k | if (stack_top_ == stack_.size()) { | 234 | 4 | stack_.push_back(new T); | 235 | 4 | flags_.push_back(false); | 236 | 4 | } | 237 | 690k | flags_[stack_top_] = true; | 238 | 690k | return stack_[stack_top_++]; | 239 | 690k | } |
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Borrow() Line | Count | Source | 231 | 3.10M | T *Borrow() { | 232 | 3.10M | std::lock_guard<std::mutex> lock(mutex_); | 233 | 3.10M | if (stack_top_ == stack_.size()) { | 234 | 8 | stack_.push_back(new T); | 235 | 8 | flags_.push_back(false); | 236 | 8 | } | 237 | 3.10M | flags_[stack_top_] = true; | 238 | 3.10M | return stack_[stack_top_++]; | 239 | 3.10M | } |
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Borrow() |
240 | | // Takes back the given item, and marks it free. Item does not have to be |
241 | | // the most recently lent out, but free slots don't get re-used until the |
242 | | // blocking item is returned. The assumption is that there will only be |
243 | | // small, temporary variations from true stack use. (Determined by the order |
244 | | // of destructors within a local scope.) |
245 | 3.79M | void Return(T *item) { |
246 | 3.79M | std::lock_guard<std::mutex> lock(mutex_); |
247 | | // Linear search will do. |
248 | 3.79M | int index = stack_top_; |
249 | 3.79M | while (--index >= 0 && stack_[index] != item) { |
250 | 0 | } |
251 | 3.79M | if (index >= 0) { |
252 | 3.79M | flags_[index] = false; |
253 | 3.79M | } |
254 | 7.59M | while (stack_top_ > 0 && !flags_[stack_top_ - 1]) { |
255 | 3.79M | --stack_top_; |
256 | 3.79M | } |
257 | 3.79M | } tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Return(tesseract::NetworkIO*) Line | Count | Source | 245 | 690k | void Return(T *item) { | 246 | 690k | std::lock_guard<std::mutex> lock(mutex_); | 247 | | // Linear search will do. | 248 | 690k | int index = stack_top_; | 249 | 690k | while (--index >= 0 && stack_[index] != item) { | 250 | 0 | } | 251 | 690k | if (index >= 0) { | 252 | 690k | flags_[index] = false; | 253 | 690k | } | 254 | 1.38M | while (stack_top_ > 0 && !flags_[stack_top_ - 1]) { | 255 | 690k | --stack_top_; | 256 | 690k | } | 257 | 690k | } |
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Return(std::__1::vector<float, std::__1::allocator<float> >*) Line | Count | Source | 245 | 3.10M | void Return(T *item) { | 246 | 3.10M | std::lock_guard<std::mutex> lock(mutex_); | 247 | | // Linear search will do. | 248 | 3.10M | int index = stack_top_; | 249 | 3.10M | while (--index >= 0 && stack_[index] != item) { | 250 | 0 | } | 251 | 3.10M | if (index >= 0) { | 252 | 3.10M | flags_[index] = false; | 253 | 3.10M | } | 254 | 6.21M | while (stack_top_ > 0 && !flags_[stack_top_ - 1]) { | 255 | 3.10M | --stack_top_; | 256 | 3.10M | } | 257 | 3.10M | } |
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Return(tesseract::TransposedArray*) |
258 | | |
259 | | private: |
260 | | std::vector<T *> stack_; |
261 | | std::vector<bool> flags_; |
262 | | unsigned stack_top_ = 0; |
263 | | std::mutex mutex_; |
264 | | }; // class Stack. |
265 | | |
266 | | private: |
267 | | // If true, the network weights are int8_t, if false, float. |
268 | | bool int_mode_; |
269 | | // Stacks of NetworkIO and vector<float>. Once allocated, they are not |
270 | | // deleted until the NetworkScratch is deleted. |
271 | | Stack<NetworkIO> int_stack_; |
272 | | Stack<NetworkIO> float_stack_; |
273 | | Stack<std::vector<TFloat>> vec_stack_; |
274 | | Stack<TransposedArray> array_stack_; |
275 | | }; |
276 | | |
277 | | } // namespace tesseract. |
278 | | |
279 | | #endif // TESSERACT_LSTM_NETWORKSCRATCH_H_ |