blob: eaeeeec83fb7b114ad2c0504af8c7e28489f8ead (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
// The utility to read checkpoints for google brain tensor ops and v3
// checkpoints for dist_belief.
//
#ifndef TENSORFLOW_UTIL_TENSOR_SLICE_READER_CACHE_H_
#define TENSORFLOW_UTIL_TENSOR_SLICE_READER_CACHE_H_
#include <unordered_map>
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/util/tensor_slice_reader.h"
#include "tensorflow/core/public/status.h"
namespace tensorflow {
namespace checkpoint {
class TensorSliceReaderCache;
// Wrapper to a lazily allocated TensorSliceReaderCache.
class TensorSliceReaderCacheWrapper {
public:
TensorSliceReaderCacheWrapper();
~TensorSliceReaderCacheWrapper();
// Same as TensorSliceReaderCache::GetReader().
const TensorSliceReader* GetReader(
const string& filepattern,
TensorSliceReader::OpenTableFunction open_function,
int preferred_shard) const;
private:
mutable mutex mu_;
mutable TensorSliceReaderCache* cache_ = nullptr;
};
// A cache of TensorSliceReaders.
class TensorSliceReaderCache {
public:
TensorSliceReaderCache();
~TensorSliceReaderCache();
// Returns the TensorSliceReader corresponding to 'filepattern' and the
// open_function. May return nullptr if we can not create a new
// TensorSliceReader for the filepattern/open_function combination.
const TensorSliceReader* GetReader(
const string& filepattern,
TensorSliceReader::OpenTableFunction open_function, int preferred_shard);
private:
// Need to use a regular function type in the key map as std::function does
// not support ==.
typedef Status (*OpenFuncType)(const string&, TensorSliceReader::Table**);
// Protects attributes below.
mutex mu_;
// Maps of opened readers.
std::unordered_map<string, std::pair<OpenFuncType, TensorSliceReader*>>
readers_;
// Set of keys that a previous GetReader() call is still trying to populate.
std::set<string> still_opening_;
// Condition variable to notify when a reader has been created.
condition_variable cv_;
};
} // namespace checkpoint
} // namespace tensorflow
#endif // TENSORFLOW_UTIL_TENSOR_SLICE_READER_CACHE_H_
|
