sequence example

Author: jakeret

copy_memory.py

@@ -0,0 +1,60 @@
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras import losses, metrics, optimizers
+from tensorflow.keras.callbacks import TensorBoard
+
+import tcn
+
+
+def load_dataset(batch_size, T):
+    x_train, y_train = generate_copy_sequence(batch_size, T)
+    x_test, y_test = generate_copy_sequence(batch_size, T)
+
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1000)
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).shuffle(1000)
+    return train_dataset, test_dataset
+
+
+def generate_copy_sequence(batch_size, sequence_length):
+    x = np.zeros((batch_size, sequence_length))
+    copy_sequence = np.random.randint(0, 8, (batch_size, 10))
+    x[:, :10] = copy_sequence
+    x[:, -11:] = 9
+
+    y = np.zeros_like(x)
+    y[:, -10:] = copy_sequence
+    return x, y
+
+
+def train():
+    depth = 6
+    filters = 25
+    block_filters = [filters] * depth
+    sequence_length = 601
+
+    train_dataset, test_dataset = load_dataset(30000, sequence_length)
+
+    model = tcn.build_model(sequence_length=sequence_length,
+                            channels=1,
+                            num_classes=10,
+                            filters=block_filters,
+                            kernel_size=8,
+                            return_sequence=True)
+
+    model.compile(optimizer=optimizers.RMSprop(lr=5e-4, clipnorm=1.),
+                  metrics=[metrics.SparseCategoricalAccuracy()],
+                  loss=losses.SparseCategoricalCrossentropy())
+
+    print(model.summary())
+
+    model.fit(train_dataset.batch(32),
+              validation_data=test_dataset.batch(32),
+              callbacks=[TensorBoard(str(Path("logs") / datetime.now().strftime("%Y-%m-%dT%H-%M_%S")))],
+              epochs=10)
+
+
+if __name__ == '__main__':
+    train()

tcn.py

@@ -66,11 +66,12 @@ def call(self, inputs, training=None, **kwargs):
 
 class TCN(layers.Layer):
 
-    def __init__(self, filters, kernel_size, **kwargs):
+    def __init__(self, filters, kernel_size, return_sequence=False, **kwargs):
         super(TCN, self).__init__(**kwargs)
         self.blocks = []
         self.depth = len(filters)
         self.kernel_size = kernel_size
+        self.return_sequence = return_sequence
 
         for i in range(self.depth):
             dilation_size = 2 ** i
@@ -81,24 +82,26 @@ def __init__(self, filters, kernel_size, **kwargs):
                               name=f"residual_block_{i}")
             )
 
-        self.slice_layer = layers.Lambda(lambda tt: tt[:, -1, :])
+        if not self.return_sequence:
+            self.slice_layer = layers.Lambda(lambda tt: tt[:, -1, :])
 
     def call(self, inputs, training=None, **kwargs):
         x = inputs
         for block in self.blocks:
             x = block(x)
 
-        x = self.slice_layer(x)
+        if not self.return_sequence:
+            x = self.slice_layer(x)
         return x
 
     @property
     def receptive_field_size(self):
         return 1 + 2 * (self.kernel_size - 1) * (2 ** self.depth - 1)
 
 
-def build_model(sequence_length, channels, filters, num_classes, kernel_size):
+def build_model(sequence_length, channels, filters, num_classes, kernel_size, return_sequence=False):
     inputs = Input(shape=(sequence_length, channels), name="inputs")
-    tcn_block = TCN(filters, kernel_size)
+    tcn_block = TCN(filters, kernel_size, return_sequence)
     x = tcn_block(inputs)
 
     outputs = layers.Dense(num_classes,