diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2017-07-25 17:48:52 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-07-25 17:53:02 -0700 |
| commit | 42eba6a7046d232de397ea2d0e627e91f095bc4c (patch) | |
| tree | c8bbfd844051a7cd4b142c9eb6aabbaa29c898f1 | |
| parent | 4ecc31ddd04f13c4d3bde48c28ff0646aa477bca (diff) | |
Implemented MemoryUsed function for hash tables.
PiperOrigin-RevId: 163149969
| -rw-r--r-- | tensorflow/core/kernels/lookup_table_op.cc | 34 |
1 files changed, 34 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/lookup_table_op.cc b/tensorflow/core/kernels/lookup_table_op.cc index b33b98975d..d721b3d542 100644 --- a/tensorflow/core/kernels/lookup_table_op.cc +++ b/tensorflow/core/kernels/lookup_table_op.cc @@ -124,6 +124,20 @@ class MutableHashTableOfScalars final : public LookupInterface { TensorShape value_shape() const override { return TensorShape(); } + int64 MemoryUsed() const override { + int64 ret = 0; + mutex_lock l(mu_); + for (unsigned i = 0; i < table_.bucket_count(); ++i) { + size_t bucket_size = table_.bucket_size(i); + if (bucket_size == 0) { + ret++; + } else { + ret += bucket_size; + } + } + return sizeof(MutableHashTableOfScalars) + ret; + } + private: // TODO(andreasst): consider using a read/write lock or a concurrent map mutable mutex mu_; @@ -239,6 +253,20 @@ class MutableHashTableOfTensors final : public LookupInterface { TensorShape value_shape() const override { return value_shape_; } + int64 MemoryUsed() const override { + int64 ret = 0; + mutex_lock l(mu_); + for (unsigned i = 0; i < table_.bucket_count(); ++i) { + size_t bucket_size = table_.bucket_size(i); + if (bucket_size == 0) { + ret++; + } else { + ret += bucket_size; + } + } + return sizeof(MutableHashTableOfTensors) + ret; + } + private: TensorShape value_shape_; // TODO(andreasst): consider using a read/write lock or a concurrent map @@ -467,6 +495,12 @@ class MutableDenseHashTable final : public LookupInterface { TensorShape value_shape() const override { return value_shape_; } + int64 MemoryUsed() const override { + mutex_lock l(mu_); + return sizeof(MutableDenseHashTable) + key_buckets_.AllocatedBytes() + + value_buckets_.AllocatedBytes() + empty_key_.AllocatedBytes(); + } + private: Status DoInsert(OpKernelContext* ctx, const Tensor& key, const Tensor& value, bool ignore_empty_key) EXCLUSIVE_LOCKS_REQUIRED(mu_) { |