import inspect
from typing import List
from tensorflow.keras import backend as K, Model, Input, optimizers
from tensorflow.keras import layers
from tensorflow.keras.layers import Activation, SpatialDropout1D, Lambda
from tensorflow.keras.layers import Layer, Conv1D, SeparableConv1D, Dense, BatchNormalization, LayerNormalization
def is_power_of_two(num: int):
return num != 0 and ((num & (num - 1)) == 0)
def adjust_dilations(dilations: list):
if all([is_power_of_two(i) for i in dilations]):
return dilations
else:
new_dilations = [2**i for i in dilations]
return new_dilations
[docs]def channel_normalization(x):
# type: (Layer) -> Layer
"""Normalize a layer to the maximum activation
This keeps a layers values between zero and one.
It helps with relu's unbounded activation
Args:
x: The layer to normalize
Returns:
A maximal normalized layer
"""
max_values = K.max(K.abs(x), 2, keepdims=True) + 1e-5
out = x / max_values
return out
[docs]class ResidualBlock(Layer):
def __init__(
self,
dilation_rate: int,
nb_filters: int,
kernel_size: int,
padding: str,
activation: str = "relu",
dropout_rate: float = 0,
kernel_initializer: str = "he_normal",
use_batch_norm: bool = False,
use_layer_norm: bool = False,
use_separable: bool = False,
**kwargs
):
"""Defines the residual block for the WaveNet TCN
Args:
x: The previous layer in the model
training: boolean indicating whether the layer should behave in training mode or in inference mode
dilation_rate: The dilation power of 2 we are using for this residual block
nb_filters: The number of convolutional filters to use in this block
kernel_size: The size of the convolutional kernel
padding: The padding used in the convolutional layers, 'same' or 'causal'.
activation: The final activation used in o = Activation(x + F(x))
dropout_rate: Float between 0 and 1. Fraction of the input units to drop.
kernel_initializer: Initializer for the kernel weights matrix (Conv1D).
use_batch_norm: Whether to use batch normalization in the residual layers or not.
use_layer_norm: Whether to use layer normalization in the residual layers or not.
use_separable: Whether to use separabel conv1d.
kwargs: Any initializers for Layer class.
"""
self.dilation_rate = dilation_rate
self.nb_filters = nb_filters
self.kernel_size = kernel_size
self.padding = padding
self.activation = activation
self.dropout_rate = dropout_rate
self.use_batch_norm = use_batch_norm
self.use_layer_norm = use_layer_norm
self.use_separable = use_separable
self.kernel_initializer = kernel_initializer
self.layers = []
self.layers_outputs = []
self.shape_match_conv = None
self.res_output_shape = None
self.final_activation = None
super(ResidualBlock, self).__init__(**kwargs)
def _add_and_activate_layer(self, layer):
"""Helper function for building layer
Args:
layer: Appends layer to internal layer list and builds it based on the current output
shape of ResidualBlocK. Updates current output shape.
"""
self.layers.append(layer)
self.layers[-1].build(self.res_output_shape)
self.res_output_shape = self.layers[-1].compute_output_shape(self.res_output_shape)
[docs] def build(self, input_shape):
with K.name_scope(self.name): # name scope used to make sure weights get unique names
self.layers = []
self.res_output_shape = input_shape
ConvThis = SeparableConv1D if self.use_separable else Conv1D
for k in range(1):
name = "conv1D_{}".format(k)
with K.name_scope(name): # name scope used to make sure weights get unique names
self._add_and_activate_layer(
ConvThis(
filters=self.nb_filters,
kernel_size=self.kernel_size,
dilation_rate=self.dilation_rate,
padding=self.padding,
name=name,
kernel_initializer=self.kernel_initializer,
)
)
# with K.name_scope('norm_{}'.format(k)):
# if self.use_batch_norm:
# self._add_and_activate_layer(BatchNormalization())
# elif self.use_layer_norm:
# self._add_and_activate_layer(LayerNormalization())
# self._add_and_activate_layer(Activation('relu'))
# if self.activation == 'norm_relu':
self._add_and_activate_layer(Activation("relu"))
self._add_and_activate_layer(Lambda(channel_normalization))
# else:
# self._add_and_activate_layer(Activation(self.activation))
self._add_and_activate_layer(SpatialDropout1D(rate=self.dropout_rate))
if self.nb_filters != input_shape[-1]:
# 1x1 conv to match the shapes (channel dimension).
name = "matching_conv1D"
with K.name_scope(name):
# make and build this layer separately because it directly uses input_shape
self.shape_match_conv = Conv1D(
filters=self.nb_filters,
kernel_size=1,
padding="same",
name=name,
kernel_initializer=self.kernel_initializer,
)
else:
name = "matching_identity"
self.shape_match_conv = Lambda(lambda x: x, name=name)
with K.name_scope(name):
self.shape_match_conv.build(input_shape)
self.res_output_shape = self.shape_match_conv.compute_output_shape(input_shape)
self.final_activation = Activation(self.activation)
self.final_activation.build(self.res_output_shape) # probably isn't necessary
# this is done to force Keras to add the layers in the list to self._layers
for layer in self.layers:
self.__setattr__(layer.name, layer)
self.__setattr__(self.shape_match_conv.name, self.shape_match_conv)
self.__setattr__(self.final_activation.name, self.final_activation)
super(ResidualBlock, self).build(input_shape) # done to make sure self.built is set True
[docs] def call(self, inputs, training=None):
"""
Returns: A tuple where the first element is the residual model tensor, and the second
is the skip connection tensor.
"""
x = inputs
self.layers_outputs = [x]
for layer in self.layers:
training_flag = "training" in dict(inspect.signature(layer.call).parameters)
x = layer(x, training=training) if training_flag else layer(x)
self.layers_outputs.append(x)
x2 = self.shape_match_conv(inputs)
self.layers_outputs.append(x2)
res_x = layers.add([x2, x])
self.layers_outputs.append(res_x)
res_act_x = self.final_activation(res_x)
self.layers_outputs.append(res_act_x)
return [res_act_x, x]
[docs] def compute_output_shape(self, input_shape):
return [self.res_output_shape, self.res_output_shape]
[docs]class TCN(Layer):
"""Creates a TCN layer.
Input shape:
A tensor of shape (batch_size, timesteps, input_dim).
Args:
nb_filters: The number of filters to use in the convolutional layers. Can be a list.
kernel_size: The size of the kernel to use in each convolutional layer.
dilations: The list of the dilations. Example is: [1, 2, 4, 8, 16, 32, 64].
nb_stacks : The number of stacks of residual blocks to use.
padding: The padding to use in the convolutional layers, 'causal' or 'same'.
use_skip_connections: Boolean. If we want to add skip connections from input to each residual blocK.
return_sequences: Boolean. Whether to return the last output in the output sequence, or the full sequence.
activation: The activation used in the residual blocks o = Activation(x + F(x)).
dropout_rate: Float between 0 and 1. Fraction of the input units to drop.
kernel_initializer: Initializer for the kernel weights matrix (Conv1D).
use_batch_norm: Whether to use batch normalization in the residual layers or not.
use_separable: Whether to use separable convolution in the residual layers or not.
kwargs: Any other arguments for configuring parent class Layer. For example "name=str", Name of the model.
Use unique names when using multiple TCN.
Returns:
A TCN layer.
"""
def __init__(
self,
nb_filters=64,
kernel_size=2,
nb_stacks=1,
dilations=(1, 2, 4, 8, 16, 32),
padding="causal",
use_skip_connections=False,
dropout_rate=0.0,
return_sequences=False,
activation="relu",
kernel_initializer="he_normal",
use_batch_norm=False,
use_layer_norm=False,
use_separable=False,
**kwargs
):
self.return_sequences = return_sequences
self.dropout_rate = dropout_rate
self.use_skip_connections = use_skip_connections
self.dilations = dilations
self.nb_stacks = nb_stacks
self.kernel_size = kernel_size
self.nb_filters = nb_filters
self.activation = activation
self.padding = padding
self.kernel_initializer = kernel_initializer
self.use_batch_norm = use_batch_norm
self.use_layer_norm = use_layer_norm
self.use_separable = use_separable
self.skip_connections = []
self.residual_blocks = []
self.layers_outputs = []
self.build_output_shape = None
self.slicer_layer = None # in case return_sequence=False
self.output_slice_index = None # in case return_sequence=False
self.padding_same_and_time_dim_unknown = False # edge case if padding='same' and time_dim = None
if isinstance(self.nb_filters, list):
assert len(self.nb_filters) == len(self.dilations)
if padding != "causal" and padding != "same":
raise ValueError("Only 'causal' or 'same' padding are compatible for this layer.")
# initialize parent class
super(TCN, self).__init__(**kwargs)
@property
def receptive_field(self):
assert_msg = "The receptive field formula works only with power of two dilations."
assert all([is_power_of_two(i) for i in self.dilations]), assert_msg
return self.kernel_size * self.nb_stacks * self.dilations[-1]
[docs] def build(self, input_shape):
# member to hold current output shape of the layer for building purposes
self.build_output_shape = input_shape
# list to hold all the member ResidualBlocks
self.residual_blocks = []
total_num_blocks = self.nb_stacks * len(self.dilations)
if not self.use_skip_connections:
total_num_blocks += 1 # cheap way to do a false case for below
for s in range(self.nb_stacks):
for i, d in enumerate(self.dilations):
res_block_filters = self.nb_filters[i] if isinstance(self.nb_filters, list) else self.nb_filters
self.residual_blocks.append(
ResidualBlock(
dilation_rate=d,
nb_filters=res_block_filters,
kernel_size=self.kernel_size,
padding=self.padding,
activation=self.activation,
dropout_rate=self.dropout_rate,
use_batch_norm=self.use_batch_norm,
use_layer_norm=self.use_layer_norm,
use_separable=self.use_separable,
kernel_initializer=self.kernel_initializer,
name="residual_block_{}".format(len(self.residual_blocks)),
)
)
# build newest residual block
self.residual_blocks[-1].build(self.build_output_shape)
self.build_output_shape = self.residual_blocks[-1].res_output_shape
# this is done to force keras to add the layers in the list to self._layers
for layer in self.residual_blocks:
self.__setattr__(layer.name, layer)
self.output_slice_index = None
if self.padding == "same":
time = self.build_output_shape.as_list()[1]
if time is not None: # if time dimension is defined. e.g. shape = (bs, 500, input_dim).
self.output_slice_index = int(self.build_output_shape.as_list()[1] / 2)
else:
# It will known at call time. c.f. self.call.
self.padding_same_and_time_dim_unknown = True
else:
self.output_slice_index = -1 # causal case.
self.slicer_layer = Lambda(lambda tt: tt[:, self.output_slice_index, :])
[docs] def compute_output_shape(self, input_shape):
"""
Overridden in case keras uses it somewhere... no idea. Just trying to avoid future errors.
"""
if not self.built:
self.build(input_shape)
if not self.return_sequences:
batch_size = self.build_output_shape[0]
batch_size = batch_size.value if hasattr(batch_size, "value") else batch_size
nb_filters = self.build_output_shape[-1]
return [batch_size, nb_filters]
else:
# Compatibility tensorflow 1.x
return [v.value if hasattr(v, "value") else v for v in self.build_output_shape]
[docs] def call(self, inputs, training=None):
x = inputs
self.layers_outputs = [x]
self.skip_connections = []
for layer in self.residual_blocks:
try:
x, skip_out = layer(x, training=training)
except TypeError: # compatibility with tensorflow 1.x
x, skip_out = layer(K.cast(x, "float32"), training=training)
self.skip_connections.append(skip_out)
self.layers_outputs.append(x)
if self.use_skip_connections:
x = layers.add(self.skip_connections)
self.layers_outputs.append(x)
# ????
x = Activation("relu")(x)
if not self.return_sequences:
# case: time dimension is unknown. e.g. (bs, None, input_dim).
if self.padding_same_and_time_dim_unknown:
self.output_slice_index = K.shape(self.layers_outputs[-1])[1] // 2
x = self.slicer_layer(x)
self.layers_outputs.append(x)
return x
[docs] def get_config(self):
"""
Returns the config of a the layer. This is used for saving and loading from a model
:return: python dictionary with specs to rebuild layer
"""
config = super(TCN, self).get_config()
config["nb_filters"] = self.nb_filters
config["kernel_size"] = self.kernel_size
config["nb_stacks"] = self.nb_stacks
config["dilations"] = self.dilations
config["padding"] = self.padding
config["use_skip_connections"] = self.use_skip_connections
config["dropout_rate"] = self.dropout_rate
config["return_sequences"] = self.return_sequences
config["activation"] = self.activation
config["use_batch_norm"] = self.use_batch_norm
config["use_layer_norm"] = self.use_layer_norm
config["kernel_initializer"] = self.kernel_initializer
return config