Tensorflow 2.0 - LSTM statefulness and input size - tensorflow

Question

For a specific problem in reinforcement learning (inspired in this paper), I'm using a RNN which is fed with data of shape (batch_size, time_steps, features) = (1,1,1), for L data-points, and then a "cycle" is over; with a LSTM cell. I'm using lstm.stateful = True, and after L feeds to the network, I call lstm.reset_states().
Between one cycle and another one, and just after calling lstm.reset_states(), I'd like to evaluate the output of the network on an input data of the shape (batch_size, time_steps, features) = (L,1,1); and then to continue using again the RNN with input of batch_size = 1.
Furthermore, I want the code to be as optimized as possible, and for this I'm trying to use AutoGraph via the #tf.function decorators.
The problem is that I encounter an error, that can be recreated with the following example (notice that if #tf.function is removed, everything works, and I don't understand why?)
import tensorflow as tf
import numpy as np
class Actor(tf.keras.Model):
def __init__(self):
super(Actor,self).__init__()
self.lstm = tf.keras.layers.LSTM(5, return_sequences=True, stateful=True, input_shape=(None,None,1))#, input_shape=(None,None,1))
def call(self, inputs):
feat= self.lstm(inputs)
return feat
actor = Actor()
#tf.function
def g(actor):
context1 = tf.reshape(np.array([0.]*10),(10,1,1))
actor(context1)
actor.reset_states()
actor.lstm.stateful=False
context = tf.reshape(np.array([0.]),(1,1,1))
actor(context)
g(actor)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-28-4487772bee64> in <module>
23 actor(context)
24
---> 25 g(actor)
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
578 xla_context.Exit()
579 else:
--> 580 result = self._call(*args, **kwds)
581
582 if tracing_count == self._get_tracing_count():
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
625 # This is the first call of __call__, so we have to initialize.
626 initializers = []
--> 627 self._initialize(args, kwds, add_initializers_to=initializers)
628 finally:
629 # At this point we know that the initialization is complete (or less
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
504 self._concrete_stateful_fn = (
505 self._stateful_fn._get_concrete_function_internal_garbage_collected( # pylint: disable=protected-access
--> 506 *args, **kwds))
507
508 def invalid_creator_scope(*unused_args, **unused_kwds):
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
2444 args, kwargs = None, None
2445 with self._lock:
-> 2446 graph_function, _, _ = self._maybe_define_function(args, kwargs)
2447 return graph_function
2448
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
2775
2776 self._function_cache.missed.add(call_context_key)
-> 2777 graph_function = self._create_graph_function(args, kwargs)
2778 self._function_cache.primary[cache_key] = graph_function
2779 return graph_function, args, kwargs
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
2665 arg_names=arg_names,
2666 override_flat_arg_shapes=override_flat_arg_shapes,
-> 2667 capture_by_value=self._capture_by_value),
2668 self._function_attributes,
2669 # Tell the ConcreteFunction to clean up its graph once it goes out of
~/.local/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
979 _, original_func = tf_decorator.unwrap(python_func)
980
--> 981 func_outputs = python_func(*func_args, **func_kwargs)
982
983 # invariant: `func_outputs` contains only Tensors, CompositeTensors,
~/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
439 # __wrapped__ allows AutoGraph to swap in a converted function. We give
440 # the function a weak reference to itself to avoid a reference cycle.
--> 441 return weak_wrapped_fn().__wrapped__(*args, **kwds)
442 weak_wrapped_fn = weakref.ref(wrapped_fn)
443
~/.local/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
966 except Exception as e: # pylint:disable=broad-except
967 if hasattr(e, "ag_error_metadata"):
--> 968 raise e.ag_error_metadata.to_exception(e)
969 else:
970 raise
ValueError: in user code:
<ipython-input-28-4487772bee64>:23 g *
actor(context)
<ipython-input-28-4487772bee64>:11 call *
feat= self.lstm(inputs)
/home/cooper-cooper/.local/lib/python3.6/site-packages/tensorflow/python/keras/layers/recurrent.py:654 __call__ **
return super(RNN, self).__call__(inputs, **kwargs)
/home/cooper-cooper/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py:886 __call__
self.name)
/home/cooper-cooper/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/input_spec.py:227 assert_input_compatibility
', found shape=' + str(shape))
ValueError: Input 0 is incompatible with layer lstm_7: expected shape=(10, None, 1), found shape=[1, 1, 1]

Answer 1

If anyone interested, I found an answer form the following posts, and the workaround solution for me was the following:
import tensorflow as tf
import numpy as np
class Actor(tf.keras.Model):
def __init__(self):
super(Actor,self).__init__()
self.lstm = tf.keras.layers.LSTM(5, return_sequences=True, stateful=True,input_shape=(1,1))#, input_shape=(None,None,1))
def call(self, inputs):
feat= self.lstm(inputs)
return feat
def reset_states_workaround(self, new_batch_size=1):
self.lstm.states = [tf.Variable(tf.random.uniform((new_batch_size,5))), tf.Variable(tf.random.uniform((new_batch_size,5)))]
self.lstm.input_spec = [tf.keras.layers.InputSpec(shape=(new_batch_size,None,1), ndim=3)]
And then, between two different calls with #tf.function, I do:
actor = Actor()
#tf.function
def g(actor):
context1 = tf.reshape(np.array([0.]*10),(10,1,1))
preds = actor(context1)
return preds
g(actor)
actor.reset_states_workaround(new_batch_size=1)
#tf.function
def g2(actor):
context1 = tf.reshape(np.array([0.]*1),(1,1,1))
preds = actor(context1)
return preds
g2(actor)
A problem with using actor.reset_states_workaround(new_batch_size=1) inside #tf.function arises though: ValueError: tf.function-decorated function tried to create variables on non-first call., and that's why i'm using it outside.