Allow tensor buffers to automatically resize as needed.

PiperOrigin-RevId: 246594454
This commit is contained in:
Steve Chien 2019-05-03 16:29:40 -07:00 committed by A. Unique TensorFlower
parent beb86c6e18
commit 28639ba0a8
3 changed files with 161 additions and 38 deletions

View file

@ -11,8 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""A lightweight fixed-sized buffer for maintaining lists. """A lightweight buffer for maintaining tensors."""
"""
from __future__ import absolute_import from __future__ import absolute_import
from __future__ import division from __future__ import division
@ -22,7 +21,7 @@ import tensorflow as tf
class TensorBuffer(object): class TensorBuffer(object):
"""A lightweight fixed-sized buffer for maintaining lists. """A lightweight buffer for maintaining lists.
The TensorBuffer accumulates tensors of the given shape into a tensor (whose The TensorBuffer accumulates tensors of the given shape into a tensor (whose
rank is one more than that of the given shape) via calls to `append`. The rank is one more than that of the given shape) via calls to `append`. The
@ -30,12 +29,12 @@ class TensorBuffer(object):
`values`. `values`.
""" """
def __init__(self, max_size, shape, dtype=tf.int32, name=None): def __init__(self, capacity, shape, dtype=tf.int32, name=None):
"""Initializes the TensorBuffer. """Initializes the TensorBuffer.
Args: Args:
max_size: The maximum size. Attempts to append more than this many rows capacity: Initial capacity. Buffer will double in capacity each time it is
will fail with an exception. filled to capacity.
shape: The shape (as tuple or list) of the tensors to accumulate. shape: The shape (as tuple or list) of the tensors to accumulate.
dtype: The type of the tensors. dtype: The type of the tensors.
name: A string name for the variable_scope used. name: A string name for the variable_scope used.
@ -45,19 +44,24 @@ class TensorBuffer(object):
""" """
shape = list(shape) shape = list(shape)
self._rank = len(shape) self._rank = len(shape)
self._name = name
self._dtype = dtype
if not self._rank: if not self._rank:
raise ValueError('Shape cannot be scalar.') raise ValueError('Shape cannot be scalar.')
shape = [max_size] + shape shape = [capacity] + shape
with tf.variable_scope(name): with tf.variable_scope(self._name):
# We need to use a placeholder as the initial value to allow resizing.
self._buffer = tf.Variable( self._buffer = tf.Variable(
initial_value=tf.zeros(shape, dtype), initial_value=tf.placeholder_with_default(
tf.zeros(shape, dtype), shape=None),
trainable=False, trainable=False,
name='buffer') name='buffer',
self._size = tf.Variable( use_resource=True)
initial_value=0, self._current_size = tf.Variable(
trainable=False, initial_value=0, trainable=False, name='current_size')
name='size') self._capacity = tf.Variable(
initial_value=capacity, trainable=False, name='capacity')
def append(self, value): def append(self, value):
"""Appends a new tensor to the end of the buffer. """Appends a new tensor to the end of the buffer.
@ -69,23 +73,59 @@ class TensorBuffer(object):
Returns: Returns:
An op appending the new tensor to the end of the buffer. An op appending the new tensor to the end of the buffer.
""" """
with tf.control_dependencies([
tf.assert_less( def _double_capacity():
self._size, """Doubles the capacity of the current tensor buffer."""
tf.shape(self._buffer)[0], padding = tf.zeros_like(self._buffer, self._buffer.dtype)
message='Appending past end of TensorBuffer.'), new_buffer = tf.concat([self._buffer, padding], axis=0)
tf.assert_equal( if tf.executing_eagerly():
tf.shape(value), with tf.variable_scope(self._name, reuse=True):
tf.shape(self._buffer)[1:], self._buffer = tf.get_variable(
message='Appending value of inconsistent shape.')]): name='buffer',
with tf.control_dependencies( dtype=self._dtype,
[tf.assign(self._buffer[self._size, :], value)]): initializer=new_buffer,
return tf.assign_add(self._size, 1) trainable=False)
return self._buffer, tf.assign(self._capacity,
tf.multiply(self._capacity, 2))
else:
return tf.assign(
self._buffer, new_buffer,
validate_shape=False), tf.assign(self._capacity,
tf.multiply(self._capacity, 2))
update_buffer, update_capacity = tf.cond(
tf.equal(self._current_size, self._capacity),
_double_capacity, lambda: (self._buffer, self._capacity))
with tf.control_dependencies([update_buffer, update_capacity]):
with tf.control_dependencies([
tf.assert_less(
self._current_size,
self._capacity,
message='Appending past end of TensorBuffer.'),
tf.assert_equal(
tf.shape(value),
tf.shape(self._buffer)[1:],
message='Appending value of inconsistent shape.')
]):
with tf.control_dependencies(
[tf.assign(self._buffer[self._current_size, :], value)]):
return tf.assign_add(self._current_size, 1)
@property @property
def values(self): def values(self):
"""Returns the accumulated tensor.""" """Returns the accumulated tensor."""
begin_value = tf.zeros([self._rank + 1], dtype=tf.int32) begin_value = tf.zeros([self._rank + 1], dtype=tf.int32)
value_size = tf.concat( value_size = tf.concat([[self._current_size],
[[self._size], tf.constant(-1, tf.int32, [self._rank])], 0) tf.constant(-1, tf.int32, [self._rank])], 0)
return tf.slice(self._buffer, begin_value, value_size) return tf.slice(self._buffer, begin_value, value_size)
@property
def current_size(self):
"""Returns the current number of tensors in the buffer."""
return self._current_size
@property
def capacity(self):
"""Returns the current capacity of the buffer."""
return self._capacity

View file

@ -11,7 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Tests for tensor_buffer.""" """Tests for tensor_buffer in eager mode."""
from __future__ import absolute_import from __future__ import absolute_import
from __future__ import division from __future__ import division
@ -25,6 +25,7 @@ tf.enable_eager_execution()
class TensorBufferTest(tf.test.TestCase): class TensorBufferTest(tf.test.TestCase):
"""Tests for TensorBuffer in eager mode."""
def test_basic(self): def test_basic(self):
size, shape = 2, [2, 3] size, shape = 2, [2, 3]
@ -53,20 +54,30 @@ class TensorBufferTest(tf.test.TestCase):
'Appending value of inconsistent shape.'): 'Appending value of inconsistent shape.'):
my_buffer.append(tf.ones(shape=[3, 4], dtype=tf.int32)) my_buffer.append(tf.ones(shape=[3, 4], dtype=tf.int32))
def test_fail_on_overflow(self): def test_resize(self):
size, shape = 2, [2, 3] size, shape = 2, [2, 3]
my_buffer = tensor_buffer.TensorBuffer(size, shape, name='my_buffer') my_buffer = tensor_buffer.TensorBuffer(size, shape, name='my_buffer')
# First two should succeed. # Append three buffers. Third one should succeed after resizing.
my_buffer.append(tf.ones(shape=shape, dtype=tf.int32)) value1 = [[1, 2, 3], [4, 5, 6]]
my_buffer.append(tf.ones(shape=shape, dtype=tf.int32)) my_buffer.append(value1)
self.assertAllEqual(my_buffer.values.numpy(), [value1])
self.assertAllEqual(my_buffer.current_size.numpy(), 1)
self.assertAllEqual(my_buffer.capacity.numpy(), 2)
# Third one should fail. value2 = [[4, 5, 6], [7, 8, 9]]
with self.assertRaisesRegex( my_buffer.append(value2)
tf.errors.InvalidArgumentError, self.assertAllEqual(my_buffer.values.numpy(), [value1, value2])
'Appending past end of TensorBuffer.'): self.assertAllEqual(my_buffer.current_size.numpy(), 2)
my_buffer.append(tf.ones(shape=shape, dtype=tf.int32)) self.assertAllEqual(my_buffer.capacity.numpy(), 2)
value3 = [[7, 8, 9], [10, 11, 12]]
my_buffer.append(value3)
self.assertAllEqual(my_buffer.values.numpy(), [value1, value2, value3])
self.assertAllEqual(my_buffer.current_size.numpy(), 3)
# Capacity should have doubled.
self.assertAllEqual(my_buffer.capacity.numpy(), 4)
if __name__ == '__main__': if __name__ == '__main__':

