Small reproduction of bug involving serialization of Lambda layer in Keras.

lambda_bug.py

"""Small reproduction of bug involving Lambda."""
import tensorflow as tf
from keras import backend as K
from keras.layers import Convolution2D, Lambda
from keras.models import Sequential, model_from_json, model_from_yaml


def space_to_depth_output_shape(input_shape):
    return (input_shape[0], input_shape[1] // 2, input_shape[2] // 2,
            4 * input_shape[3]) if input_shape[1] else (
                input_shape[0], None, None, 4* input_shape[3])


if __name__ == '__main__':
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, input_shape=(224, 224, 3)))
    # model.add(Lambda(K.resize_images, arguments={'height_factor':112, 'width_factor':112, 'dim_ordering':'tf'}))
    model.add(
        Lambda(
            tf.space_to_depth,
            output_shape=space_to_depth_output_shape,
            arguments={'block_size': 2},
            name='space_to_depth'))
    m_json = model.to_json()
    m_yaml = model.to_yaml()
    model = model_from_json(m_json)
    model = model_from_yaml(m_yaml)