Merge pull request #39 from Yoctol/pick_layer

Pick layer

README.md

@@ -429,8 +429,15 @@ outputs = RNNDecoder(
             units=decoding_size
         )
     ),
-    time_steps=decoding_size
+    time_steps=decoding_length
 )(encoded)
 model = Model(inputs, outputs)
 model.compile('sgd', 'mean_squared_error')
 ```
+
+For more examples you could visit our seq2vec repository.
+
+https://github.com/Yoctol/seq2vec
+
+The seq2vec repository contains auto-encoder models which encode 
+sequences into fixed length feature vector.

test/recurrent/test_rnn_base.py

@@ -82,6 +82,7 @@ def test_mask(self):
         self.assertTrue(np.all(mask[:, :self.mask_start_point]))
 
     def test_save_load(self):
+        answer = self.model.predict(self.data)
         model_name = self.__class__.__name__ + '_temp.model'
         self.model.save(model_name)
         self.model = load_model(
@@ -94,3 +95,7 @@ def test_save_load(self):
             result.shape,
             (self.data_size, self.max_length, self.encoding_size)
         )
+        np.testing.assert_array_almost_equal(
+            answer,
+            result
+        )

test/util/test_pick.py

@@ -0,0 +1,93 @@
+'''Pick layer test case'''
+import os
+from unittest import TestCase
+
+import numpy as np
+import keras.backend as K
+from keras.models import Model, Input
+from keras.layers.core import Masking
+from keras.layers import LSTM
+from keras.models import load_model
+
+from yklz import RNNEncoder, Pick
+
+class TestPickClass(TestCase):
+
+    def setUp(self):
+        self.max_length = 10
+        self.feature_size = 30
+        self.encoding_size = 20
+        self.data_size = 100
+
+        self.mask_start_point = 7
+        self.data = np.random.rand(
+            self.data_size,
+            self.max_length,
+            self.feature_size
+        )
+        self.data[:, self.mask_start_point:, :] = 0.0
+        self.y = np.random.rand(
+            self.data_size,
+            self.encoding_size
+        )
+        self.custom_objects = {}
+        self.custom_objects['RNNEncoder'] = RNNEncoder
+        self.custom_objects['Pick'] = Pick
+        self.model = self.create_model()
+
+    def create_model(self):
+        inputs = Input(shape=(self.max_length, self.feature_size))
+        masked_inputs = Masking(0.0)(inputs)
+        encoded = RNNEncoder(
+            LSTM(
+                self.encoding_size,
+                return_sequences=True
+            )
+        )(masked_inputs)
+        outputs = Pick()(encoded)
+        model = Model(inputs, outputs)
+        model.compile('sgd', 'mean_squared_error')
+        return model
+
+    def test_output_shape(self):
+        result = self.model.predict(self.data)
+        self.assertEqual(
+            result.shape,
+            (self.data_size, self.encoding_size)
+        )
+
+    def test_output_value_not_zero(self):
+        result = self.model.predict(self.data)
+        self.assertTrue(
+            np.sum(result, dtype=bool)
+        )
+
+    def test_mask(self):
+        mask_cache_key = str(id(self.model.input)) + '_' + str(id(None))
+        mask_tensor = self.model._output_mask_cache[mask_cache_key]
+        mask = mask_tensor.eval(
+            session=K.get_session(),
+            feed_dict={self.model.input: self.data}
+        )
+        self.assertTrue(
+            np.all(mask)
+        )
+
+    def test_save_load(self):
+        answer = self.model.predict(self.data)
+        model_name = self.__class__.__name__ + '_temp.model'
+        self.model.save(model_name)
+        self.model = load_model(
+            model_name,
+            custom_objects=self.custom_objects
+        )
+        os.remove(model_name)
+        result = self.model.predict(self.data)
+        self.assertEqual(
+            result.shape,
+            (self.data_size, self.encoding_size)
+        )
+        np.testing.assert_array_almost_equal(
+            answer,
+            result
+        )

yklz/__init__.py

@@ -18,6 +18,7 @@
 '''Utility Layers'''
 from .util.padding_layer import PaddingZero
 from .util.flatten import MaskFlatten
+from .util.pick import Pick
 
 '''Wrapper Layers'''
 from .wrapper.mask_to_seq import MaskToSeq

yklz/util/pick.py

@@ -0,0 +1,48 @@
+'''Pick specific slice of tensor from input'''
+import keras.backend as K
+from keras.engine.topology import Layer
+import tensorflow as tf
+
+class Pick(Layer):
+    def __init__(
+        self,
+        timestamp=0,
+        **kwargs
+    ):
+        super(Pick, self).__init__(**kwargs)
+        self.timestamp = timestamp
+        self.supports_masking = True
+
+    def build(self, input_shape):
+        super(Pick, self).build(input_shape)
+
+    def compute_output_shape(self, input_shape):
+        return (input_shape[0], input_shape[2])
+
+    def compute_mask(self, inputs, mask):
+        if mask is None:
+            return mask
+        else:
+            return tf.slice(
+                mask,
+                [0, self.timestamp],
+                [-1, 1]
+            )
+
+    def call(self, inputs, mask=None):
+        input_shape = K.int_shape(inputs)
+        return K.reshape(
+            tf.slice(
+                inputs,
+                [0, self.timestamp, 0],
+                [-1, 1, -1]
+            ),
+            (-1, input_shape[2])
+        )
+
+    def get_config(self):
+        config = {
+            'timestamp': self.timestamp,
+        }
+        base_config = super(Pick, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))