View file

@ -0,0 +1,72 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for tensor_buffer in graph mode."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf
from privacy.analysis import tensor_buffer
class TensorBufferTest(tf.test.TestCase):
"""Tests for TensorBuffer in graph mode."""
def test_noresize(self):
"""Test buffer does not resize if capacity is not exceeded."""
with self.cached_session() as sess:
size, shape = 2, [2, 3]
my_buffer = tensor_buffer.TensorBuffer(size, shape, name='my_buffer')
value1 = [[1, 2, 3], [4, 5, 6]]
with tf.control_dependencies([my_buffer.append(value1)]):
value2 = [[7, 8, 9], [10, 11, 12]]
with tf.control_dependencies([my_buffer.append(value2)]):
values = my_buffer.values
current_size = my_buffer.current_size
capacity = my_buffer.capacity
self.evaluate(tf.global_variables_initializer())
v, cs, cap = sess.run([values, current_size, capacity])
self.assertAllEqual(v, [value1, value2])
self.assertEqual(cs, 2)
self.assertEqual(cap, 2)
def test_resize(self):
"""Test buffer resizes if capacity is exceeded."""
with self.cached_session() as sess:
size, shape = 2, [2, 3]
my_buffer = tensor_buffer.TensorBuffer(size, shape, name='my_buffer')
value1 = [[1, 2, 3], [4, 5, 6]]
with tf.control_dependencies([my_buffer.append(value1)]):
value2 = [[7, 8, 9], [10, 11, 12]]
with tf.control_dependencies([my_buffer.append(value2)]):
value3 = [[13, 14, 15], [16, 17, 18]]
with tf.control_dependencies([my_buffer.append(value3)]):
values = my_buffer.values
current_size = my_buffer.current_size
capacity = my_buffer.capacity
self.evaluate(tf.global_variables_initializer())
v, cs, cap = sess.run([values, current_size, capacity])
self.assertAllEqual(v, [value1, value2, value3])
self.assertEqual(cs, 3)
self.assertEqual(cap, 4)
if __name__ == '__main__':
tf.test.main()