Workaround for lambda bug seen here: https://gist.github.com/allanzelener/b3e365e68485c965aaa26e1bdf644098

lambda_workaround.py

"""Fix + workaround for bug involving Lambda."""
import numpy as np
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 resize_im(x, height_factor, width_factor, dim_ordering):
    return K.resize_images(x, height_factor, width_factor, dim_ordering)


def space_to_depth(x, block_size):
    import tensorflow as tf
    return tf.space_to_depth(x, block_size=block_size)


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,
                                                    input_shape[3] * 4)


if __name__ == '__main__':
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, input_shape=(224, 224, 3)))
    # model.add(Lambda(resize_im, arguments={'height_factor':2, 'width_factor':2, 'dim_ordering':'tf'}))
    model.add(
        Lambda(
            space_to_depth,
            arguments={'block_size': 2},
            name='space_to_depth',
            output_shape=space_to_depth_output_shape))

    m_json = model.to_json()
    m_yaml = model.to_yaml()
    m2 = model_from_json(m_json)
    m2y = model_from_yaml(m_yaml)

    x = np.random.randn(5, 224, 224, 3)
    print(m2.predict(x).shape)