Serialization of Attention Layer

This might be a Keras problem but have you tried serializing some of the layers? I tried the following to save a model that contains SoftAttention:

    from keras.engine import Input, Model
    from keras.layers import Embedding, SimpleRNN
    from keras.models import model_from_json, model_from_yaml

    embedding_matrix = np.random.randn(100).reshape((10, 10))
    num_vocab, embedding_size = embedding_matrix.shape

    input_layer = Input(shape=(10, ), dtype='int32', name="data_padded")
    embedding_layer = Embedding(num_vocab, embedding_size, input_length=10, mask_zero=True)(input_layer)
    rnn_layer = SimpleRNN(5, return_sequences=True)(embedding_layer)
    attention_layer = SoftAttention()(rnn_layer)
    loss_layer = Dense(1, activation='sigmoid')(attention_layer)

    model = Model(input_layer, loss_layer)
    model.summary()
    model.compile('adam', loss='binary_crossentropy', metrics=['accuracy'])

    x = np.random.randint(0, 10, 10).reshape(1, -1)
    y = np.random.randint(0, 2, 1)
    model.fit(x, y, nb_epoch=1)

    path = ""

    # -- save
    config_text = model.to_json()
    with open(path + '_config', 'w') as f:
        f.write(config_text)

    model.save_weights(path + '_model', overwrite=True)

    # -- load
    m = model_from_json(open(path + '_config', 'r').read(), custom_objects={
        'SoftAttention': SoftAttention
    })
    m.load_weights(path + '_model')

And I get this error:

Traceback (most recent call last):
  File "~/src/deep_learning/layers/attention.py", line 329, in <module>
    'SoftAttention': SoftAttention
  File "~/lib/python2.7/site-packages/keras/models.py", line 209, in model_from_json
    return layer_from_config(config, custom_objects=custom_objects)
  File "~/lib/python2.7/site-packages/keras/utils/layer_utils.py", line 34, in layer_from_config
    return layer_class.from_config(config['config'])
  File "~/lib/python2.7/site-packages/keras/engine/topology.py", line 2395, in from_config
    process_layer(layer_data)
  File "~/lib/python2.7/site-packages/keras/engine/topology.py", line 2390, in process_layer
    layer(input_tensors[0])
  File "~/lib/python2.7/site-packages/keras/engine/topology.py", line 517, in __call__
    self.add_inbound_node(inbound_layers, node_indices, tensor_indices)
  File "~/lib/python2.7/site-packages/keras/engine/topology.py", line 571, in add_inbound_node
    Node.create_node(self, inbound_layers, node_indices, tensor_indices)
  File "~/lib/python2.7/site-packages/keras/engine/topology.py", line 174, in create_node
    str(output_shapes))
Exception: The `get_output_shape_for` method of layer "softattention_1"" should return one shape tuple per output tensor of the layer. Found: [(None, 5)]

I think it's related to JSON not serializing optional arguments.

Thanks

braingineer / ikelos

Serialization of Attention Layer #4