Keras LSTM predict 1 timestep at a time - keras

Question

Edited to add:
I found what I think is a working solution: https://bleyddyn.github.io/posts/2017/10/keras-lstm/
I'm trying to use a Conv/LSTM network for controlling a robot. I think I have everything set up so I could start training it on batches of data from a replay memory, but I can't figure out how to actually use it to control a robot. Simplified test code is below.
import numpy as np
from keras.models import Sequential
from keras.layers import Dense, Flatten, Input
from keras.layers import Convolution2D
from keras.layers.recurrent import LSTM
from keras.layers.wrappers import TimeDistributed
from keras.utils import to_categorical
def make_model(num_actions, timesteps, input_dim, l2_reg=0.005 ):
input_shape=(timesteps,) + input_dim
model = Sequential()
model.add(TimeDistributed( Convolution2D(8, (3, 3), strides=(2,2), activation='relu' ), input_shape=input_shape) )
model.add(TimeDistributed( Convolution2D(16, (3, 3), strides=(2,2), activation='relu', ) ))
model.add(TimeDistributed( Convolution2D(32, (3, 3), strides=(2,2), activation='relu', ) ))
model.add(TimeDistributed(Flatten()))
model.add(LSTM(512, return_sequences=True, activation='relu', unroll=True))
model.add(Dense(num_actions, activation='softmax', ))
model.compile(loss='categorical_crossentropy', optimizer='adam' )
return model
batch_size = 16
timesteps = 10
num_actions = 6
model = make_model( num_actions, timesteps, (84,84,3) )
model.summary()
# Fake training batch. Would be pulled from a replay memory
batch = np.random.uniform( low=0, high=255, size=(batch_size,timesteps,84,84,3) )
y = np.random.randint( 0, high=5, size=(160) )
y = to_categorical( y, num_classes=num_actions )
y = y.reshape( batch_size, timesteps, num_actions )
# stateful should be false here
pred = model.train_on_batch( batch, y )
# move trained network to robot
# This works, but it isn't practical to not get outputs (actions) until after 10 timesteps and I don't think the LSTM internal state would be correct if I tried a rolling queue of input images.
batch = np.random.uniform( low=0, high=255, size=(1,timesteps,84,84,3) )
pred = model.predict( batch, batch_size=1 )
# This is what I would need to do on my robot, with the LSTM keeping state between calls to predict
max_time = 10 # or 100000, or forever, etc.
for i in range(max_time) :
image = np.random.uniform( low=0, high=255, size=(1,1,84,84,3) ) # pull one image from camera
# stateful should be true here
pred = model.predict( image, batch_size=1 )
# take action based on pred
The error I get on the "model.predict( image..." line is:
ValueError: Error when checking : expected time_distributed_1_input to have shape (None, 10, 84, 84, 3) but got array with shape (1, 1, 84, 84, 3)
Which is understandable, but I can't find a way around it.
I don't know Keras well enough to even know if I'm using the TimeDistributed layers correctly.
So, is this even possible in Keras? If so, how?
If not, is it possible in TF or PyTorch?
Thanks for any suggestions!
Edited to add running code, although it's not necessarily correct. Still needs to be tested on an OpenAI gym task.
import numpy as np
from keras.models import Sequential
from keras.layers import Dense, Flatten, Input
from keras.layers import Convolution2D
from keras.layers.recurrent import LSTM
from keras.layers.wrappers import TimeDistributed
from keras.utils import to_categorical
def make_model(num_actions, timesteps, input_dim, l2_reg=0.005 ):
input_shape=(1,None) + input_dim
model = Sequential()
model.add(TimeDistributed( Convolution2D(8, (3, 3), strides=(2,2), activation='relu' ), batch_input_shape=input_shape) )
model.add(TimeDistributed( Convolution2D(16, (3, 3), strides=(2,2), activation='relu', ) ))
model.add(TimeDistributed( Convolution2D(32, (3, 3), strides=(2,2), activation='relu', ) ))
model.add(TimeDistributed(Flatten()))
model.add(LSTM(512, return_sequences=True, activation='relu', stateful=True))
model.add(Dense(num_actions, activation='softmax', ))
model.compile(loss='categorical_crossentropy', optimizer='adam' )
return model
batch_size = 16
timesteps = 10
num_actions = 6
model = make_model( num_actions, 1, (84,84,3) )
model.summary()
# Fake training batch. Would be pulled from a replay memory
batch = np.random.uniform( low=0, high=255, size=(batch_size,timesteps,84,84,3) )
y = np.random.randint( 0, high=5, size=(160) )
y = to_categorical( y, num_classes=num_actions )
y = y.reshape( batch_size, timesteps, num_actions )
# Need to find a way to prevent the optimizer from updating every b, but accumulate updates over an entire batch (batch_size).
for b in range(batch_size):
pred = model.train_on_batch( np.reshape(batch[b,:], (1,timesteps,84,84,3)), np.reshape(y[b,:], (1,timesteps,num_actions)) )
#for t in range(timesteps):
# pred = model.train_on_batch( np.reshape(batch[b,t,:], (1,1,84,84,3)), np.reshape(y[b,t,:], (1,1,num_actions)) )
model.reset_states() # Don't carry internal state between batches
# move trained network to robot
# This works, but it isn't practical to not get outputs (actions) until after 10 timesteps
#batch = np.random.uniform( low=0, high=255, size=(1,timesteps,84,84,3) )
#pred = model.predict( batch, batch_size=1 )
# This is what I would need to do on my robot, with the LSTM keeping state between calls to predict
max_time = 10 # or 100000, or forever, etc.
for i in range(max_time) :
image = np.random.uniform( low=0, high=255, size=(1,1,84,84,3) ) # pull one image from camera
# stateful should be true here
pred = model.predict( image, batch_size=1 )
# take action based on pred
print( pred )

