mirror of
https://github.com/microsoft/qlib.git
synced 2026-07-09 14:00:55 +08:00
1392 lines
44 KiB
Python
1392 lines
44 KiB
Python
# coding=utf-8
|
|
# Copyright 2020 The Google Research Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Lint as: python3
|
|
"""Temporal Fusion Transformer Model.
|
|
|
|
Contains the full TFT architecture and associated components. Defines functions
|
|
for training, evaluation and prediction using simple Pandas Dataframe inputs.
|
|
"""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import gc
|
|
import json
|
|
import os
|
|
import shutil
|
|
|
|
import data_formatters.base
|
|
import libs.utils as utils
|
|
import numpy as np
|
|
import pandas as pd
|
|
import tensorflow as tf
|
|
|
|
# Layer definitions.
|
|
concat = tf.keras.backend.concatenate
|
|
stack = tf.keras.backend.stack
|
|
K = tf.keras.backend
|
|
Add = tf.keras.layers.Add
|
|
LayerNorm = tf.keras.layers.LayerNormalization
|
|
Dense = tf.keras.layers.Dense
|
|
Multiply = tf.keras.layers.Multiply
|
|
Dropout = tf.keras.layers.Dropout
|
|
Activation = tf.keras.layers.Activation
|
|
Lambda = tf.keras.layers.Lambda
|
|
|
|
# Default input types.
|
|
InputTypes = data_formatters.base.InputTypes
|
|
|
|
|
|
# Layer utility functions.
|
|
def linear_layer(size,
|
|
activation=None,
|
|
use_time_distributed=False,
|
|
use_bias=True):
|
|
"""Returns simple Keras linear layer.
|
|
|
|
Args:
|
|
size: Output size
|
|
activation: Activation function to apply if required
|
|
use_time_distributed: Whether to apply layer across time
|
|
use_bias: Whether bias should be included in layer
|
|
"""
|
|
linear = tf.keras.layers.Dense(size, activation=activation, use_bias=use_bias)
|
|
if use_time_distributed:
|
|
linear = tf.keras.layers.TimeDistributed(linear)
|
|
return linear
|
|
|
|
|
|
def apply_mlp(inputs,
|
|
hidden_size,
|
|
output_size,
|
|
output_activation=None,
|
|
hidden_activation='tanh',
|
|
use_time_distributed=False):
|
|
"""Applies simple feed-forward network to an input.
|
|
|
|
Args:
|
|
inputs: MLP inputs
|
|
hidden_size: Hidden state size
|
|
output_size: Output size of MLP
|
|
output_activation: Activation function to apply on output
|
|
hidden_activation: Activation function to apply on input
|
|
use_time_distributed: Whether to apply across time
|
|
|
|
Returns:
|
|
Tensor for MLP outputs.
|
|
"""
|
|
if use_time_distributed:
|
|
hidden = tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(hidden_size, activation=hidden_activation))(
|
|
inputs)
|
|
return tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(output_size, activation=output_activation))(
|
|
hidden)
|
|
else:
|
|
hidden = tf.keras.layers.Dense(
|
|
hidden_size, activation=hidden_activation)(
|
|
inputs)
|
|
return tf.keras.layers.Dense(
|
|
output_size, activation=output_activation)(
|
|
hidden)
|
|
|
|
|
|
def apply_gating_layer(x,
|
|
hidden_layer_size,
|
|
dropout_rate=None,
|
|
use_time_distributed=True,
|
|
activation=None):
|
|
"""Applies a Gated Linear Unit (GLU) to an input.
|
|
|
|
Args:
|
|
x: Input to gating layer
|
|
hidden_layer_size: Dimension of GLU
|
|
dropout_rate: Dropout rate to apply if any
|
|
use_time_distributed: Whether to apply across time
|
|
activation: Activation function to apply to the linear feature transform if
|
|
necessary
|
|
|
|
Returns:
|
|
Tuple of tensors for: (GLU output, gate)
|
|
"""
|
|
|
|
if dropout_rate is not None:
|
|
x = tf.keras.layers.Dropout(dropout_rate)(x)
|
|
|
|
if use_time_distributed:
|
|
activation_layer = tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(hidden_layer_size, activation=activation))(
|
|
x)
|
|
gated_layer = tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(hidden_layer_size, activation='sigmoid'))(
|
|
x)
|
|
else:
|
|
activation_layer = tf.keras.layers.Dense(
|
|
hidden_layer_size, activation=activation)(
|
|
x)
|
|
gated_layer = tf.keras.layers.Dense(
|
|
hidden_layer_size, activation='sigmoid')(
|
|
x)
|
|
|
|
return tf.keras.layers.Multiply()([activation_layer,
|
|
gated_layer]), gated_layer
|
|
|
|
|
|
def add_and_norm(x_list):
|
|
"""Applies skip connection followed by layer normalisation.
|
|
|
|
Args:
|
|
x_list: List of inputs to sum for skip connection
|
|
|
|
Returns:
|
|
Tensor output from layer.
|
|
"""
|
|
tmp = Add()(x_list)
|
|
tmp = LayerNorm()(tmp)
|
|
return tmp
|
|
|
|
|
|
def gated_residual_network(x,
|
|
hidden_layer_size,
|
|
output_size=None,
|
|
dropout_rate=None,
|
|
use_time_distributed=True,
|
|
additional_context=None,
|
|
return_gate=False):
|
|
"""Applies the gated residual network (GRN) as defined in paper.
|
|
|
|
Args:
|
|
x: Network inputs
|
|
hidden_layer_size: Internal state size
|
|
output_size: Size of output layer
|
|
dropout_rate: Dropout rate if dropout is applied
|
|
use_time_distributed: Whether to apply network across time dimension
|
|
additional_context: Additional context vector to use if relevant
|
|
return_gate: Whether to return GLU gate for diagnostic purposes
|
|
|
|
Returns:
|
|
Tuple of tensors for: (GRN output, GLU gate)
|
|
"""
|
|
|
|
# Setup skip connection
|
|
if output_size is None:
|
|
output_size = hidden_layer_size
|
|
skip = x
|
|
else:
|
|
linear = Dense(output_size)
|
|
if use_time_distributed:
|
|
linear = tf.keras.layers.TimeDistributed(linear)
|
|
skip = linear(x)
|
|
|
|
# Apply feedforward network
|
|
hidden = linear_layer(
|
|
hidden_layer_size,
|
|
activation=None,
|
|
use_time_distributed=use_time_distributed)(
|
|
x)
|
|
if additional_context is not None:
|
|
hidden = hidden + linear_layer(
|
|
hidden_layer_size,
|
|
activation=None,
|
|
use_time_distributed=use_time_distributed,
|
|
use_bias=False)(
|
|
additional_context)
|
|
hidden = tf.keras.layers.Activation('elu')(hidden)
|
|
hidden = linear_layer(
|
|
hidden_layer_size,
|
|
activation=None,
|
|
use_time_distributed=use_time_distributed)(
|
|
hidden)
|
|
|
|
gating_layer, gate = apply_gating_layer(
|
|
hidden,
|
|
output_size,
|
|
dropout_rate=dropout_rate,
|
|
use_time_distributed=use_time_distributed,
|
|
activation=None)
|
|
|
|
if return_gate:
|
|
return add_and_norm([skip, gating_layer]), gate
|
|
else:
|
|
return add_and_norm([skip, gating_layer])
|
|
|
|
|
|
# Attention Components.
|
|
def get_decoder_mask(self_attn_inputs):
|
|
"""Returns causal mask to apply for self-attention layer.
|
|
|
|
Args:
|
|
self_attn_inputs: Inputs to self attention layer to determine mask shape
|
|
"""
|
|
len_s = tf.shape(self_attn_inputs)[1]
|
|
bs = tf.shape(self_attn_inputs)[:1]
|
|
mask = K.cumsum(tf.eye(len_s, batch_shape=bs), 1)
|
|
return mask
|
|
|
|
|
|
class ScaledDotProductAttention():
|
|
"""Defines scaled dot product attention layer.
|
|
|
|
Attributes:
|
|
dropout: Dropout rate to use
|
|
activation: Normalisation function for scaled dot product attention (e.g.
|
|
softmax by default)
|
|
"""
|
|
|
|
def __init__(self, attn_dropout=0.0):
|
|
self.dropout = Dropout(attn_dropout)
|
|
self.activation = Activation('softmax')
|
|
|
|
def __call__(self, q, k, v, mask):
|
|
"""Applies scaled dot product attention.
|
|
|
|
Args:
|
|
q: Queries
|
|
k: Keys
|
|
v: Values
|
|
mask: Masking if required -- sets softmax to very large value
|
|
|
|
Returns:
|
|
Tuple of (layer outputs, attention weights)
|
|
"""
|
|
temper = tf.sqrt(tf.cast(tf.shape(k)[-1], dtype='float32'))
|
|
attn = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[2, 2]) / temper)(
|
|
[q, k]) # shape=(batch, q, k)
|
|
if mask is not None:
|
|
mmask = Lambda(lambda x: (-1e+9) * (1. - K.cast(x, 'float32')))(
|
|
mask) # setting to infinity
|
|
attn = Add()([attn, mmask])
|
|
attn = self.activation(attn)
|
|
attn = self.dropout(attn)
|
|
output = Lambda(lambda x: K.batch_dot(x[0], x[1]))([attn, v])
|
|
return output, attn
|
|
|
|
|
|
class InterpretableMultiHeadAttention():
|
|
"""Defines interpretable multi-head attention layer.
|
|
|
|
Attributes:
|
|
n_head: Number of heads
|
|
d_k: Key/query dimensionality per head
|
|
d_v: Value dimensionality
|
|
dropout: Dropout rate to apply
|
|
qs_layers: List of queries across heads
|
|
ks_layers: List of keys across heads
|
|
vs_layers: List of values across heads
|
|
attention: Scaled dot product attention layer
|
|
w_o: Output weight matrix to project internal state to the original TFT
|
|
state size
|
|
"""
|
|
|
|
def __init__(self, n_head, d_model, dropout):
|
|
"""Initialises layer.
|
|
|
|
Args:
|
|
n_head: Number of heads
|
|
d_model: TFT state dimensionality
|
|
dropout: Dropout discard rate
|
|
"""
|
|
|
|
self.n_head = n_head
|
|
self.d_k = self.d_v = d_k = d_v = d_model // n_head
|
|
self.dropout = dropout
|
|
|
|
self.qs_layers = []
|
|
self.ks_layers = []
|
|
self.vs_layers = []
|
|
|
|
# Use same value layer to facilitate interp
|
|
vs_layer = Dense(d_v, use_bias=False)
|
|
|
|
for _ in range(n_head):
|
|
self.qs_layers.append(Dense(d_k, use_bias=False))
|
|
self.ks_layers.append(Dense(d_k, use_bias=False))
|
|
self.vs_layers.append(vs_layer) # use same vs_layer
|
|
|
|
self.attention = ScaledDotProductAttention()
|
|
self.w_o = Dense(d_model, use_bias=False)
|
|
|
|
def __call__(self, q, k, v, mask=None):
|
|
"""Applies interpretable multihead attention.
|
|
|
|
Using T to denote the number of time steps fed into the transformer.
|
|
|
|
Args:
|
|
q: Query tensor of shape=(?, T, d_model)
|
|
k: Key of shape=(?, T, d_model)
|
|
v: Values of shape=(?, T, d_model)
|
|
mask: Masking if required with shape=(?, T, T)
|
|
|
|
Returns:
|
|
Tuple of (layer outputs, attention weights)
|
|
"""
|
|
n_head = self.n_head
|
|
|
|
heads = []
|
|
attns = []
|
|
for i in range(n_head):
|
|
qs = self.qs_layers[i](q)
|
|
ks = self.ks_layers[i](k)
|
|
vs = self.vs_layers[i](v)
|
|
head, attn = self.attention(qs, ks, vs, mask)
|
|
|
|
head_dropout = Dropout(self.dropout)(head)
|
|
heads.append(head_dropout)
|
|
attns.append(attn)
|
|
head = K.stack(heads) if n_head > 1 else heads[0]
|
|
attn = K.stack(attns)
|
|
|
|
outputs = K.mean(head, axis=0) if n_head > 1 else head
|
|
outputs = self.w_o(outputs)
|
|
outputs = Dropout(self.dropout)(outputs) # output dropout
|
|
|
|
return outputs, attn
|
|
|
|
|
|
class TFTDataCache(object):
|
|
"""Caches data for the TFT."""
|
|
|
|
_data_cache = {}
|
|
|
|
@classmethod
|
|
def update(cls, data, key):
|
|
"""Updates cached data.
|
|
|
|
Args:
|
|
data: Source to update
|
|
key: Key to dictionary location
|
|
"""
|
|
cls._data_cache[key] = data
|
|
|
|
@classmethod
|
|
def get(cls, key):
|
|
"""Returns data stored at key location."""
|
|
return cls._data_cache[key].copy()
|
|
|
|
@classmethod
|
|
def contains(cls, key):
|
|
"""Retuns boolean indicating whether key is present in cache."""
|
|
|
|
return key in cls._data_cache
|
|
|
|
|
|
# TFT model definitions.
|
|
class TemporalFusionTransformer(object):
|
|
"""Defines Temporal Fusion Transformer.
|
|
|
|
Attributes:
|
|
name: Name of model
|
|
time_steps: Total number of input time steps per forecast date (i.e. Width
|
|
of Temporal fusion decoder N)
|
|
input_size: Total number of inputs
|
|
output_size: Total number of outputs
|
|
category_counts: Number of categories per categorical variable
|
|
n_multiprocessing_workers: Number of workers to use for parallel
|
|
computations
|
|
column_definition: List of tuples of (string, DataType, InputType) that
|
|
define each column
|
|
quantiles: Quantiles to forecast for TFT
|
|
use_cudnn: Whether to use Keras CuDNNLSTM or standard LSTM layers
|
|
hidden_layer_size: Internal state size of TFT
|
|
dropout_rate: Dropout discard rate
|
|
max_gradient_norm: Maximum norm for gradient clipping
|
|
learning_rate: Initial learning rate of ADAM optimizer
|
|
minibatch_size: Size of minibatches for training
|
|
num_epochs: Maximum number of epochs for training
|
|
early_stopping_patience: Maximum number of iterations of non-improvement
|
|
before early stopping kicks in
|
|
num_encoder_steps: Size of LSTM encoder -- i.e. number of past time steps
|
|
before forecast date to use
|
|
num_stacks: Number of self-attention layers to apply (default is 1 for basic
|
|
TFT)
|
|
num_heads: Number of heads for interpretable mulit-head attention
|
|
model: Keras model for TFT
|
|
"""
|
|
|
|
def __init__(self, raw_params, use_cudnn=False):
|
|
"""Builds TFT from parameters.
|
|
|
|
Args:
|
|
raw_params: Parameters to define TFT
|
|
use_cudnn: Whether to use CUDNN GPU optimised LSTM
|
|
"""
|
|
|
|
self.name = self.__class__.__name__
|
|
|
|
params = dict(raw_params) # copy locally
|
|
|
|
# Data parameters
|
|
self.time_steps = int(params['total_time_steps'])
|
|
self.input_size = int(params['input_size'])
|
|
self.output_size = int(params['output_size'])
|
|
self.category_counts = json.loads(str(params['category_counts']))
|
|
self.n_multiprocessing_workers = int(params['multiprocessing_workers'])
|
|
|
|
# Relevant indices for TFT
|
|
self._input_obs_loc = json.loads(str(params['input_obs_loc']))
|
|
self._static_input_loc = json.loads(str(params['static_input_loc']))
|
|
self._known_regular_input_idx = json.loads(
|
|
str(params['known_regular_inputs']))
|
|
self._known_categorical_input_idx = json.loads(
|
|
str(params['known_categorical_inputs']))
|
|
|
|
self.column_definition = params['column_definition']
|
|
|
|
# Network params
|
|
self.quantiles = [0.1, 0.5, 0.9]
|
|
self.use_cudnn = use_cudnn # Whether to use GPU optimised LSTM
|
|
self.hidden_layer_size = int(params['hidden_layer_size'])
|
|
self.dropout_rate = float(params['dropout_rate'])
|
|
self.max_gradient_norm = float(params['max_gradient_norm'])
|
|
self.learning_rate = float(params['learning_rate'])
|
|
self.minibatch_size = int(params['minibatch_size'])
|
|
self.num_epochs = int(params['num_epochs'])
|
|
self.early_stopping_patience = int(params['early_stopping_patience'])
|
|
|
|
self.num_encoder_steps = int(params['num_encoder_steps'])
|
|
self.num_stacks = int(params['stack_size'])
|
|
self.num_heads = int(params['num_heads'])
|
|
|
|
# Serialisation options
|
|
self._temp_folder = os.path.join(params['model_folder'], 'tmp')
|
|
self.reset_temp_folder()
|
|
|
|
# Extra components to store Tensorflow nodes for attention computations
|
|
self._input_placeholder = None
|
|
self._attention_components = None
|
|
self._prediction_parts = None
|
|
|
|
print('*** {} params ***'.format(self.name))
|
|
for k in params:
|
|
print('# {} = {}'.format(k, params[k]))
|
|
|
|
# Build model
|
|
self.model = self.build_model()
|
|
|
|
def get_tft_embeddings(self, all_inputs):
|
|
"""Transforms raw inputs to embeddings.
|
|
|
|
Applies linear transformation onto continuous variables and uses embeddings
|
|
for categorical variables.
|
|
|
|
Args:
|
|
all_inputs: Inputs to transform
|
|
|
|
Returns:
|
|
Tensors for transformed inputs.
|
|
"""
|
|
|
|
time_steps = self.time_steps
|
|
|
|
# Sanity checks
|
|
for i in self._known_regular_input_idx:
|
|
if i in self._input_obs_loc:
|
|
raise ValueError('Observation cannot be known a priori!')
|
|
for i in self._input_obs_loc:
|
|
if i in self._static_input_loc:
|
|
raise ValueError('Observation cannot be static!')
|
|
|
|
if all_inputs.get_shape().as_list()[-1] != self.input_size:
|
|
raise ValueError(
|
|
'Illegal number of inputs! Inputs observed={}, expected={}'.format(
|
|
all_inputs.get_shape().as_list()[-1], self.input_size))
|
|
|
|
num_categorical_variables = len(self.category_counts)
|
|
num_regular_variables = self.input_size - num_categorical_variables
|
|
|
|
embedding_sizes = [
|
|
self.hidden_layer_size for i, size in enumerate(self.category_counts)
|
|
]
|
|
|
|
embeddings = []
|
|
for i in range(num_categorical_variables):
|
|
|
|
embedding = tf.keras.Sequential([
|
|
tf.keras.layers.InputLayer([time_steps]),
|
|
tf.keras.layers.Embedding(
|
|
self.category_counts[i],
|
|
embedding_sizes[i],
|
|
input_length=time_steps,
|
|
dtype=tf.float32)
|
|
])
|
|
embeddings.append(embedding)
|
|
|
|
regular_inputs, categorical_inputs \
|
|
= all_inputs[:, :, :num_regular_variables], \
|
|
all_inputs[:, :, num_regular_variables:]
|
|
|
|
embedded_inputs = [
|
|
embeddings[i](categorical_inputs[Ellipsis, i])
|
|
for i in range(num_categorical_variables)
|
|
]
|
|
|
|
# Static inputs
|
|
if self._static_input_loc:
|
|
static_inputs = [tf.keras.layers.Dense(self.hidden_layer_size)(
|
|
regular_inputs[:, 0, i:i + 1]) for i in range(num_regular_variables)
|
|
if i in self._static_input_loc] \
|
|
+ [embedded_inputs[i][:, 0, :]
|
|
for i in range(num_categorical_variables)
|
|
if i + num_regular_variables in self._static_input_loc]
|
|
static_inputs = tf.keras.backend.stack(static_inputs, axis=1)
|
|
|
|
else:
|
|
static_inputs = None
|
|
|
|
def convert_real_to_embedding(x):
|
|
"""Applies linear transformation for time-varying inputs."""
|
|
return tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(self.hidden_layer_size))(
|
|
x)
|
|
|
|
# Targets
|
|
obs_inputs = tf.keras.backend.stack([
|
|
convert_real_to_embedding(regular_inputs[Ellipsis, i:i + 1])
|
|
for i in self._input_obs_loc
|
|
],
|
|
axis=-1)
|
|
|
|
# Observed (a prioir unknown) inputs
|
|
wired_embeddings = []
|
|
for i in range(num_categorical_variables):
|
|
if i not in self._known_categorical_input_idx \
|
|
and i + num_regular_variables not in self._input_obs_loc:
|
|
e = embeddings[i](categorical_inputs[:, :, i])
|
|
wired_embeddings.append(e)
|
|
|
|
unknown_inputs = []
|
|
for i in range(regular_inputs.shape[-1]):
|
|
if i not in self._known_regular_input_idx \
|
|
and i not in self._input_obs_loc:
|
|
e = convert_real_to_embedding(regular_inputs[Ellipsis, i:i + 1])
|
|
unknown_inputs.append(e)
|
|
|
|
if unknown_inputs + wired_embeddings:
|
|
unknown_inputs = tf.keras.backend.stack(
|
|
unknown_inputs + wired_embeddings, axis=-1)
|
|
else:
|
|
unknown_inputs = None
|
|
|
|
# A priori known inputs
|
|
known_regular_inputs = [
|
|
convert_real_to_embedding(regular_inputs[Ellipsis, i:i + 1])
|
|
for i in self._known_regular_input_idx
|
|
if i not in self._static_input_loc
|
|
]
|
|
known_categorical_inputs = [
|
|
embedded_inputs[i]
|
|
for i in self._known_categorical_input_idx
|
|
if i + num_regular_variables not in self._static_input_loc
|
|
]
|
|
|
|
known_combined_layer = tf.keras.backend.stack(
|
|
known_regular_inputs + known_categorical_inputs, axis=-1)
|
|
|
|
return unknown_inputs, known_combined_layer, obs_inputs, static_inputs
|
|
|
|
def _get_single_col_by_type(self, input_type):
|
|
"""Returns name of single column for input type."""
|
|
|
|
return utils.get_single_col_by_input_type(input_type,
|
|
self.column_definition)
|
|
|
|
def training_data_cached(self):
|
|
"""Returns boolean indicating if training data has been cached."""
|
|
|
|
return TFTDataCache.contains('train') and TFTDataCache.contains('valid')
|
|
|
|
def cache_batched_data(self, data, cache_key, num_samples=-1):
|
|
"""Batches and caches data once for using during training.
|
|
|
|
Args:
|
|
data: Data to batch and cache
|
|
cache_key: Key used for cache
|
|
num_samples: Maximum number of samples to extract (-1 to use all data)
|
|
"""
|
|
|
|
if num_samples > 0:
|
|
TFTDataCache.update(
|
|
self._batch_sampled_data(data, max_samples=num_samples), cache_key)
|
|
else:
|
|
TFTDataCache.update(self._batch_data(data), cache_key)
|
|
|
|
print('Cached data "{}" updated'.format(cache_key))
|
|
|
|
def _batch_sampled_data(self, data, max_samples):
|
|
"""Samples segments into a compatible format.
|
|
|
|
Args:
|
|
data: Sources data to sample and batch
|
|
max_samples: Maximum number of samples in batch
|
|
|
|
Returns:
|
|
Dictionary of batched data with the maximum samples specified.
|
|
"""
|
|
|
|
if max_samples < 1:
|
|
raise ValueError(
|
|
'Illegal number of samples specified! samples={}'.format(max_samples))
|
|
|
|
id_col = self._get_single_col_by_type(InputTypes.ID)
|
|
time_col = self._get_single_col_by_type(InputTypes.TIME)
|
|
|
|
data.sort_values(by=[id_col, time_col], inplace=True)
|
|
|
|
print('Getting valid sampling locations.')
|
|
valid_sampling_locations = []
|
|
split_data_map = {}
|
|
for identifier, df in data.groupby(id_col):
|
|
print('Getting locations for {}'.format(identifier))
|
|
num_entries = len(df)
|
|
if num_entries >= self.time_steps:
|
|
valid_sampling_locations += [
|
|
(identifier, self.time_steps + i)
|
|
for i in range(num_entries - self.time_steps + 1)
|
|
]
|
|
split_data_map[identifier] = df
|
|
|
|
inputs = np.zeros((max_samples, self.time_steps, self.input_size))
|
|
outputs = np.zeros((max_samples, self.time_steps, self.output_size))
|
|
time = np.empty((max_samples, self.time_steps, 1), dtype=object)
|
|
identifiers = np.empty((max_samples, self.time_steps, 1), dtype=object)
|
|
|
|
if max_samples > 0 and len(valid_sampling_locations) > max_samples:
|
|
print('Extracting {} samples...'.format(max_samples))
|
|
ranges = [
|
|
valid_sampling_locations[i] for i in np.random.choice(
|
|
len(valid_sampling_locations), max_samples, replace=False)
|
|
]
|
|
else:
|
|
print('Max samples={} exceeds # available segments={}'.format(
|
|
max_samples, len(valid_sampling_locations)))
|
|
ranges = valid_sampling_locations
|
|
|
|
id_col = self._get_single_col_by_type(InputTypes.ID)
|
|
time_col = self._get_single_col_by_type(InputTypes.TIME)
|
|
target_col = self._get_single_col_by_type(InputTypes.TARGET)
|
|
input_cols = [
|
|
tup[0]
|
|
for tup in self.column_definition
|
|
if tup[2] not in {InputTypes.ID, InputTypes.TIME}
|
|
]
|
|
|
|
for i, tup in enumerate(ranges):
|
|
if (i + 1 % 1000) == 0:
|
|
print(i + 1, 'of', max_samples, 'samples done...')
|
|
identifier, start_idx = tup
|
|
sliced = split_data_map[identifier].iloc[start_idx -
|
|
self.time_steps:start_idx]
|
|
inputs[i, :, :] = sliced[input_cols]
|
|
outputs[i, :, :] = sliced[[target_col]]
|
|
time[i, :, 0] = sliced[time_col]
|
|
identifiers[i, :, 0] = sliced[id_col]
|
|
|
|
sampled_data = {
|
|
'inputs': inputs,
|
|
'outputs': outputs[:, self.num_encoder_steps:, :],
|
|
'active_entries': np.ones_like(outputs[:, self.num_encoder_steps:, :]),
|
|
'time': time,
|
|
'identifier': identifiers
|
|
}
|
|
|
|
return sampled_data
|
|
|
|
def _batch_data(self, data):
|
|
"""Batches data for training.
|
|
|
|
Converts raw dataframe from a 2-D tabular format to a batched 3-D array
|
|
to feed into Keras model.
|
|
|
|
Args:
|
|
data: DataFrame to batch
|
|
|
|
Returns:
|
|
Batched Numpy array with shape=(?, self.time_steps, self.input_size)
|
|
"""
|
|
|
|
# Functions.
|
|
def _batch_single_entity(input_data):
|
|
time_steps = len(input_data)
|
|
lags = self.time_steps
|
|
x = input_data.values
|
|
if time_steps >= lags:
|
|
return np.stack(
|
|
[x[i:time_steps - (lags - 1) + i, :] for i in range(lags)], axis=1)
|
|
|
|
else:
|
|
return None
|
|
|
|
id_col = self._get_single_col_by_type(InputTypes.ID)
|
|
time_col = self._get_single_col_by_type(InputTypes.TIME)
|
|
target_col = self._get_single_col_by_type(InputTypes.TARGET)
|
|
input_cols = [
|
|
tup[0]
|
|
for tup in self.column_definition
|
|
if tup[2] not in {InputTypes.ID, InputTypes.TIME}
|
|
]
|
|
|
|
data_map = {}
|
|
for _, sliced in data.groupby(id_col):
|
|
|
|
col_mappings = {
|
|
'identifier': [id_col],
|
|
'time': [time_col],
|
|
'outputs': [target_col],
|
|
'inputs': input_cols
|
|
}
|
|
|
|
for k in col_mappings:
|
|
cols = col_mappings[k]
|
|
arr = _batch_single_entity(sliced[cols].copy())
|
|
|
|
if k not in data_map:
|
|
data_map[k] = [arr]
|
|
else:
|
|
data_map[k].append(arr)
|
|
|
|
# Combine all data
|
|
for k in data_map:
|
|
# Wendi: Avoid returning None when the length is not enough
|
|
data_map[k] = np.concatenate([i for i in data_map[k] if i is not None], axis=0)
|
|
|
|
# Shorten target so we only get decoder steps
|
|
data_map['outputs'] = data_map['outputs'][:, self.num_encoder_steps:, :]
|
|
|
|
active_entries = np.ones_like(data_map['outputs'])
|
|
if 'active_entries' not in data_map:
|
|
data_map['active_entries'] = active_entries
|
|
else:
|
|
data_map['active_entries'].append(active_entries)
|
|
|
|
return data_map
|
|
|
|
def _get_active_locations(self, x):
|
|
"""Formats sample weights for Keras training."""
|
|
return (np.sum(x, axis=-1) > 0.0) * 1.0
|
|
|
|
def _build_base_graph(self):
|
|
"""Returns graph defining layers of the TFT."""
|
|
|
|
# Size definitions.
|
|
time_steps = self.time_steps
|
|
combined_input_size = self.input_size
|
|
encoder_steps = self.num_encoder_steps
|
|
|
|
# Inputs.
|
|
all_inputs = tf.keras.layers.Input(
|
|
shape=(
|
|
time_steps,
|
|
combined_input_size,
|
|
))
|
|
|
|
unknown_inputs, known_combined_layer, obs_inputs, static_inputs \
|
|
= self.get_tft_embeddings(all_inputs)
|
|
|
|
# Isolate known and observed historical inputs.
|
|
if unknown_inputs is not None:
|
|
historical_inputs = concat([
|
|
unknown_inputs[:, :encoder_steps, :],
|
|
known_combined_layer[:, :encoder_steps, :],
|
|
obs_inputs[:, :encoder_steps, :]
|
|
],
|
|
axis=-1)
|
|
else:
|
|
historical_inputs = concat([
|
|
known_combined_layer[:, :encoder_steps, :],
|
|
obs_inputs[:, :encoder_steps, :]
|
|
],
|
|
axis=-1)
|
|
|
|
# Isolate only known future inputs.
|
|
future_inputs = known_combined_layer[:, encoder_steps:, :]
|
|
|
|
def static_combine_and_mask(embedding):
|
|
"""Applies variable selection network to static inputs.
|
|
|
|
Args:
|
|
embedding: Transformed static inputs
|
|
|
|
Returns:
|
|
Tensor output for variable selection network
|
|
"""
|
|
|
|
# Add temporal features
|
|
_, num_static, _ = embedding.get_shape().as_list()
|
|
|
|
flatten = tf.keras.layers.Flatten()(embedding)
|
|
|
|
# Nonlinear transformation with gated residual network.
|
|
mlp_outputs = gated_residual_network(
|
|
flatten,
|
|
self.hidden_layer_size,
|
|
output_size=num_static,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False,
|
|
additional_context=None)
|
|
|
|
sparse_weights = tf.keras.layers.Activation('softmax')(mlp_outputs)
|
|
sparse_weights = K.expand_dims(sparse_weights, axis=-1)
|
|
|
|
trans_emb_list = []
|
|
for i in range(num_static):
|
|
e = gated_residual_network(
|
|
embedding[:, i:i + 1, :],
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False)
|
|
trans_emb_list.append(e)
|
|
|
|
transformed_embedding = concat(trans_emb_list, axis=1)
|
|
|
|
combined = tf.keras.layers.Multiply()(
|
|
[sparse_weights, transformed_embedding])
|
|
|
|
static_vec = K.sum(combined, axis=1)
|
|
|
|
return static_vec, sparse_weights
|
|
|
|
static_encoder, static_weights = static_combine_and_mask(static_inputs)
|
|
|
|
static_context_variable_selection = gated_residual_network(
|
|
static_encoder,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False)
|
|
static_context_enrichment = gated_residual_network(
|
|
static_encoder,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False)
|
|
static_context_state_h = gated_residual_network(
|
|
static_encoder,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False)
|
|
static_context_state_c = gated_residual_network(
|
|
static_encoder,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=False)
|
|
|
|
def lstm_combine_and_mask(embedding):
|
|
"""Apply temporal variable selection networks.
|
|
|
|
Args:
|
|
embedding: Transformed inputs.
|
|
|
|
Returns:
|
|
Processed tensor outputs.
|
|
"""
|
|
|
|
# Add temporal features
|
|
_, time_steps, embedding_dim, num_inputs = embedding.get_shape().as_list()
|
|
|
|
flatten = K.reshape(embedding,
|
|
[-1, time_steps, embedding_dim * num_inputs])
|
|
|
|
expanded_static_context = K.expand_dims(
|
|
static_context_variable_selection, axis=1)
|
|
|
|
# Variable selection weights
|
|
mlp_outputs, static_gate = gated_residual_network(
|
|
flatten,
|
|
self.hidden_layer_size,
|
|
output_size=num_inputs,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=True,
|
|
additional_context=expanded_static_context,
|
|
return_gate=True)
|
|
|
|
sparse_weights = tf.keras.layers.Activation('softmax')(mlp_outputs)
|
|
sparse_weights = tf.expand_dims(sparse_weights, axis=2)
|
|
|
|
# Non-linear Processing & weight application
|
|
trans_emb_list = []
|
|
for i in range(num_inputs):
|
|
grn_output = gated_residual_network(
|
|
embedding[Ellipsis, i],
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=True)
|
|
trans_emb_list.append(grn_output)
|
|
|
|
transformed_embedding = stack(trans_emb_list, axis=-1)
|
|
|
|
combined = tf.keras.layers.Multiply()(
|
|
[sparse_weights, transformed_embedding])
|
|
temporal_ctx = K.sum(combined, axis=-1)
|
|
|
|
return temporal_ctx, sparse_weights, static_gate
|
|
|
|
historical_features, historical_flags, _ = lstm_combine_and_mask(
|
|
historical_inputs)
|
|
future_features, future_flags, _ = lstm_combine_and_mask(future_inputs)
|
|
|
|
# LSTM layer
|
|
def get_lstm(return_state):
|
|
"""Returns LSTM cell initialized with default parameters."""
|
|
if self.use_cudnn:
|
|
lstm = tf.keras.layers.CuDNNLSTM(
|
|
self.hidden_layer_size,
|
|
return_sequences=True,
|
|
return_state=return_state,
|
|
stateful=False,
|
|
)
|
|
else:
|
|
lstm = tf.keras.layers.LSTM(
|
|
self.hidden_layer_size,
|
|
return_sequences=True,
|
|
return_state=return_state,
|
|
stateful=False,
|
|
# Additional params to ensure LSTM matches CuDNN, See TF 2.0 :
|
|
# (https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM)
|
|
activation='tanh',
|
|
recurrent_activation='sigmoid',
|
|
recurrent_dropout=0,
|
|
unroll=False,
|
|
use_bias=True)
|
|
return lstm
|
|
|
|
history_lstm, state_h, state_c \
|
|
= get_lstm(return_state=True)(historical_features,
|
|
initial_state=[static_context_state_h,
|
|
static_context_state_c])
|
|
|
|
future_lstm = get_lstm(return_state=False)(
|
|
future_features, initial_state=[state_h, state_c])
|
|
|
|
lstm_layer = concat([history_lstm, future_lstm], axis=1)
|
|
|
|
# Apply gated skip connection
|
|
input_embeddings = concat([historical_features, future_features], axis=1)
|
|
|
|
lstm_layer, _ = apply_gating_layer(
|
|
lstm_layer, self.hidden_layer_size, self.dropout_rate, activation=None)
|
|
temporal_feature_layer = add_and_norm([lstm_layer, input_embeddings])
|
|
|
|
# Static enrichment layers
|
|
expanded_static_context = K.expand_dims(static_context_enrichment, axis=1)
|
|
enriched, _ = gated_residual_network(
|
|
temporal_feature_layer,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=True,
|
|
additional_context=expanded_static_context,
|
|
return_gate=True)
|
|
|
|
# Decoder self attention
|
|
self_attn_layer = InterpretableMultiHeadAttention(
|
|
self.num_heads, self.hidden_layer_size, dropout=self.dropout_rate)
|
|
|
|
mask = get_decoder_mask(enriched)
|
|
x, self_att \
|
|
= self_attn_layer(enriched, enriched, enriched,
|
|
mask=mask)
|
|
|
|
x, _ = apply_gating_layer(
|
|
x,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
activation=None)
|
|
x = add_and_norm([x, enriched])
|
|
|
|
# Nonlinear processing on outputs
|
|
decoder = gated_residual_network(
|
|
x,
|
|
self.hidden_layer_size,
|
|
dropout_rate=self.dropout_rate,
|
|
use_time_distributed=True)
|
|
|
|
# Final skip connection
|
|
decoder, _ = apply_gating_layer(
|
|
decoder, self.hidden_layer_size, activation=None)
|
|
transformer_layer = add_and_norm([decoder, temporal_feature_layer])
|
|
|
|
# Attention components for explainability
|
|
attention_components = {
|
|
# Temporal attention weights
|
|
'decoder_self_attn': self_att,
|
|
# Static variable selection weights
|
|
'static_flags': static_weights[Ellipsis, 0],
|
|
# Variable selection weights of past inputs
|
|
'historical_flags': historical_flags[Ellipsis, 0, :],
|
|
# Variable selection weights of future inputs
|
|
'future_flags': future_flags[Ellipsis, 0, :]
|
|
}
|
|
|
|
return transformer_layer, all_inputs, attention_components
|
|
|
|
def build_model(self):
|
|
"""Build model and defines training losses.
|
|
|
|
Returns:
|
|
Fully defined Keras model.
|
|
"""
|
|
|
|
with tf.variable_scope(self.name):
|
|
|
|
transformer_layer, all_inputs, attention_components \
|
|
= self._build_base_graph()
|
|
|
|
outputs = tf.keras.layers.TimeDistributed(
|
|
tf.keras.layers.Dense(self.output_size * len(self.quantiles))) \
|
|
(transformer_layer[Ellipsis, self.num_encoder_steps:, :])
|
|
|
|
self._attention_components = attention_components
|
|
|
|
adam = tf.keras.optimizers.Adam(
|
|
lr=self.learning_rate, clipnorm=self.max_gradient_norm)
|
|
|
|
model = tf.keras.Model(inputs=all_inputs, outputs=outputs)
|
|
|
|
print(model.summary())
|
|
|
|
valid_quantiles = self.quantiles
|
|
output_size = self.output_size
|
|
|
|
class QuantileLossCalculator(object):
|
|
"""Computes the combined quantile loss for prespecified quantiles.
|
|
|
|
Attributes:
|
|
quantiles: Quantiles to compute losses
|
|
"""
|
|
|
|
def __init__(self, quantiles):
|
|
"""Initializes computer with quantiles for loss calculations.
|
|
|
|
Args:
|
|
quantiles: Quantiles to use for computations.
|
|
"""
|
|
self.quantiles = quantiles
|
|
|
|
def quantile_loss(self, a, b):
|
|
"""Returns quantile loss for specified quantiles.
|
|
|
|
Args:
|
|
a: Targets
|
|
b: Predictions
|
|
"""
|
|
quantiles_used = set(self.quantiles)
|
|
|
|
loss = 0.
|
|
for i, quantile in enumerate(valid_quantiles):
|
|
if quantile in quantiles_used:
|
|
loss += utils.tensorflow_quantile_loss(
|
|
a[Ellipsis, output_size * i:output_size * (i + 1)],
|
|
b[Ellipsis, output_size * i:output_size * (i + 1)], quantile)
|
|
return loss
|
|
|
|
quantile_loss = QuantileLossCalculator(valid_quantiles).quantile_loss
|
|
|
|
model.compile(
|
|
loss=quantile_loss, optimizer=adam, sample_weight_mode='temporal')
|
|
|
|
self._input_placeholder = all_inputs
|
|
|
|
return model
|
|
|
|
def fit(self, train_df=None, valid_df=None):
|
|
"""Fits deep neural network for given training and validation data.
|
|
|
|
Args:
|
|
train_df: DataFrame for training data
|
|
valid_df: DataFrame for validation data
|
|
"""
|
|
|
|
print('*** Fitting {} ***'.format(self.name))
|
|
|
|
# Add relevant callbacks
|
|
callbacks = [
|
|
tf.keras.callbacks.EarlyStopping(
|
|
monitor='val_loss',
|
|
patience=self.early_stopping_patience,
|
|
min_delta=1e-4),
|
|
tf.keras.callbacks.ModelCheckpoint(
|
|
filepath=self.get_keras_saved_path(self._temp_folder),
|
|
monitor='val_loss',
|
|
save_best_only=True,
|
|
save_weights_only=True),
|
|
tf.keras.callbacks.TerminateOnNaN()
|
|
]
|
|
|
|
print('Getting batched_data')
|
|
if train_df is None:
|
|
print('Using cached training data')
|
|
train_data = TFTDataCache.get('train')
|
|
else:
|
|
train_data = self._batch_data(train_df)
|
|
|
|
if valid_df is None:
|
|
print('Using cached validation data')
|
|
valid_data = TFTDataCache.get('valid')
|
|
else:
|
|
valid_data = self._batch_data(valid_df)
|
|
|
|
print('Using keras standard fit')
|
|
|
|
def _unpack(data):
|
|
return data['inputs'], data['outputs'], \
|
|
self._get_active_locations(data['active_entries'])
|
|
|
|
# Unpack without sample weights
|
|
data, labels, active_flags = _unpack(train_data)
|
|
val_data, val_labels, val_flags = _unpack(valid_data)
|
|
|
|
all_callbacks = callbacks
|
|
|
|
self.model.fit(
|
|
x=data,
|
|
y=np.concatenate([labels, labels, labels], axis=-1),
|
|
sample_weight=active_flags,
|
|
epochs=self.num_epochs,
|
|
batch_size=self.minibatch_size,
|
|
validation_data=(val_data,
|
|
np.concatenate([val_labels, val_labels, val_labels],
|
|
axis=-1), val_flags),
|
|
callbacks=all_callbacks,
|
|
shuffle=True,
|
|
use_multiprocessing=True,
|
|
workers=self.n_multiprocessing_workers)
|
|
|
|
# Load best checkpoint again
|
|
tmp_checkpont = self.get_keras_saved_path(self._temp_folder)
|
|
if os.path.exists(tmp_checkpont):
|
|
self.load(
|
|
self._temp_folder,
|
|
use_keras_loadings=True)
|
|
|
|
else:
|
|
print('Cannot load from {}, skipping ...'.format(self._temp_folder))
|
|
|
|
def evaluate(self, data=None, eval_metric='loss'):
|
|
"""Applies evaluation metric to the training data.
|
|
|
|
Args:
|
|
data: Dataframe for evaluation
|
|
eval_metric: Evaluation metic to return, based on model definition.
|
|
|
|
Returns:
|
|
Computed evaluation loss.
|
|
"""
|
|
|
|
if data is None:
|
|
print('Using cached validation data')
|
|
raw_data = TFTDataCache.get('valid')
|
|
else:
|
|
raw_data = self._batch_data(data)
|
|
|
|
inputs = raw_data['inputs']
|
|
outputs = raw_data['outputs']
|
|
active_entries = self._get_active_locations(raw_data['active_entries'])
|
|
|
|
metric_values = self.model.evaluate(
|
|
x=inputs,
|
|
y=np.concatenate([outputs, outputs, outputs], axis=-1),
|
|
sample_weight=active_entries,
|
|
workers=16,
|
|
use_multiprocessing=True)
|
|
|
|
metrics = pd.Series(metric_values, self.model.metrics_names)
|
|
|
|
return metrics[eval_metric]
|
|
|
|
def predict(self, df, return_targets=False):
|
|
"""Computes predictions for a given input dataset.
|
|
|
|
Args:
|
|
df: Input dataframe
|
|
return_targets: Whether to also return outputs aligned with predictions to
|
|
faciliate evaluation
|
|
|
|
Returns:
|
|
Input dataframe or tuple of (input dataframe, algined output dataframe).
|
|
"""
|
|
|
|
data = self._batch_data(df)
|
|
|
|
inputs = data['inputs']
|
|
time = data['time']
|
|
identifier = data['identifier']
|
|
outputs = data['outputs']
|
|
|
|
combined = self.model.predict(
|
|
inputs,
|
|
workers=16,
|
|
use_multiprocessing=True,
|
|
batch_size=self.minibatch_size)
|
|
|
|
# Format output_csv
|
|
if self.output_size != 1:
|
|
raise NotImplementedError('Current version only supports 1D targets!')
|
|
|
|
def format_outputs(prediction):
|
|
"""Returns formatted dataframes for prediction."""
|
|
|
|
flat_prediction = pd.DataFrame(
|
|
prediction[:, :, 0],
|
|
columns=[
|
|
't+{}'.format(i)
|
|
for i in range(self.time_steps - self.num_encoder_steps)
|
|
])
|
|
cols = list(flat_prediction.columns)
|
|
flat_prediction['forecast_time'] = time[:, self.num_encoder_steps - 1, 0]
|
|
flat_prediction['identifier'] = identifier[:, 0, 0]
|
|
|
|
# Arrange in order
|
|
return flat_prediction[['forecast_time', 'identifier'] + cols]
|
|
|
|
# Extract predictions for each quantile into different entries
|
|
process_map = {
|
|
'p{}'.format(int(q * 100)):
|
|
combined[Ellipsis, i * self.output_size:(i + 1) * self.output_size]
|
|
for i, q in enumerate(self.quantiles)
|
|
}
|
|
|
|
if return_targets:
|
|
# Add targets if relevant
|
|
process_map['targets'] = outputs
|
|
|
|
return {k: format_outputs(process_map[k]) for k in process_map}
|
|
|
|
def get_attention(self, df):
|
|
"""Computes TFT attention weights for a given dataset.
|
|
|
|
Args:
|
|
df: Input dataframe
|
|
|
|
Returns:
|
|
Dictionary of numpy arrays for temporal attention weights and variable
|
|
selection weights, along with their identifiers and time indices
|
|
"""
|
|
|
|
data = self._batch_data(df)
|
|
inputs = data['inputs']
|
|
identifiers = data['identifier']
|
|
time = data['time']
|
|
|
|
def get_batch_attention_weights(input_batch):
|
|
"""Returns weights for a given minibatch of data."""
|
|
input_placeholder = self._input_placeholder
|
|
attention_weights = {}
|
|
for k in self._attention_components:
|
|
attention_weight = tf.keras.backend.get_session().run(
|
|
self._attention_components[k],
|
|
{input_placeholder: input_batch.astype(np.float32)})
|
|
attention_weights[k] = attention_weight
|
|
return attention_weights
|
|
|
|
# Compute number of batches
|
|
batch_size = self.minibatch_size
|
|
n = inputs.shape[0]
|
|
num_batches = n // batch_size
|
|
if n - (num_batches * batch_size) > 0:
|
|
num_batches += 1
|
|
|
|
# Split up inputs into batches
|
|
batched_inputs = [
|
|
inputs[i * batch_size:(i + 1) * batch_size, Ellipsis]
|
|
for i in range(num_batches)
|
|
]
|
|
|
|
# Get attention weights, while avoiding large memory increases
|
|
attention_by_batch = [
|
|
get_batch_attention_weights(batch) for batch in batched_inputs
|
|
]
|
|
attention_weights = {}
|
|
for k in self._attention_components:
|
|
attention_weights[k] = []
|
|
for batch_weights in attention_by_batch:
|
|
attention_weights[k].append(batch_weights[k])
|
|
|
|
if len(attention_weights[k][0].shape) == 4:
|
|
tmp = np.concatenate(attention_weights[k], axis=1)
|
|
else:
|
|
tmp = np.concatenate(attention_weights[k], axis=0)
|
|
|
|
del attention_weights[k]
|
|
gc.collect()
|
|
attention_weights[k] = tmp
|
|
|
|
attention_weights['identifiers'] = identifiers[:, 0, 0]
|
|
attention_weights['time'] = time[:, :, 0]
|
|
|
|
return attention_weights
|
|
|
|
# Serialisation.
|
|
def reset_temp_folder(self):
|
|
"""Deletes and recreates folder with temporary Keras training outputs."""
|
|
print('Resetting temp folder...')
|
|
utils.create_folder_if_not_exist(self._temp_folder)
|
|
shutil.rmtree(self._temp_folder)
|
|
os.makedirs(self._temp_folder)
|
|
|
|
def get_keras_saved_path(self, model_folder):
|
|
"""Returns path to keras checkpoint."""
|
|
return os.path.join(model_folder, '{}.check'.format(self.name))
|
|
|
|
def save(self, model_folder):
|
|
"""Saves optimal TFT weights.
|
|
|
|
Args:
|
|
model_folder: Location to serialze model.
|
|
"""
|
|
# Allows for direct serialisation of tensorflow variables to avoid spurious
|
|
# issue with Keras that leads to different performance evaluation results
|
|
# when model is reloaded (https://github.com/keras-team/keras/issues/4875).
|
|
|
|
utils.save(
|
|
tf.keras.backend.get_session(),
|
|
model_folder,
|
|
cp_name=self.name,
|
|
scope=self.name)
|
|
|
|
def load(self, model_folder, use_keras_loadings=False):
|
|
"""Loads TFT weights.
|
|
|
|
Args:
|
|
model_folder: Folder containing serialized models.
|
|
use_keras_loadings: Whether to load from Keras checkpoint.
|
|
|
|
Returns:
|
|
|
|
"""
|
|
if use_keras_loadings:
|
|
# Loads temporary Keras model saved during training.
|
|
serialisation_path = self.get_keras_saved_path(model_folder)
|
|
print('Loading model from {}'.format(serialisation_path))
|
|
self.model.load_weights(serialisation_path)
|
|
else:
|
|
# Loads tensorflow graph for optimal models.
|
|
utils.load(
|
|
tf.keras.backend.get_session(),
|
|
model_folder,
|
|
cp_name=self.name,
|
|
scope=self.name)
|
|
|
|
@classmethod
|
|
def get_hyperparm_choices(cls):
|
|
"""Returns hyperparameter ranges for random search."""
|
|
return {
|
|
'dropout_rate': [0.1, 0.2, 0.3, 0.4, 0.5, 0.7, 0.9],
|
|
'hidden_layer_size': [10, 20, 40, 80, 160, 240, 320],
|
|
'minibatch_size': [64, 128, 256],
|
|
'learning_rate': [1e-4, 1e-3, 1e-2],
|
|
'max_gradient_norm': [0.01, 1.0, 100.0],
|
|
'num_heads': [1, 4],
|
|
'stack_size': [1],
|
|
}
|