mirror of
https://github.com/wassname/keras-language-modeling.git
synced 2026-09-11 12:20:57 +08:00
made the attentional lstm code much more compact
This commit is contained in:
codekansas
+8
-225
@@ -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()))
|
||||
|
||||
Reference in New Issue
Block a user