Answer 1

The first thing you need is to understand your data.
Do these 5 dimensions mean anything?
I'll try to guess:
- 1 learning example
- 1 time step (this is added by TimeDistributed, normal 2D convolutions don't take this)
- 84 image side
- 84 another image side
- 3 channels (RGB)
The purpose of TimeDistributed is to add that extra timesteps dimension, so you can simulate a sequence in layers that are not supposed to work with sequences.
Your error message is telling you this:
Your input_shape parameter is (None, 10, 84, 84, 3), where None is the batch size (number of samples/examples).
Your input data, which is batch in your code is (1, 1, 84, 84, 3).
There is a mismatch, you are supposed to use batches containing 10 time steps (as defined by your input_shape). It's ok for the stateful=False model to pack 10 images in a batch and train with that.
But later, in the stateful=True case, you will need that input_shape to be just one step. (You either create a new model just for predicting and copy all weights from the training model to the predicting model, or you can try to use None in that time steps dimension, meaning you can train and predict with different amounts of time steps)
Now, differently from the convolutionals, the LSTM layer is already expecting time steps. So you should find a way to squeeze your data in less dimensions.
The LSTM will expect (None, timeSteps, features). The time steps are the same as the previous. 10 for training, 1 for predicting, and you could try to go with None there.
So, instead of a Flatten() inside a TimeDistributed, you should simply reshape the data, condensing the dimensions that are not batch size or steps:
model.add(Reshape((8,9*9*32))) #the batch size doesn't participate in this definition, and it will remain as it is.
The 9*9*32 are the sides of the preceding convolutional and its 32 filters. (I'm just not sure the sides are 9, maybe they're 8, you can see in the current model.summary()).
Finally, for the stateful=True case, you will have to define the model with batch_shape instead of input_shape. The amount of samples in a batch must be a fixed number, because the model will assume the samples in the second batch are new steps belonging to the samples in the previous batch. (The number of samples will then need to be the same for all batches).