Tensorflow: function.defun with a a while loop in the body is throwing shape error - python-3.x

Question

I am using a while loop to calculate a cost function for memory reasons. When calculating the gradient, tensorflow will store Nm tensors where Nm is the number of iterations in my while loop (this cuases the same memory issues I had with the original energy functions). I do not want that as I don't have enough memory. So I want to register a new op along with a gradient function that both use a while loop. However I am having issues with using function.defun and a while loop. To simplify things, I have a small test example below:
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.framework import function
def _run(tensor):
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
res = sess.run(tensor)
return res
#function.Defun(tf.float32,tf.float32,func_name ='tf_test_log')#,grad_func=tf_test_logGrad)
def tf_test_log(t_x,t_y):
#N = t_x.shape[0].value
condition = lambda i,m1: i<N
def body(index,x):
#return[(index+1),tf.concat([x, tf.expand_dims(tf.exp( tf.add( t_x[:,index],t_y[:,index]) ),1) ],1 ) ]
return[(index+1),tf.add(x, tf.exp( tf.add( t_x[:,0],t_y[:,0]) ) ) ]
i0 = tf.constant(0,dtype=tf.int32)
m0 = tf.zeros([N,1],dType)
ijk_0 = [i0,m0]
L,t_log_x = tf.while_loop(condition,body,ijk_0,
shape_invariants=[i0.get_shape(),
tf.TensorShape([N,None])]
)
return t_log_x
dType = tf.float32
N = np.int32(100)
t_N = tf.constant(N,dtype = tf.int32)
t_x = tf.constant(np.random.randn(N,N),dtype = dType)
t_y = tf.constant(np.random.randn(N,N),dtype = dType)
ys = _run(tf_test_log(t_x,t_y))
I then try to test the new op:
I get a Value error: The shape for while/Merge_1:0 is not an invariant for the loop. It enters the loop with shape (100, ?), but has shape after one iteration. Provide shape invariants using either the shape_invariants argument of tf.while_loop or set_shape() on the loop variables.
Note that calling
If i use a concatenate operation (instead of the add operation that gets returned by my while loop), I do not get any issues.
However, If I do not set N as a global variable (i.e. I do N = t_x.shape[0]) inside the body of the tf_test_log function, I get a Value error.
ValueError: Cannot convert a partially known TensorShape to a Tensor: (?, 1)
What is wrong with my code? Any help is greatly appreciated!
I am using python 3.5 on ubuntu 16.04 and tensorflow 1.4
full output:
ValueError Traceback (most recent call last)
~/Documents/TheEffingPhDHatersGonnaHate/PAM/defun_while.py in <module>()
51 t_x = tf.constant(np.random.randn(N,N),dtype = dType)
52 t_y = tf.constant(np.random.randn(N,N),dtype = dType)
---> 53 ys = _run(tf_test_log(t_x,t_y))
54
55
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in __call__(self, *args, **kwargs)
503
504 def __call__(self, *args, **kwargs):
--> 505 self.add_to_graph(ops.get_default_graph())
506 args = [ops.convert_to_tensor(_) for _ in args] + self._extra_inputs
507 ret, op = _call(self._signature, *args, **kwargs)
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in add_to_graph(self, g)
484 def add_to_graph(self, g):
485 """Adds this function into the graph g."""
--> 486 self._create_definition_if_needed()
487
488 # Adds this function into 'g'.
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in _create_definition_if_needed(self)
319 """Creates the function definition if it's not created yet."""
320 with context.graph_mode():
--> 321 self._create_definition_if_needed_impl()
322
323 def _create_definition_if_needed_impl(self):
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in _create_definition_if_needed_impl(self)
336 # Call func and gather the output tensors.
337 with vs.variable_scope("", custom_getter=temp_graph.getvar):
--> 338 outputs = self._func(*inputs)
339
340 # There is no way of distinguishing between a function not returning
~/Documents/TheEffingPhDHatersGonnaHate/PAM/defun_while.py in tf_test_log(t_x, t_y)
39 L,t_log_x = tf.while_loop(condition,body,ijk_0,
40 shape_invariants=[i0.get_shape(),
---> 41 tf.TensorShape([N,None])]
42 )
43 return t_log_x
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name)
2814 loop_context = WhileContext(parallel_iterations, back_prop, swap_memory) # pylint: disable=redefined-outer-name
2815 ops.add_to_collection(ops.GraphKeys.WHILE_CONTEXT, loop_context)
-> 2816 result = loop_context.BuildLoop(cond, body, loop_vars, shape_invariants)
2817 return result
2818
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py in BuildLoop(self, pred, body, loop_vars, shape_invariants)
2638 self.Enter()
2639 original_body_result, exit_vars = self._BuildLoop(
-> 2640 pred, body, original_loop_vars, loop_vars, shape_invariants)
2641 finally:
2642 self.Exit()
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py in _BuildLoop(self, pred, body, original_loop_vars, loop_vars, shape_invariants)
2619 for m_var, n_var in zip(merge_vars, next_vars):
2620 if isinstance(m_var, ops.Tensor):
-> 2621 _EnforceShapeInvariant(m_var, n_var)
2622
2623 # Exit the loop.
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py in _EnforceShapeInvariant(merge_var, next_var)
576 "Provide shape invariants using either the `shape_invariants` "
577 "argument of tf.while_loop or set_shape() on the loop variables."
--> 578 % (merge_var.name, m_shape, n_shape))
579 else:
580 if not isinstance(var, (ops.IndexedSlices, sparse_tensor.SparseTensor)):
ValueError: The shape for while/Merge_1:0 is not an invariant for the loop. It enters the loop with shape (100, ?), but has shape <unknown> after one iteration. Provide shape invariants using either the `shape_invariants` argument of tf.while_loop or set_shape() on the loop variables.

