GeoMAN.py

# -*- coding: utf-8 -*-
"""
Yuxuan Liang, Songyu Ke, Junbo Zhang, Xiuwen Yi, Yu Zheng
GeoMAN: Multi-level Attention Networks for Geo-sensory Time Series Prediction
27th International Joint Conference on Artificial Intelligence (IJCAI 2018)
-------------------------------------------------
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
from base_model import BaseModel
from tensorflow.contrib import rnn
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.util import nest
from six.moves import xrange
from six.moves import zip
from utils import Linear

def input_transform(local_inputs,
                    global_inputs,
                    external_inputs,
                    local_attn_states,
                    global_attn_states,
                    labels):
    n_steps_decoder = labels.get_shape()[1].value
    n_output_decoder = labels.get_shape()[2].value
    n_sensors = global_inputs.get_shape()[2].value
    n_steps_encoder = local_inputs.get_shape()[1].value
    n_input_encoder = local_inputs.get_shape()[2].value
    n_external_input = external_inputs.get_shape()[2].value

    # a tuple composed of the local and global attention states
    encoder_attention_states = (local_attn_states,
                                global_attn_states)

    # transform the inputs from local and global view into encoder_inputs
    _local_inputs = tf.transpose(local_inputs, [1, 0, 2])
    _local_inputs = tf.reshape(_local_inputs, [-1, n_input_encoder])
    _local_inputs = tf.split(_local_inputs, n_steps_encoder, 0)
    _global_inputs = tf.transpose(global_inputs, [1, 0, 2])
    _global_inputs = tf.reshape(_global_inputs, [-1, n_sensors])
    _global_inputs = tf.split(_global_inputs, n_steps_encoder, 0)
    encoder_inputs = (_local_inputs, _global_inputs)

    # transform the variables into lists as the input of different function
    _labels = tf.transpose(labels, [1, 0, 2])
    _labels = tf.reshape(_labels, [-1, n_output_decoder])
    _labels = tf.split(_labels, n_steps_decoder, 0)
    _external_inputs = tf.transpose(external_inputs, [1, 0, 2])
    _external_inputs = tf.reshape(_external_inputs, [-1, n_external_input])
    _external_inputs = tf.split(_external_inputs, n_steps_decoder, 0)

    # not useful when the loop function is employed
    decoder_inputs = [tf.zeros_like(
        _labels[0], dtype=tf.float32, name="GO")] + _labels[:-1]
    return encoder_attention_states, encoder_inputs, _labels, _external_inputs, decoder_inputs

class GeoMAN(BaseModel):
    def __init__(self, hps, mode='train'):
        super(GeoMAN, self).__init__(hps, mode)
        preds = self.mod_fn()
        self.phs['preds'] = preds
        self.phs['loss'] = self.get_loss()  # see at eq.[11]
        tf.add_to_collection('loss', self.phs['loss'])
        self.phs['train_op'] = self.train_op()
        self.phs['summary'] = self.summary()

    def spatial_attention(self,
                          encoder_inputs,
                          attention_states,
                          cell,
                          s_attn_flag=2,
                          output_size=None,
                          dtype=dtypes.float32,
                          scope=None):
        """ Spaital attention in GeoMAN
            Args:
                encoder_inputs: A tuple consisting of
                  1) local_inputs: the inputs of local spatial attention, i.e., a list of 2D tensors with the shape of
                    [batch_size, n_inputs_encoder]
                  2) global_inputs: the inputs of local spatial attention, i.e., a list of 2D tensors with the shape of
                    [batch_size, n_sensors]
                attention_states: A tuple consisting of
                  1) local_attention_states: 3D tensor [batch_size, n_input_encoder, n_steps_encoder]
                  2) global_attention_states: 4D tensor [batch_size, n_sensors, n_input_encoder, n_steps_encoder]
                cell: core_rnn_cell.RNNCell defining the cell function and size.
                s_attn_flag: 0: only local. 1: only global. 2: local + global.
                output_size: Size of the output vectors; if None, we use cell.output_size.
                loop_function: the loop function we use.
                dtype: The dtype to use for the RNN initial state (default: tf.float32).
                scope: VariableScope for the created subgraph; default: "spatial_attention".
            Return:
                A tuple of the form (outputs, state), where:
            Raises:
                ValueError: when num_heads is not positive, there are no inputs, shapes
                  of attention_states are not set, or input size cannot be inferred from the
                  input.
        """
        # check inputs
        if not encoder_inputs:
            raise ValueError(
                "Must provide at least 1 input to attention encoder.")
        local_inputs = encoder_inputs[0]
        global_inputs = encoder_inputs[1]
        if output_size is None:
            output_size = cell.output_size
        local_attention_states = attention_states[0]
        global_attention_states = attention_states[1]
        if not local_attention_states.get_shape()[1:2].is_fully_defined():
            raise ValueError("Shape[1] and [2] of local_attention_states must be known: %s"
                             % local_attention_states.get_shape())
        if not global_attention_states.get_shape()[1:3].is_fully_defined():
            raise ValueError("Shape[1] to [3] of global_attention_states must be known: %s"
                             % global_attention_states.get_shape())
        batch_size = array_ops.shape(local_inputs[0])[0]

        # decide whether to use local/global attention
        # s_attn_flag: 0: only local. 1: only global. 2: local + global
        local_flag = True
        global_flag = True
        if s_attn_flag == 0:
            global_flag = False
        elif s_attn_flag == 1:
            local_flag = False

        with vs.variable_scope(scope or "spatial_attn"):
            if local_flag:
                # implement of local spatial attention
                with tf.variable_scope('local_spatial_attn'):
                    local_attn_length = local_attention_states.get_shape()[
                        1].value  # n_input_encoder
                    local_attn_size = local_attention_states.get_shape()[
                        2].value  # n_steps_encoder

                    # A trick: to calculate U_l * x^{i,k} by a 1-by-1 convolution
                    # refer: https://github.com/tensorflow/tensorflow/blob/r1.8/tensorflow/contrib/legacy_seq2seq/python/ops/seq2seq.py
                    # See at eq.[1] in the paper
                    local_hidden = array_ops.reshape(
                        local_attention_states, [-1, local_attn_length, 1, local_attn_size])
                    # Size of query vectors for attention.
                    local_attention_vec_size = local_attn_size
                    local_u = vs.get_variable("AttnUl",
                                              [1, 1, local_attn_size, local_attention_vec_size])  # U_l
                    local_hidden_features = nn_ops.conv2d(
                        local_hidden, local_u, [1, 1, 1, 1], "SAME")  # U_l * x^{i,k}
                    local_v = vs.get_variable(
                        "AttnVl", [local_attention_vec_size])  # v_l
                    batch_attn_size = array_ops.stack(
                        [batch_size, local_attn_length])
                    local_attn = array_ops.zeros(batch_attn_size, dtype=dtype)

                    def local_attention(query):
                        """Put attention masks on local_hidden using local_hidden_features and query."""
                        # If the query is a tuple (when stacked RNN/LSTM), flatten it
                        if nest.is_sequence(query):
                            query_list = nest.flatten(query)
                            for q in query_list:
                                ndims = q.get_shape().ndims
                                if ndims:
                                    assert ndims == 2
                            query = array_ops.concat(query_list, 1)
                        with tf.variable_scope("AttnWl"):
                            # linear map
                            y = Linear(query, local_attention_vec_size, True)
                            y = array_ops.reshape(
                                y, [-1, 1, 1, local_attention_vec_size])
                            # Attention mask is a softmax of v_l^{\top} * tanh(...)
                            s = math_ops.reduce_sum(
                                local_v * math_ops.tanh(local_hidden_features + y), [2, 3])
                            # Now calculate the attention-weighted vector, i.e., alpha in eq.[2]
                            a = nn_ops.softmax(s)
                        return a

            # implement of global spatial attention
            if global_flag:
                with tf.variable_scope('global_spatial_attn'):
                    global_attn_length = global_attention_states.get_shape()[
                        1].value  # n_sensor
                    global_n_input = global_attention_states.get_shape()[
                        2].value  # n_input_dim
                    global_attn_size = global_attention_states.get_shape()[
                        3].value  # n_steps_encoder, T

                    # This implement is a little bit different from the paper at IJCAI-18.
                    # See at eq.[3.5] (we have no place to label the equation T-T) in the paper
                    # Note that our global input X^l include the target series y^l at l-th sensor,
                    # we calculate W'_g * X^l * u_g by a convolution and omit U_g * y^l for simplicity.
                    # You can easily add U_g * y^l here, where y^l is the first column of local inputs.
                    global_hidden = array_ops.reshape(global_attention_states,
                                                      [-1, global_attn_length, global_n_input, global_attn_size])
                    # Size of query vectors for attention.
                    global_attention_vec_size = global_attn_size
                    global_k = vs.get_variable("AttnW_and_u",
                                               [1, global_n_input, global_attn_size,
                                                global_attention_vec_size])
                    # global_hidden_features with the shape (batch_size, global_attn_length, 1, global_attn_size)
                    global_hidden_features = nn_ops.conv2d(
                        global_hidden, global_k, [1, 1, 1, 1], "SAME")
                    global_v = vs.get_variable(
                        "AttnVg", [global_attention_vec_size])
                    batch_attn_size = array_ops.stack(
                        [batch_size, global_attn_length])
                    global_attn = array_ops.zeros(batch_attn_size, dtype=dtype)

                    def global_attention(query):
                        """Put attention masks on global_hidden using global_hidden_features and query."""
                        # If the query is a tuple (when stacked RNN/LSTM), flatten it
                        if nest.is_sequence(query):
                            query_list = nest.flatten(query)
                            for q in query_list:  # Check that ndims == 2 if specified.
                                ndims = q.get_shape().ndims
                                if ndims:
                                    assert ndims == 2
                            query = array_ops.concat(query_list, 1)
                        with tf.variable_scope("AttnWg"):
                            # linear map
                            y = Linear(query, global_attention_vec_size, True)
                            y = array_ops.reshape(
                                y, [-1, 1, 1, global_attention_vec_size])
                            # Attention mask is a softmax of v_g^{\top} * tanh(...)
                            s = math_ops.reduce_sum(
                                global_v * math_ops.tanh(global_hidden_features + y), [2, 3])
                            # Sometimes it's not easy to find a measurement to denote similarity between sensors,
                            # here we omit such prior knowledge in eq.[4].
                            # You can use "a = nn_ops.softmax((1-lambda)*s + lambda*sim)" to encode similarity info,
                            # where:
                            #     sim: a vector with length n_sensors, describing the sim between the target sensor and the others
                            #     lambda: a trade-off.
                            a = nn_ops.softmax(s)
                            # a = nn_ops.softmax((1 - lambda) * s + lambda * sim)
                        return a

            # how to get the initial_state
            initial_state_size = array_ops.stack([batch_size, output_size])
            initial_state_one = [array_ops.zeros(
                initial_state_size, dtype=dtype) for _ in xrange(2)]
            initial_state = [
                initial_state_one for _ in range(len(cell._cells))]
            state = initial_state

            outputs = []
            attn_weights = []
            i = 0
            # i is the index of the which time step
            # local_inp is numpy.array and the shape of local_inp is (batch_size, n_feature)
            for local_inp, global_inp in zip(local_inputs, global_inputs):
                if i > 0:
                    vs.get_variable_scope().reuse_variables()
                input_size = local_inp.get_shape().with_rank(2)[1]
                if input_size.value is None:
                    raise ValueError(
                        "Could not infer input size from input: %s" % local_inp.name)

                if local_flag and global_flag:
                    # multiply attention weights with the original input
                    local_x = local_attn * local_inp
                    global_x = global_attn * global_inp
                    # Run the BasicLSTM with the newly input
                    cell_output, state = cell(
                        tf.concat([local_x, global_x], axis=1), state)
                    # Run the attention mechanism.
                    with tf.variable_scope('local_spatial_attn'):
                        local_attn = local_attention(state)
                    with tf.variable_scope('global_spatial_attn'):
                        global_attn = global_attention(state)
                    attn_weights.append((local_attn, global_attn))
                elif local_flag:
                    local_x = local_attn * local_inp
                    cell_output, state = cell(local_x, state)
                    with tf.variable_scope('local_spatial_attn'):
                        local_attn = local_attention(state)
                    attn_weights.append(local_attn)
                elif global_flag:
                    global_x = global_attn * global_inp
                    cell_output, state = cell(global_x, state)
                    with tf.variable_scope('global_spatial_attn'):
                        global_attn = global_attention(state)
                    attn_weights.append(global_attn)
                # Attention output projection
                with vs.variable_scope("AttnOutputProjection"):
                    output = cell_output
                outputs.append(output)
                i += 1
        return outputs, state, attn_weights

    def temporal_attention(self,
                           decoder_inputs,
                           external_inputs,
                           initial_state,
                           attention_states,
                           cell,
                           output_size=None,
                           loop_function=None,
                           dtype=tf.float32,
                           scope=None,
                           initial_state_attention=False,
                           external_flag=True):
        """ Temporal attention in GeoMAN
        Args:
            decoder_inputs: A list (length: n_steps_decoder) of 2D Tensors [batch_size, n_input_decoder].
            external_inputs: A list (length: n_steps_decoder) of 2D Tensors [batch_size, n_external_input].
            initial_state: 2D Tensor [batch_size, cell.state_size].
            attention_states: 3D Tensor [batch_size, n_step_encoder, n_hidden_encoder].
            cell: core_rnn_cell.RNNCell defining the cell function and size.
            output_size: Size of the output vectors; if None, we use cell.output_size.
            loop_function: the loop function we use. 
            dtype: The dtype to use for the RNN initial state (default: tf.float32).
            scope: VariableScope for the created subgraph; default: "tempotal_attention".
            initial_state_attention: If False (default), initial attentions are zero.
            external_flag: whether to use external factors

        Return:
            A tuple of the form (outputs, state), where:
                outputs: A list of the same length as the inputs of decoder of 2D Tensors of
                         shape [batch_size x output_size]
                state: The state of each decoder cell the final time-step.
        """
        # check inputs
        if not decoder_inputs:
            raise ValueError(
                "Must provide at least 1 input to attention decoder.")
        if not external_inputs:
            raise ValueError(
                "Must provide at least 1 ext_input to attention decoder.")
        if attention_states.get_shape()[2].value is None:
            raise ValueError("Shape[2] of attention_states must be known: %s" %
                             attention_states.get_shape())
        if output_size is None:
            output_size = cell.output_size

        # implement of temporal attention
        with vs.variable_scope(
                scope or "temporal_attn", dtype=dtype) as scope:
            dtype = scope.dtype
            # Needed for reshaping.
            batch_size = array_ops.shape(decoder_inputs[0])[0]
            attn_length = attention_states.get_shape()[1].value
            if attn_length is None:
                attn_length = array_ops.shape(attention_states)[1]
            attn_size = attention_states.get_shape()[2].value

            # A trick: to calculate W_d * h_o by a 1-by-1 convolution
            # See at eq.[6] in the paper
            hidden = array_ops.reshape(attention_states,
                                       [-1, attn_length, 1, attn_size])  # need to reshape before
            # Size of query vectors for attention.
            attention_vec_size = attn_size
            w = vs.get_variable(
                "Attn_Wd", [1, 1, attn_size, attention_vec_size])  # W_d
            hidden_features = nn_ops.conv2d(
                hidden, w, [1, 1, 1, 1], "SAME")  # W_d * h_o
            v = vs.get_variable("Attn_v", [attention_vec_size])  # v_d
            state = initial_state

            def attention(query):
                """Put attention masks on local_hidden using local_hidden_features and query."""
                # If the query is a tuple (when stacked RNN/LSTM), flatten it
                if nest.is_sequence(query):
                    query_list = nest.flatten(query)
                    for q in query_list:  # Check that ndims == 2 if specified.
                        ndims = q.get_shape().ndims
                        if ndims:
                            assert ndims == 2
                    query = array_ops.concat(query_list, 1)
                with vs.variable_scope("Attn_Wpd"):
                    # linear map
                    y = Linear(query, attention_vec_size, True)
                    y = array_ops.reshape(
                        y, [-1, 1, 1, attention_vec_size])
                    # Attention mask is a softmax of v_d^{\top} * tanh(...).
                    s = math_ops.reduce_sum(v * math_ops.tanh(hidden_features + y),
                                            [2, 3])
                    # Now calculate the attention-weighted vector, i.e., gamma in eq.[7]
                    a = nn_ops.softmax(s)
                    # eq. [8]
                    d = math_ops.reduce_sum(
                        array_ops.reshape(a, [-1, attn_length, 1, 1]) * hidden, [1, 2])
                return array_ops.reshape(d, [-1, attn_size])

            if initial_state_attention:
                attn = attention(initial_state)
            else:
                batch_attn_size = array_ops.stack([batch_size, attn_size])
                attn = array_ops.zeros(batch_attn_size, dtype=dtype)
                attn.set_shape([None, attn_size])

            i = 0
            outputs = []
            prev = None
            for inp, ext_inp in zip(decoder_inputs, external_inputs):
                if i > 0:
                    vs.get_variable_scope().reuse_variables()
                # If loop_function is set, we use it instead of decoder_inputs.
                if loop_function is not None and prev is not None:
                    with vs.variable_scope("loop_function", reuse=True):
                        inp = loop_function(prev, i)
                # Merge input and previous attentions into one vector of the right size.
                input_size = inp.get_shape().with_rank(2)[1]
                if input_size.value is None:
                    raise ValueError(
                        "Could not infer input size from input: %s" % inp.name)
                # we map the concatenation to shape [batch_size, input_size]
                if external_flag:
                    x = Linear([inp] + [ext_inp] + [attn], input_size, True)
                else:
                    x = Linear([inp] + [attn], input_size, True)
                # Run the RNN.
                cell_output, state = cell(x, state)
                # Run the attention mechanism.
                if i == 0 and initial_state_attention:
                    with vs.variable_scope(vs.get_variable_scope(), reuse=True):
                        attn = attention(state)
                else:
                    attn = attention(state)
                # Attention output projection
                with vs.variable_scope("AttnOutputProjection"):
                    output = Linear([cell_output] + [attn], output_size, True)
                if loop_function is not None:
                    prev = output
                outputs.append(output)
                i += 1
        return outputs, state

    def _loop_function(self, prev, _):
        """loop function used in the decoder to generate the next inupt"""
        return tf.matmul(prev, self.phs['w_out']) + self.phs['b_out']

    def mod_fn(self):
        encoder_attention_states, encoder_inputs, _labels, _external_inputs, decoder_inputs \
            = input_transform(self.phs['local_inputs'],
                              self.phs['global_inputs'],
                              self.phs['external_inputs'],
                              self.phs['local_attn_states'],
                              self.phs['global_attn_states'],
                              self.phs['labels'])

        n_stacked_layers = self.hps.n_stacked_layers  # num of layer stacked in RNN
        # dimension of encoder hidden/cell state
        n_hidden_encoder = self.hps.n_hidden_encoder
        # dimension of decoder hidden/cell state
        n_hidden_decoder = self.hps.n_hidden_decoder
        dropout_rate = self.hps.dropout_rate  # dropout rate in RNN unit
        n_output_decoder = self.hps.n_output_decoder

        # Define weights in the transformation layer of decoder
        self.phs['w_out'] = tf.get_variable('Weights_out',
                                            [n_hidden_decoder, n_output_decoder],
                                            dtype=tf.float32,
                                            initializer=tf.truncated_normal_initializer())

        self.phs['b_out'] = tf.get_variable('Biases_out',
                                            shape=[n_output_decoder],
                                            dtype=tf.float32,
                                            initializer=tf.constant_initializer(0.))

        with tf.variable_scope('GeoMAN'):
            # the implement of encoder
            with tf.variable_scope('Encoder'):
                cells = []
                for i in range(n_stacked_layers):
                    with tf.variable_scope('LSTM_{}'.format(i)):
                        cell = rnn.BasicLSTMCell(
                            n_hidden_encoder, forget_bias=1.0, state_is_tuple=True)
                        cell = tf.nn.rnn_cell.DropoutWrapper(
                            cell, output_keep_prob=1.0 - dropout_rate)
                        cells.append(cell)
                encoder_cell = tf.contrib.rnn.MultiRNNCell(cells)
                encoder_outputs, encoder_state, attn_weights = self.spatial_attention(encoder_inputs,
                                                                                      encoder_attention_states,
                                                                                      encoder_cell,
                                                                                      self.hps.s_attn_flag)

            # Calculate a concatenation of encoder outputs to put attention on.
            top_states = [tf.reshape(e, [-1, 1, encoder_cell.output_size])
                          for e in encoder_outputs]
            attention_states = tf.concat(top_states, 1)

            # the implement of decoder
            with tf.variable_scope('Decoder'):
                cells = []
                for i in range(n_stacked_layers):
                    with tf.variable_scope('LSTM_{}'.format(i)):
                        cell = rnn.BasicLSTMCell(n_hidden_decoder,
                                                 forget_bias=1.0,
                                                 state_is_tuple=True)
                        cell = tf.nn.rnn_cell.DropoutWrapper(cell,
                                                             output_keep_prob=1.0 - dropout_rate)
                        cells.append(cell)
                decoder_cell = tf.contrib.rnn.MultiRNNCell(cells)
                decoder_outputs, states = self.temporal_attention(decoder_inputs,
                                                                  _external_inputs,
                                                                  encoder_state,
                                                                  attention_states,
                                                                  decoder_cell,
                                                                  loop_function=self._loop_function,
                                                                  external_flag=self.hps.ext_flag)
            # generate outputs
            with tf.variable_scope('Prediction'):
                preds = [tf.matmul(i, self.phs['w_out']) +
                         self.phs['b_out'] for i in decoder_outputs]

        return preds

    def get_loss(self):
        """MSE loss"""
        # reshape
        n_steps_decoder = self.phs['labels'].get_shape()[1].value
        n_output_decoder = self.phs['labels'].get_shape()[2].value
        labels = tf.transpose(self.phs['labels'], [1, 0, 2])
        labels = tf.reshape(labels, [-1, n_output_decoder])
        labels = tf.split(labels, n_steps_decoder, 0)

        # compute empirical loss
        empirical_loss = 0
        # Extra: we can also get separate error at each future time slot
        for _y, _Y in zip(self.phs['preds'], labels):
            empirical_loss += tf.reduce_mean(tf.pow(_y - _Y, 2))
        self.phs['empirical_loss'] = empirical_loss
        return empirical_loss

    def get_l2reg_loss(self):
        """l2 reg loss"""
        reg_loss = 0
        for tf_var in tf.trainable_variables():
            if 'kernel:' in tf_var.name or 'bias:' in tf_var.name:
                reg_loss += tf.reduce_mean(tf.nn.l2_loss(tf_var))
        return self.lambda_l2_reg * reg_loss

    def train_op(self):
        # Training optimizer
        with tf.variable_scope('Optimizer'):
            global_step = tf.Variable(
                initial_value=0,
                name="global_step",
                trainable=False,
                collections=[tf.GraphKeys.GLOBAL_STEP,
                             tf.GraphKeys.GLOBAL_VARIABLES])
            optimizer = tf.contrib.layers.optimize_loss(
                loss=self.phs['loss'],
                learning_rate=self.hps.learning_rate,
                global_step=global_step,
                optimizer="Adam",
                clip_gradients=self.hps.gc_rate)
        return optimizer
    
    def summary(self):
        tf.summary.scalar("loss", self.phs['loss'])
        for var in tf.trainable_variables():
            tf.summary.histogram(var.name, var)
        return tf.summary.merge_all()