You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
80 lines
3.0 KiB
80 lines
3.0 KiB
import tensorflow as tf
|
|
|
|
try:
|
|
import tensorflow.python.keras as keras
|
|
from tensorflow.python.keras import layers
|
|
from tensorflow.python.keras import backend as K
|
|
except:
|
|
import tensorflow.keras as keras
|
|
from tensorflow.keras import layers
|
|
from tensorflow.keras import backend as K
|
|
|
|
|
|
class Attention_layer(layers.Layer):
|
|
def __init__(self, **kwargs):
|
|
super(Attention_layer, self).__init__(**kwargs)
|
|
|
|
def build(self, inputs_shape):
|
|
assert isinstance(inputs_shape, list)
|
|
|
|
# inputs: --> size
|
|
# [(None,max_contents,code_vector_size),(None,max_contents)]
|
|
# the second input is optional
|
|
if (len(inputs_shape) < 1 or len(inputs_shape) > 2):
|
|
raise ValueError("AttentionLayer expect one or two inputs.")
|
|
|
|
# (None,max_contents,code_vector_size)
|
|
input_shape = inputs_shape[0]
|
|
|
|
if (len(input_shape) != 3):
|
|
raise ValueError("Input shape for AttentionLayer shoud be of 3 dimensions.")
|
|
|
|
self.input_length = int(input_shape[1])
|
|
self.input_dim = int(input_shape[2])
|
|
|
|
attention_param_shape = (self.input_dim, 1)
|
|
|
|
self.attention_param = self.add_weight(
|
|
name='attention_param',
|
|
shape=attention_param_shape,
|
|
initializer='uniform',
|
|
trainable=True,
|
|
dtype=tf.float32
|
|
)
|
|
|
|
super(Attention_layer, self).build(input_shape)
|
|
|
|
def call(self, inputs, **kwargs):
|
|
assert isinstance(inputs, list)
|
|
|
|
# inputs:
|
|
# [(None,max_contents,code_vector_size),(None,max_contents)]
|
|
# the second input is optional
|
|
if (len(inputs) < 1 or len(inputs) > 2):
|
|
raise ValueError("AttentionLayer expect one or two inputs.")
|
|
|
|
actual_input = inputs[0]
|
|
mask = inputs[1] if (len(inputs) > 1) else None
|
|
|
|
if mask is not None and not (((len(mask.shape) == 3 and mask.shape[2] == 1) or (len(mask.shape) == 2)) and (
|
|
mask.shape[1] == self.input_length)):
|
|
raise ValueError(
|
|
"`mask` shoud be of shape (batch, input_length) or (batch, input_length, 1) when calling AttentionLayer.")
|
|
|
|
assert actual_input.shape[-1] == self.attention_param.shape[0]
|
|
|
|
# (batch, input_length, input_dim) * (input_dim, 1) ==> (batch, input_length, 1)
|
|
attention_weights = K.dot(actual_input, self.attention_param)
|
|
|
|
if mask is not None:
|
|
if (len(mask.shape) == 2):
|
|
mask = K.expand_dims(mask, axis=2) # (batch, input_dim, 1)
|
|
mask = K.log(mask) # e.g. K.exp(K.log(0.)) = 0 K.exp(K.log(1.)) =1
|
|
attention_weights += mask
|
|
|
|
attention_weights = K.softmax(attention_weights, axis=1)
|
|
result = K.sum(actual_input * attention_weights, axis=1)
|
|
return result, attention_weights
|
|
|
|
def compute_output_shape(self, input_shape):
|
|
return input_shape[0], input_shape[2] # (batch,input_length,input_dim) --> (batch,input_dim)
|