一、Keras實現Attention
from keras import backend as K
from keras.engine.topology import Layer
class Position_Embedding_Attention(Layer):
def __init__(self, size=None, mode='sum', **kwargs):
self.size = size #必須爲偶數
self.mode = mode
super(Position_Embedding, self).__init__(**kwargs)
def call(self, x):
if (self.size == None) or (self.mode == 'sum'):
self.size = int(x.shape[-1])
batch_size,seq_len = K.shape(x)[0],K.shape(x)[1]
position_j = 1. / K.pow(10000., 2 * K.arange(self.size / 2, dtype='float32') / self.size)
position_j = K.expand_dims(position_j, 0)
position_i = K.cumsum(K.ones_like(x[:,:,0]), 1)-1 #K.arange不支持變長，只好用這種方法生成
position_i = K.expand_dims(position_i, 2)
position_ij = K.dot(position_i, position_j)
position_ij = K.concatenate([K.cos(position_ij), K.sin(position_ij)], 2)
if self.mode == 'sum':
return position_ij + x
elif self.mode == 'concat':
return K.concatenate([position_ij, x], 2)
def compute_output_shape(self, input_shape):
if self.mode == 'sum':
return input_shape
elif self.mode == 'concat':
return (input_shape[0], input_shape[1], input_shape[2]+self.size)
class Multi_Head_AttentionAttention(Layer):
def __init__(self, nb_head, size_per_head, **kwargs):
self.nb_head = nb_head
self.size_per_head = size_per_head
self.output_dim = nb_head*size_per_head
super(Attention, self).__init__(**kwargs)
def build(self, input_shape):
self.WQ = self.add_weight(name='WQ',
shape=(input_shape[0][-1], self.output_dim),
initializer='glorot_uniform',
trainable=True)
self.WK = self.add_weight(name='WK',
shape=(input_shape[1][-1], self.output_dim),
initializer='glorot_uniform',
trainable=True)
self.WV = self.add_weight(name='WV',
shape=(input_shape[2][-1], self.output_dim),
initializer='glorot_uniform',
trainable=True)
super(Attention, self).build(input_shape)
def Mask(self, inputs, seq_len, mode='mul'):
if seq_len == None:
return inputs
else:
mask = K.one_hot(seq_len[:,0], K.shape(inputs)[1])
mask = 1 - K.cumsum(mask, 1)
for _ in range(len(inputs.shape)-2):
mask = K.expand_dims(mask, 2)
if mode == 'mul':
return inputs * mask
if mode == 'add':
return inputs - (1 - mask) * 1e12
def call(self, x):
#如果只傳入Q_seq,K_seq,V_seq，那麼就不做Mask
#如果同時傳入Q_seq,K_seq,V_seq,Q_len,V_len，那麼對多餘部分做Mask
if len(x) == 3:
Q_seq,K_seq,V_seq = x
Q_len,V_len = None,None
elif len(x) == 5:
Q_seq,K_seq,V_seq,Q_len,V_len = x
#對Q、K、V做線性變換
Q_seq = K.dot(Q_seq, self.WQ)
Q_seq = K.reshape(Q_seq, (-1, K.shape(Q_seq)[1], self.nb_head, self.size_per_head))
Q_seq = K.permute_dimensions(Q_seq, (0,2,1,3))
K_seq = K.dot(K_seq, self.WK)
K_seq = K.reshape(K_seq, (-1, K.shape(K_seq)[1], self.nb_head, self.size_per_head))
K_seq = K.permute_dimensions(K_seq, (0,2,1,3))
V_seq = K.dot(V_seq, self.WV)
V_seq = K.reshape(V_seq, (-1, K.shape(V_seq)[1], self.nb_head, self.size_per_head))
V_seq = K.permute_dimensions(V_seq, (0,2,1,3))
#計算內積，然後mask，然後softmax
A = K.batch_dot(Q_seq, K_seq, axes=[3,3]) / self.size_per_head**0.5
A = K.permute_dimensions(A, (0,3,2,1))
A = self.Mask(A, V_len, 'add')
A = K.permute_dimensions(A, (0,3,2,1))
A = K.softmax(A)
#輸出並mask
O_seq = K.batch_dot(A, V_seq, axes=[3,2])
O_seq = K.permute_dimensions(O_seq, (0,2,1,3))
O_seq = K.reshape(O_seq, (-1, K.shape(O_seq)[1], self.output_dim))
O_seq = self.Mask(O_seq, Q_len, 'mul')
return O_seq
def compute_output_shape(self, input_shape):
return (input_shape[0][0], input_shape[0][1], self.output_dim)
二、Tensrflow實現Attention機制
import tensorflow as tf
'''
inputs是一個形如(batch_size, seq_len, word_size)的張量；
函數返回一個形如(batch_size, seq_len, position_size)的位置張量。
'''
def Position_Embedding_Attention(inputs, position_size):
batch_size,seq_len = tf.shape(inputs)[0],tf.shape(inputs)[1]
position_j = 1. / tf.pow(10000.,2 * tf.range(position_size / 2, dtype=tf.float32 ) / position_size)
position_j = tf.expand_dims(position_j, 0)
position_i = tf.range(tf.cast(seq_len, tf.float32), dtype=tf.float32)
position_i = tf.expand_dims(position_i, 1)
position_ij = tf.matmul(position_i, position_j)
position_ij = tf.concat([tf.cos(position_ij), tf.sin(position_ij)], 1)
position_embedding = tf.expand_dims(position_ij, 0) + tf.zeros((batch_size, seq_len, position_size))
return position_embedding
'''
inputs是一個二階以上的張量，代表輸入序列，比如形如(batch_size, seq_len, input_size)的張量；
seq_len是一個形如(batch_size,)的張量，代表每個序列的實際長度，多出部分都被忽略；
mode分爲mul和add，mul是指把多出部分全部置零，一般用於全連接層之前；
add是指把多出部分全部減去一個大的常數，一般用於softmax之前。
'''
def Mask(inputs, seq_len, mode='mul'):
if seq_len == None:
return inputs
else:
mask = tf.cast(tf.sequence_mask(seq_len), tf.float32)
for _ in range(len(inputs.shape)-2):
mask = tf.expand_dims(mask, 2)
if mode == 'mul':
return inputs * mask
if mode == 'add':
return inputs - (1 - mask) * 1e12
'''
普通的全連接
inputs是一個二階或二階以上的張量，即形如(batch_size,...,input_size)。
只對最後一個維度做矩陣乘法，即輸出一個形如(batch_size,...,ouput_size)的張量。
'''
def Dense(inputs, ouput_size, bias=True, seq_len=None):
input_size = int(inputs.shape[-1])
W = tf.Variable(tf.random_uniform([input_size, ouput_size], -0.05, 0.05))
if bias:
b = tf.Variable(tf.random_uniform([ouput_size], -0.05, 0.05))
else:
b = 0
outputs = tf.matmul(tf.reshape(inputs, (-1, input_size)), W) + b
outputs = tf.reshape(outputs, \
tf.concat([tf.shape(inputs)[:-1], [ouput_size]], 0)
)
if seq_len != None:
outputs = Mask(outputs, seq_len, 'mul')
return outputs
'''
Multi-Head Attention的實現
'''
def Multi_Head_AttentionAttention(Q, K, V, nb_head, size_per_head, Q_len=None, V_len=None):
#對Q、K、V分別作線性映射
Q = Dense(Q, nb_head * size_per_head, False)
Q = tf.reshape(Q, (-1, tf.shape(Q)[1], nb_head, size_per_head))
Q = tf.transpose(Q, [0, 2, 1, 3])
K = Dense(K, nb_head * size_per_head, False)
K = tf.reshape(K, (-1, tf.shape(K)[1], nb_head, size_per_head))
K = tf.transpose(K, [0, 2, 1, 3])
V = Dense(V, nb_head * size_per_head, False)
V = tf.reshape(V, (-1, tf.shape(V)[1], nb_head, size_per_head))
V = tf.transpose(V, [0, 2, 1, 3])
#計算內積，然後mask，然後softmax
A = tf.matmul(Q, K, transpose_b=True) / tf.sqrt(float(size_per_head))
A = tf.transpose(A, [0, 3, 2, 1])
A = Mask(A, V_len, mode='add')
A = tf.transpose(A, [0, 3, 2, 1])
A = tf.nn.softmax(A)
#輸出並mask
O = tf.matmul(A, V)
O = tf.transpose(O, [0, 2, 1, 3])
O = tf.reshape(O, (-1, tf.shape(O)[1], nb_head * size_per_head))
O = Mask(O, Q_len, 'mul')
return O
if __name__ == '__main__':
import numpy as np
a = np.array([[1,1],[2,2],[3,3]])
b = 3 * a
print(b)