From 2d691fe77da04492146b65f6d700bdf843902e4a Mon Sep 17 00:00:00 2001 From: Eugene Brevdo Date: Fri, 8 Apr 2016 12:34:56 -0800 Subject: Re-enable write-once, read-many semantics for TensorArray. This implementation is a bit more efficient than the previous one because the first write just performs a shallow copy. Only on an aggregation is any new memory allocated. For read-many semantics, the operations read, pack, and concat must be called with parameter clear_after_read=False. By default, the flag is set True; this means a read will remove the reference to the underlying Tensor in the TensorArray to reclaim memory in the runtime. Change: 119404140 --- tensorflow/python/ops/tensor_array_ops.py | 200 +++++++++++++++++++++--------- 1 file changed, 143 insertions(+), 57 deletions(-) (limited to 'tensorflow/python/ops/tensor_array_ops.py') diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py index 35f63db23f..b5bc6e2570 100644 --- a/tensorflow/python/ops/tensor_array_ops.py +++ b/tensorflow/python/ops/tensor_array_ops.py @@ -23,6 +23,7 @@ from __future__ import print_function from tensorflow.python.framework import dtypes as _dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import array_ops from tensorflow.python.ops import constant_op from tensorflow.python.ops import gen_data_flow_ops from tensorflow.python.ops import math_ops @@ -53,7 +54,7 @@ class TensorArray(object): """ def __init__(self, dtype, size=None, dynamic_size=None, - tensor_array_name=None, + clear_after_read=None, tensor_array_name=None, handle=None, flow=None, name=None): """Construct a new TensorArray or wrap an existing TensorArray handle. @@ -63,6 +64,9 @@ class TensorArray(object): Required if handle is not provided. dynamic_size: (optional) Python bool: If true, writes to the TensorArray can grow the TensorArray past its initial size. Default: False. + clear_after_read: Boolean (optional, default: True). If True, clear + TensorArray values after reading them. This disables read-many + semantics, but allows early release of memory. tensor_array_name: (optional) Python string: the name of the TensorArray. This is used when creating the TensorArray handle. If this value is set, handle should be None. @@ -89,7 +93,12 @@ class TensorArray(object): if handle is not None and dynamic_size is not None: raise ValueError("Cannot provide both a handle and dynamic_size " "at the same time") + if handle is not None and clear_after_read is not None: + raise ValueError("Cannot provide both a handle and clear_after_read " + "at the same time") + if clear_after_read is None: + clear_after_read = True dynamic_size = dynamic_size or False self._dtype = dtype @@ -97,13 +106,22 @@ class TensorArray(object): if handle is not None: self._handle = handle else: - self._handle = gen_data_flow_ops._tensor_array( - dtype=dtype, size=size, dynamic_size=dynamic_size, - tensor_array_name=tensor_array_name, name=scope) - if flow is not None: - self._flow = flow - else: - self._flow = constant_op.constant(0, dtype=_dtypes.float32) + if flow is not None: + with ops.colocate_with(flow): + self._handle = gen_data_flow_ops._tensor_array( + dtype=dtype, size=size, dynamic_size=dynamic_size, + clear_after_read=clear_after_read, + tensor_array_name=tensor_array_name, name=scope) + else: + self._handle = gen_data_flow_ops._tensor_array( + dtype=dtype, size=size, dynamic_size=dynamic_size, + clear_after_read=clear_after_read, + tensor_array_name=tensor_array_name, name=scope) + if flow is not None: + self._flow = flow + else: + with ops.colocate_with(self._handle): + self._flow = constant_op.constant(0, dtype=_dtypes.float32) @property def flow(self): @@ -120,80 +138,148 @@ class TensorArray(object): """The reference to the TensorArray.""" return self._handle - def grad(self, source, flow=None): + def grad(self, source, flow=None, name=None): # tensor_array_grad requires a flow input when forward # TensorArrays are dynamically sized. This forces the creation # of the grad TensorArray only once the final forward array's size # is fixed. if flow is None: flow = self.flow - g_handle = gen_data_flow_ops._tensor_array_grad( - handle=self._handle, source=source, flow_in=flow) - g = TensorArray(dtype=self._dtype, handle=g_handle, flow=flow) - return g + with ops.op_scope([self._handle], name, "TensorArrayGrad"): + with ops.colocate_with(self._handle): + g_handle = gen_data_flow_ops._tensor_array_grad( + handle=self._handle, source=source, flow_in=flow, name=name) + with ops.control_dependencies([g_handle]): + flow = array_ops.identity(flow, name="gradient_flow") + g = TensorArray(dtype=self._dtype, handle=g_handle, flow=flow) + return g def read(self, index, name=None): - """Read the value at location `index` in the TensorArray.""" - value = gen_data_flow_ops._tensor_array_read( - handle=self._handle, index=index, flow_in=self._flow, dtype=self._dtype, - name=name) - return value + """Read the value at location `index` in the TensorArray. + + Args: + index: 0-D. int32 tensor with the index to read from. + name: A name for the operation (optional). + + Returns: + The tensor at index `index`. + """ + with ops.colocate_with(self._handle): + value = gen_data_flow_ops._tensor_array_read( + handle=self._handle, index=index, flow_in=self._flow, + dtype=self._dtype, name=name) + return value def write(self, index, value, name=None): - """Write `value` into index `index` of the TensorArray.""" - flow_out = gen_data_flow_ops._tensor_array_write( - handle=self._handle, index=index, value=value, flow_in=self._flow, - name=name) - # Size below is ignored - ta = TensorArray(dtype=self._dtype, handle=self._handle) - ta._flow = flow_out - return ta + """Write `value` into index `index` of the TensorArray. + + Args: + index: 0-D. int32 scalar with the index to write to. + value: N-D. Tensor of type `dtype`. The Tensor to write to this index. + name: A name for the operation (optional). + + Returns: + A new TensorArray object with flow that ensures the write occurs. + Use this object all for subsequent operations. + """ + with ops.colocate_with(self._handle): + flow_out = gen_data_flow_ops._tensor_array_write( + handle=self._handle, index=index, value=value, flow_in=self._flow, + name=name) + ta = TensorArray(dtype=self._dtype, handle=self._handle) + ta._flow = flow_out + return ta def pack(self, name=None): - """Return the values in the TensorArray as a packed `Tensor`.""" - value = gen_data_flow_ops._tensor_array_pack( - handle=self._handle, flow_in=self._flow, dtype=self._dtype, - name=name) + """Return the values in the TensorArray as a packed `Tensor`. + + All of the values must have been written and their shapes must all match. + + Args: + name: A name for the operation (optional). + + Returns: + All the tensors in the TensorArray packed into one tensor. + """ + with ops.colocate_with(self._handle): + value = gen_data_flow_ops._tensor_array_pack( + handle=self._handle, flow_in=self._flow, dtype=self._dtype, + name=name) - return value + return value def concat(self, name=None): - """Return the values in the TensorArray as a concatenated `Tensor`.""" - value, _ = gen_data_flow_ops._tensor_array_concat( - handle=self._handle, flow_in=self._flow, dtype=self._dtype, - name=name) - return value + """Return the values in the TensorArray as a concatenated `Tensor`. + + All of the values must have been written, their ranks must match, and + and their shapes must all match for all dimensions except the first. + + Args: + name: A name for the operation (optional). + + Returns: + All the tensors in the TensorArray concatenated into one tensor. + """ + with ops.colocate_with(self._handle): + value, _ = gen_data_flow_ops._tensor_array_concat( + handle=self._handle, flow_in=self._flow, dtype=self._dtype, + name=name) + return value def unpack(self, value, name=None): - """Pack the values of a `Tensor` in the TensorArray.""" - flow_out = gen_data_flow_ops._tensor_array_unpack( - handle=self._handle, value=value, flow_in=self._flow, - name=name) - ta = TensorArray(dtype=self._dtype, handle=self._handle) - ta._flow = flow_out - return ta + """Pack the values of a `Tensor` in the TensorArray. + + Args: + value: (N+1)-D. Tensor of type `dtype`. The Tensor to unpack. + name: A name for the operation (optional). + + Returns: + A new TensorArray object with flow that ensures the unpack occurs. + Use this object all for subsequent operations. + """ + with ops.colocate_with(self._handle): + flow_out = gen_data_flow_ops._tensor_array_unpack( + handle=self._handle, value=value, flow_in=self._flow, + name=name) + ta = TensorArray(dtype=self._dtype, handle=self._handle) + ta._flow = flow_out + return ta def split(self, value, lengths, name=None): - """Split the values of a `Tensor` into the TensorArray.""" - with ops.op_scope( - [self._handle, value, lengths], name, "TensorArraySplit"): - lengths = math_ops.to_int64(lengths) - flow_out = gen_data_flow_ops._tensor_array_split( - handle=self._handle, value=value, lengths=lengths, flow_in=self._flow, - name=name) - ta = TensorArray(dtype=self._dtype, handle=self._handle) - ta._flow = flow_out - return ta + """Split the values of a `Tensor` into the TensorArray. + + Args: + value: (N+1)-D. Tensor of type `dtype`. The Tensor to split. + lengths: 1-D. int32 vector with the lengths to use when splitting + `value` along its first dimension. + name: A name for the operation (optional). + + Returns: + A new TensorArray object with flow that ensures the split occurs. + Use this object all for subsequent operations. + """ + with ops.colocate_with(self._handle): + with ops.op_scope( + [self._handle, value, lengths], name, "TensorArraySplit"): + lengths = math_ops.to_int64(lengths) + flow_out = gen_data_flow_ops._tensor_array_split( + handle=self._handle, value=value, lengths=lengths, flow_in=self._flow, + name=name) + ta = TensorArray(dtype=self._dtype, handle=self._handle) + ta._flow = flow_out + return ta def size(self, name=None): """Return the size of the TensorArray.""" - return gen_data_flow_ops._tensor_array_size( - handle=self._handle, flow_in=self.flow, name=name) + with ops.colocate_with(self._handle): + return gen_data_flow_ops._tensor_array_size( + handle=self._handle, flow_in=self.flow, name=name) def close(self, name=None): """Close the current TensorArray.""" - return gen_data_flow_ops._tensor_array_close( - handle=self._handle, name=name) + with ops.colocate_with(self._handle): + return gen_data_flow_ops._tensor_array_close( + handle=self._handle, name=name) @ops.RegisterShape("TensorArray") -- cgit v1.2.3