Answer 1

Thanks #Alexandre Passos for the suggestion in the comment above!
The following piece of code is a modification of the original with a set_shape function added inside the body.
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.framework import function
def _run(tensor):
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
res = sess.run(tensor)
return res
#function.Defun(tf.float32,tf.float32,tf.float32,func_name ='tf_test_logGrad')
def tf_test_logGrad(t_x,t_y,grad):
return grad
#function.Defun(tf.float32,tf.float32,func_name ='tf_test_log')#,grad_func=tf_test_logGrad)
def tf_test_log(t_x,t_y):
#N = t_x.shape[0].value
condition = lambda i,m1: i<N
def body(index,x):
#return[(index+1),tf.concat([x, tf.expand_dims(tf.exp( tf.add( t_x[:,index],t_y[:,index]) ),1) ],1 ) ]
x = tf.add(x, tf.exp( tf.add( t_x[:,0],t_y[:,0]) ) )
x.set_shape([N])
return[(index+1), x]
i0 = tf.constant(0,dtype=tf.int32)
m0 = tf.zeros([N],dType)
ijk_0 = [i0,m0]
L,t_log_x = tf.while_loop(condition,body,ijk_0,
shape_invariants=[i0.get_shape(),
tf.TensorShape([N])]
)
return t_log_x
dType = tf.float32
N = np.int32(100)
t_N = tf.constant(N,dtype = tf.int32)
t_x = tf.constant(np.random.randn(N,N),dtype = dType)
t_y = tf.constant(np.random.randn(N,N),dtype = dType)
ys = _run(tf_test_log(t_x,t_y))
The Issue of global N still persists.
You still need to set the shape of the loop tensors as a global variable outside of the defun decorator. If you try to get it from the shape of the inputs of the defun decorator, you get:
TypeError Traceback (most recent call last)
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/array_ops.py in zeros(shape, dtype, name)
1438 shape = tensor_shape.as_shape(shape)
-> 1439 output = constant(zero, shape=shape, dtype=dtype, name=name)
1440 except (TypeError, ValueError):
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/constant_op.py in constant(value, dtype, shape, name, verify_shape)
207 tensor_util.make_tensor_proto(
--> 208 value, dtype=dtype, shape=shape, verify_shape=verify_shape))
209 dtype_value = attr_value_pb2.AttrValue(type=tensor_value.tensor.dtype)
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/tensor_util.py in make_tensor_proto(values, dtype, shape, verify_shape)
379 # exception when dtype is set to np.int64
--> 380 if shape is not None and np.prod(shape, dtype=np.int64) == 0:
381 nparray = np.empty(shape, dtype=np_dt)
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/numpy/core/fromnumeric.py in prod(a, axis, dtype, out, keepdims)
2517 return _methods._prod(a, axis=axis, dtype=dtype,
-> 2518 out=out, **kwargs)
2519
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/numpy/core/_methods.py in _prod(a, axis, dtype, out, keepdims)
34 def _prod(a, axis=None, dtype=None, out=None, keepdims=False):
---> 35 return umr_prod(a, axis, dtype, out, keepdims)
36
TypeError: __int__ returned non-int (type NoneType)
During handling of the above exception, another exception occurred:
ValueError Traceback (most recent call last)
~/Documents/TheEffingPhDHatersGonnaHate/PAM/defun_while.py in <module>()
52 t_x = tf.constant(np.random.randn(N,N),dtype = dType)
53 t_y = tf.constant(np.random.randn(N,N),dtype = dType)
---> 54 ys = _run(tf_test_log(t_x,t_y))
55
56
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in __call__(self, *args, **kwargs)
503
504 def __call__(self, *args, **kwargs):
--> 505 self.add_to_graph(ops.get_default_graph())
506 args = [ops.convert_to_tensor(_) for _ in args] + self._extra_inputs
507 ret, op = _call(self._signature, *args, **kwargs)
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in add_to_graph(self, g)
484 def add_to_graph(self, g):
485 """Adds this function into the graph g."""
--> 486 self._create_definition_if_needed()
487
488 # Adds this function into 'g'.
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in _create_definition_if_needed(self)
319 """Creates the function definition if it's not created yet."""
320 with context.graph_mode():
--> 321 self._create_definition_if_needed_impl()
322
323 def _create_definition_if_needed_impl(self):
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/function.py in _create_definition_if_needed_impl(self)
336 # Call func and gather the output tensors.
337 with vs.variable_scope("", custom_getter=temp_graph.getvar):
--> 338 outputs = self._func(*inputs)
339
340 # There is no way of distinguishing between a function not returning
~/Documents/TheEffingPhDHatersGonnaHate/PAM/defun_while.py in tf_test_log(t_x, t_y)
33
34 i0 = tf.constant(0,dtype=tf.int32)
---> 35 m0 = tf.zeros([N],dType)
36
37
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/ops/array_ops.py in zeros(shape, dtype, name)
1439 output = constant(zero, shape=shape, dtype=dtype, name=name)
1440 except (TypeError, ValueError):
-> 1441 shape = ops.convert_to_tensor(shape, dtype=dtypes.int32, name="shape")
1442 output = fill(shape, constant(zero, dtype=dtype), name=name)
1443 assert output.dtype.base_dtype == dtype
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/ops.py in convert_to_tensor(value, dtype, name, preferred_dtype)
834 name=name,
835 preferred_dtype=preferred_dtype,
--> 836 as_ref=False)
837
838
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/ops.py in internal_convert_to_tensor(value, dtype, name, as_ref, preferred_dtype, ctx)
924
925 if ret is None:
--> 926 ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
927
928 if ret is NotImplemented:
~/environments/tf_1_4_gpu/lib/python3.5/site-packages/tensorflow/python/framework/constant_op.py in _tensor_shape_tensor_conversion_function(s, dtype, name, as_ref)
248 if not s.is_fully_defined():
249 raise ValueError(
--> 250 "Cannot convert a partially known TensorShape to a Tensor: %s" % s)
251 s_list = s.as_list()
252 int64_value = 0
ValueError: Cannot convert a partially known TensorShape to a Tensor: (?,)