diff --git a/attention_lstm.py b/attention_lstm.py index 7587f42..c6217ad 100644 --- a/attention_lstm.py +++ b/attention_lstm.py @@ -2,109 +2,23 @@ from __future__ import absolute_import import numpy as np from keras import backend as K -from keras import activations, initializations, regularizers -from keras.engine import InputSpec -from keras.layers import Recurrent, time_distributed_dense +from keras.layers import LSTM -class AttentionLSTM(Recurrent): - '''Attentional LSTM - Tan et. al. 2016 - # Arguments - output_dim: dimension of the internal projections and the final output. - attention_vec: source layer of the attention vector. should be a flat layer. - init: weight initialization function. - Can be the name of an existing function (str), - or a Theano function (see: [initializations](../initializations.md)). - inner_init: initialization function of the inner cells. - forget_bias_init: initialization function for the bias of the forget gate. - [Jozefowicz et al.](http://www.jmlr.org/proceedings/papers/v37/jozefowicz15.pdf) - recommend initializing with ones. - activation: activation function. - Can be the name of an existing function (str), - or a Theano function (see: [activations](../activations.md)). - inner_activation: activation function for the inner cells. - W_regularizer: instance of [WeightRegularizer](../regularizers.md) - (eg. L1 or L2 regularization), applied to the input weights matrices. - U_regularizer: instance of [WeightRegularizer](../regularizers.md) - (eg. L1 or L2 regularization), applied to the recurrent weights matrices. - b_regularizer: instance of [WeightRegularizer](../regularizers.md), - applied to the bias. - dropout_W: float between 0 and 1. Fraction of the input units to drop for input gates. - dropout_U: float between 0 and 1. Fraction of the input units to drop for recurrent connections. - # References - - [LSTM-based deep learning model for non-factoid answer selection](http://arxiv.org/pdf/1511.04108.pdf) - ''' - def __init__(self, output_dim, attention_vec, - init='glorot_uniform', inner_init='orthogonal', - forget_bias_init='one', activation='tanh', - inner_activation='hard_sigmoid', - W_regularizer=None, U_regularizer=None, b_regularizer=None, - dropout_W=0., dropout_U=0., **kwargs): +class AttentionLSTM(LSTM): + def __init__(self, output_dim, attention_vec, **kwargs): self.attention_vec = attention_vec - self.output_dim = output_dim - self.init = initializations.get(init) - self.inner_init = initializations.get(inner_init) - self.forget_bias_init = initializations.get(forget_bias_init) - self.activation = activations.get(activation) - self.inner_activation = activations.get(inner_activation) - self.W_regularizer = regularizers.get(W_regularizer) - self.U_regularizer = regularizers.get(U_regularizer) - self.b_regularizer = regularizers.get(b_regularizer) - self.dropout_W, self.dropout_U = dropout_W, dropout_U - self.attention_units = list() - - if self.dropout_W or self.dropout_U: - self.uses_learning_phase = True - super(AttentionLSTM, self).__init__(**kwargs) + super(AttentionLSTM, self).__init__(output_dim, **kwargs) def build(self, input_shape): - self.input_spec = [InputSpec(shape=input_shape)] - - input_dim = input_shape[2] + super(AttentionLSTM, self).build(input_shape) if hasattr(self.attention_vec, '_keras_shape'): attention_dim = self.attention_vec._keras_shape[1] else: raise Exception('Layer could not be build: No information about expected input shape.') - self.input_dim = input_dim - - if self.stateful: - self.reset_states() - else: - # initial states: 2 all-zero tensors of shape (output_dim) - self.states = [None, None] - - self.W_i = self.init((input_dim, self.output_dim), - name='{}_W_i'.format(self.name)) - self.U_i = self.inner_init((self.output_dim, self.output_dim), - name='{}_U_i'.format(self.name)) - self.b_i = K.zeros((self.output_dim,), name='{}_b_i'.format(self.name)) - - self.W_f = self.init((input_dim, self.output_dim), - name='{}_W_f'.format(self.name)) - self.U_f = self.inner_init((self.output_dim, self.output_dim), - name='{}_U_f'.format(self.name)) - self.b_f = self.forget_bias_init((self.output_dim,), - name='{}_b_f'.format(self.name)) - - self.W_c = self.init((input_dim, self.output_dim), - name='{}_W_c'.format(self.name)) - self.U_c = self.inner_init((self.output_dim, self.output_dim), - name='{}_U_c'.format(self.name)) - self.b_c = K.zeros((self.output_dim,), name='{}_b_c'.format(self.name)) - - self.W_o = self.init((input_dim, self.output_dim), - name='{}_W_o'.format(self.name)) - self.U_o = self.inner_init((self.output_dim, self.output_dim), - name='{}_U_o'.format(self.name)) - self.b_o = K.zeros((self.output_dim,), name='{}_b_o'.format(self.name)) - - ######################## - # Attention parameters # - ######################## - self.U_a = self.inner_init((self.output_dim, self.output_dim), name='{}_U_a'.format(self.name)) self.b_a = K.zeros((self.output_dim,), name='{}_b_a'.format(self.name)) @@ -117,109 +31,16 @@ class AttentionLSTM(Recurrent): name='{}_U_s'.format(self.name)) self.b_s = K.zeros((self.output_dim,), name='{}_b_s'.format(self.name)) - self.regularizers = [] - if self.W_regularizer: - self.W_regularizer.set_param(K.concatenate([self.W_i, - self.W_f, - self.W_c, - self.W_o])) - self.regularizers.append(self.W_regularizer) - if self.U_regularizer: - self.U_regularizer.set_param(K.concatenate([self.U_a, - self.U_m, - self.U_s, - self.U_i, - self.U_f, - self.U_c, - self.U_o])) - self.regularizers.append(self.U_regularizer) - if self.b_regularizer: - self.b_regularizer.set_param(K.concatenate([self.b_a, - self.b_m, - self.b_s, - self.b_i, - self.b_f, - self.b_c, - self.b_o])) - self.regularizers.append(self.b_regularizer) - - self.trainable_weights = [self.W_i, self.U_i, self.b_i, - self.W_c, self.U_c, self.b_c, - self.W_f, self.U_f, self.b_f, - self.W_o, self.U_o, self.b_o, - self.U_a, self.U_m, self.U_s, - self.b_a, self.b_m, self.b_s] + self.trainable_weights += [self.U_a, self.U_m, self.U_s, self.b_a, self.b_m, self.b_s] if self.initial_weights is not None: self.set_weights(self.initial_weights) del self.initial_weights - def reset_states(self): - assert self.stateful, 'Layer must be stateful.' - input_shape = self.input_spec[0].shape - if not input_shape[0]: - raise Exception('If a RNN is stateful, a complete ' + - 'input_shape must be provided (including batch size).') - if hasattr(self, 'states'): - K.set_value(self.states[0], - np.zeros((input_shape[0], self.output_dim))) - K.set_value(self.states[1], - np.zeros((input_shape[0], self.output_dim))) - else: - self.states = [K.zeros((input_shape[0], self.output_dim)), - K.zeros((input_shape[0], self.output_dim))] - - def preprocess_input(self, x, train=False): - if self.consume_less == 'cpu': - if train and (0 < self.dropout_W < 1): - dropout = self.dropout_W - else: - dropout = 0 - input_shape = self.input_spec[0].shape - input_dim = input_shape[2] - timesteps = input_shape[1] - - x_i = time_distributed_dense(x, self.W_i, self.b_i, dropout, - input_dim, self.output_dim, timesteps) - x_f = time_distributed_dense(x, self.W_f, self.b_f, dropout, - input_dim, self.output_dim, timesteps) - x_c = time_distributed_dense(x, self.W_c, self.b_c, dropout, - input_dim, self.output_dim, timesteps) - x_o = time_distributed_dense(x, self.W_o, self.b_o, dropout, - input_dim, self.output_dim, timesteps) - return K.concatenate([x_i, x_f, x_c, x_o], axis=2) - else: - return x - def step(self, x, states): - h_tm1 = states[0] - c_tm1 = states[1] - B_U = states[2] - B_W = states[3] + h, [h, c] = super(AttentionLSTM, self).step(x, states) attention = states[4] - if self.consume_less == 'cpu': - x_i = x[:, :self.output_dim] - x_f = x[:, self.output_dim: 2 * self.output_dim] - x_c = x[:, 2 * self.output_dim: 3 * self.output_dim] - x_o = x[:, 3 * self.output_dim:] - else: - x_i = K.dot(x * B_W[0], self.W_i) + self.b_i - x_f = K.dot(x * B_W[1], self.W_f) + self.b_f - x_c = K.dot(x * B_W[2], self.W_c) + self.b_c - x_o = K.dot(x * B_W[3], self.W_o) + self.b_o - - i = self.inner_activation(x_i + K.dot(h_tm1 * B_U[0], self.U_i)) - f = self.inner_activation(x_f + K.dot(h_tm1 * B_U[1], self.U_f)) - c = f * c_tm1 + i * self.activation(x_c + K.dot(h_tm1 * B_U[2], self.U_c)) - o = self.inner_activation(x_o + K.dot(h_tm1 * B_U[3], self.U_o)) - - h = o * self.activation(c) - - ################## - # Attention gate # - ################## - m = K.tanh(K.dot(h, self.U_a) + attention + self.b_a) # Intuitively it makes more sense to use a sigmoid (was getting some NaN problems # which I think might have been caused by the exponential function -> gradients blow up) @@ -230,44 +51,6 @@ class AttentionLSTM(Recurrent): return h, [h, c] def get_constants(self, x): - constants = [] - if 0 < self.dropout_U < 1: - ones = K.ones_like(K.reshape(x[:, 0, 0], (-1, 1))) - ones = K.concatenate([ones] * self.output_dim, 1) - B_U = [K.dropout(ones, self.dropout_U) for _ in range(4)] - constants.append(B_U) - else: - constants.append([K.cast_to_floatx(1.) for _ in range(4)]) - - if self.consume_less == 'cpu' and 0 < self.dropout_W < 1: - input_shape = self.input_spec[0].shape - input_dim = input_shape[-1] - ones = K.ones_like(K.reshape(x[:, 0, 0], (-1, 1))) - ones = K.concatenate([ones] * input_dim, 1) - B_W = [K.dropout(ones, self.dropout_W) for _ in range(4)] - constants.append(B_W) - else: - constants.append([K.cast_to_floatx(1.) for _ in range(4)]) - - #################### - # Attention vector # - #################### - + constants = super(AttentionLSTM, self).get_constants(x) constants.append(K.dot(self.attention_vec, self.U_m) + self.b_m) - return constants - - def get_config(self): - config = {"output_dim": self.output_dim, - "init": self.init.__name__, - "inner_init": self.inner_init.__name__, - "forget_bias_init": self.forget_bias_init.__name__, - "activation": self.activation.__name__, - "inner_activation": self.inner_activation.__name__, - "W_regularizer": self.W_regularizer.get_config() if self.W_regularizer else None, - "U_regularizer": self.U_regularizer.get_config() if self.U_regularizer else None, - "b_regularizer": self.b_regularizer.get_config() if self.b_regularizer else None, - "dropout_W": self.dropout_W, - "dropout_U": self.dropout_U} - base_config = super(AttentionLSTM, self).get_config() - return dict(list(base_config.items()) + list(config.items()))