Function to convert a Keras LSTM model trained as stateless to a stateful model expecting a single sample and time step as input to use in inference.

convert_to_inference_model.py

import json
from keras.models import model_from_json

def convert_to_inference_model(original_model):
    original_model_json = original_model.to_json()
    inference_model_dict = json.loads(original_model_json)

    layers = inference_model_dict['config']
    for layer in layers:
        if 'stateful' in layer['config']:
            layer['config']['stateful'] = True

        if 'batch_input_shape' in layer['config']:
            layer['config']['batch_input_shape'][0] = 1
            layer['config']['batch_input_shape'][1] = None

    inference_model = model_from_json(json.dumps(inference_model_dict))
    inference_model.set_weights(original_model.get_weights())

    return inference_model

I use the same way to make the stateless model to stateful, while it doesn't change its behaviour while testing.

Works in keras 2.3.0-tf if you change:

layers = inference_model_dict['config']

to

layers = inference_model_dict['config']['layers']

I also added a change to a Reshape layer, which was reshaping to the number of timesteps I used in training.