【TensorFlow】Seq2Seq模型的代码实现 & attention机制

import tensorflow as tf

SRC_TRAIN_DATA = '/path/to/data/train.en'  # 源语言输入文件
TRG_TRAIN_DATA = '/path/to/data/train.zh'  # 目标语言输入文件
CHECKPOINT_PATH = '/path/to/seq2seq_ckpt' 
HIDDEN_SIZE = 1024
NUM_LAYERS = 2
SRC_VOCAB_SIZE = 10000
TRG_VOCAB_SIZE = 4000
BATCH_SIZE = 100
NUM_EPOCH = 5
KEEP_PROB = 0.8
MAX_GRAD_NPRM = 5 # 用于控制梯度膨胀的梯度大小上限
SHARE_EMB_AND_SOFTMAX = TRUE # 在Softmax层和词向量层之间共享参数

class NMTModel(object):
	def __init__(self):
		# 定义编码器和解码器所使用的LSTM结构
		self.enc_cell = tf.nn.rnn_cell.MultiRNNCell(
			[tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE) for _ in range((NUM_LAYERS)])
		self.dec_cell = tf.nn.rnn_cell.MultiRNNCell(
			[tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE) for _ in range((NUM_LAYERS)])
		
		# 为源语言和目标语言分别定义词向量
		self.src_embedding = tf.get_variable("src_emb", [SRC_VOCAB_SIZE, HIDDEN_SIZE])
		self.trg_embedding = tf.get_variable("trg_emb", [TRG_VOCAB_SIZE, HIDDEN_SIZE])

		# 定义softmax层的变量
		# ! 词向量是(vocab_size, hidden_size), softmax层是(hidden_size, vocab_size),所以要转秩。因为共享词向量层和softmax层的参数,不仅能大幅减少参数数量,还能提高最终模型效果
		if SHARE_EMB_AND_SOFTMAX:
			self.softmax_weight = tf.transpose(self.trg_embedding) 
		else:
			self.softmax_weight = tf.get_variable("weight", [HIDDEN_SIZE, TRG_VOCAB_SIZE])
		self.softmax_biases = tf.get_variable("biases", [TRG_VOCAB_SIZE])

	# 在forward函数中定义模型的前向计算图
	def forward(self, src_input, src_size, trg_input, trg_label, trg_size):
		batch_size = tf.shape(src_input)[0]

		# 将输入和输出单词编号转为词向量
		src_emb = tf.nn.embedding_lookup(self.src_embedding, src_input)
		trg_emb = tf.nn.embedding_lookup(self.trg_embedding, trg_input)

		# 在词向量上进行dropout
		src_emb = tf.nn.dropout(src_emb, KEEP_PROB)
		trg_emb = tf.nn.dropout(trg_emb, KEEP_PROB)

		# 使用dynamic_rnn构造编码器,dynamic_rnn对每一个batch的数据读取两个输入,输入数据的内通和输入数据的长度。对于batch里的每一条数据,在读取了相应长度的内容后,dynamic_rnn就跳过后面的输入,直接把前一步的计算结果复制到后面的时刻。
		# 编码器读取源句子每个位置的词向量,输出最后一步的隐藏状态enc_state。因为编码器是一个双层lstm,因此enc_state是一个包含两个LSTMStateTuple类的tuple,每个LSTMStateTuple对应编码器中一层的状态
		# enc_outputs是顶层LSTM在每一步的输出,它的维度是[batch_size, max_time, HIDDEN_SIZE], seq2seq模型中不需要用到enc_outputs,而后面介绍的attention模型会用到它。
		with tf.variable_scope("encoder"):
			enc_outputs, enc_state = tf.nn.dynamic_rnn(self.enc_cell, src_emb, src_size, dtype=tf.float32)

		# 使用dynamic_rnn构造解码器
		# 解码器读取目标句子每个位置的词向量,输出的sec_outputs为每一步顶层lstm的输出,其维度是[batch_size, max_time, HIDDEN_SIZE]
		# initialize_state = enc_state表示用编码器的输出来初始化第一步的隐藏状态
		with tf.variable_scope("decoder"):
			dec_outputs, _ = tf.nn.dynamc_rnn(self.dec_cell, trg_emb, trg_size, initial_state=enc_state)

		# 计算解码器每一步的log perplexity
		output = tf.reshape(dec_outputs, [-1, HIDDEN_SIZE])
		logits = tf.matmul(output, self.softmax_weight)+self.softmax_biases
		loss = t.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.reshape(trg_label, [-1]), logits=logits)

		# 在计算平均损失时,需要将填充位置的权重设置为0,以避免无效位置的预测干扰模型的训练
		label_weights = tf.sequence_mask(trg_size, maxlen=tf.shape(trg_label)[1], dtype=tf.float32)
		label_weights = tf.reshape(label_weights, [-1])
		cost = tf.reduce_sum(loss * label_weights)
		cost_per_token = cost/tf.reduce_sum(label_weights)

		trainable_variables = tf.trainable_variables()

		# 控制梯度大小
		grads = tf.gradients(cost/tf.to_float(batch_size), trainable_variables)
		grads, _ = tf.clip_by_global_norm(grads, MAX_GRAD_NORM)

		optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0)
		train_op = optimizer.apply_gradients(zip(grads, trainable_variables))

		return cost_per_token, train_op

def run_epoch(session, cost_op, train_op, saver, step):
	while True:
		try:
			cost, _ = session.run([cost_op, train_op])
			if step%200==0:
				saver.save(session, CHECKPOINT_PATH, global_step=step)
			step += 1
		except tf.errors.OutOfRangeError:
			break
	return step

def main():
	initializer = tf.random_uniform_initializer(-0.05, 0.05)

	with tf.variable_scope("nmt_model", reuse=None. initializer=initializer):
		train_model = NMTModel

	data = MakeSrcTrgDataset(SRC_TRAIN_DATA, TRG_TRAIN_DATA, BATCH_SIZE)
	iterator = data.make_initializable_iterator()
	(src, src_size), (trg_input, trg_label, trg_size) = iterator.get_next()

	cost_op, train_op = train_model.forward(src, src_size, trg_input, trg_label, trg_size)

	saver = tf.train.Saver()
	step=0
	with tf.Session() as sess:
		tf.global_variables_initializer().run()
		for i in range(NUM_EPOCH):
			sess.run(iterator.initializer)
			step = run_epoch(sess, cost_op, train_op, saver, step)

if __name__ == "__main__":
	main()
	

上面的程序完成了机器翻译模型的训练步骤,并将训练好的模型保存到checkpoint中,下面是怎样从checkpoint中读取模型并对一个新的句子进行翻译。
在解码的过程中,解码器的实现与训练时有很大不同。这是因为训练时解码器可以从输入中读取完整的目标训练句子,因此可以用dynamic_rnn简单地展开成前馈网络。而在解码过程中,模型只能看到输入句子,却看不到目标句子。解码器在第一步读取符,预测目标句子的第一个单词。然后需要将这个预测的单词复制到第二步作为输入,再预测第二个单词,直到预测的单词为止,这个过程需要使用一个循环结构来实现,在tensorflow中,循环结构是由tf.while_loop来实现的。tf.while_loop的使用方法如下:

# cood是一个函数,负责判断继续执行循环的条件。
# loop_body是每个循环体内执行的操作,负责对循环状态进行更新
# init_state为循环的起始状态,它可以包含多个Tensor或者TensorArray
# 返回的结果是循环结束时的循环状态
final_state = tf.while_loop(cood, loop_body, init_state)

下面的代码展示了如何使用tf.while_loop来实现解码的过程

import tensorflow as tf
CHECKPOINT_PATH = '/path/to/seq2seq_ckpt-9000'

# 模型参数,必须与训练时的模型参数保持一致
HIDDEN_SIZE = 1024
NUM_LAYERS = 2
SRC_VOCAB_SIZE = 10000
TRG_VOCAB_SIZE = 4000
SHARE_EMB_AND_SOFTMAX = True

SOS_ID = 1
EOS_ID = 2

class NMTModel(object):
	def __init__(self):
		# 与训练时的__init__函数相同,通常在训练程序和解码程序中复用NMTModel类及其init函数,以确保解码时和训练时定义的变量是相同的

	def inference(self, src_input):
		src_size = tf.convert_to_tensor([len(src_input)], dtype=tf.int32)
		src_input = tf.convert_to_tensor([src_input], dtype=tf.int32)
		src_emb = tf.nn.embedding_lookup(self.src_embedding, src_input)

		# 使用dynamic_rnn构造编码器,这一步与训练时相同
		with tf.variable_scope("encoder"):
			enc_outputs, enc_state = tf.nn.dynamic_rnn(self.enc_cell, src_emb, src_size, dtype=tf.float32)

		# 设置解码的最大步数,这是为了避免在极端情况下出现无限循环的问题
		MAX_DEC_LEN = 100
		
		with tf.variable_scope("decoder/rnn/multi_rnn_cell"):
			# 使用一个变长的TensorArray来存储生成的句子
			init_array = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True, clear_after_read=False)

			# 填入第一个单词<sos>作为解码器的输入
			init_array = init_array.write(0, SOS_ID)
			# 构建初始的循环状态。循环状态包含循环神经网络的隐藏状态,保存生成句子的TensorArray,以及记录解码步数的一个整数step
			init_loop_var = (enc_state, init_array, 0)

			# tf.while_loop的循环条件:循环直到解码器输出<eos>,或者达到最大步数为止
			def continue_loop_condition(state, trg_ids, step):
				return tf.reduce_all(tf.logical_and(
					tf.not_equal(trg_ids.read(step), EOS_ID),
					tf.less(step, MAX_DEC_LEN-1)))

			def loop_body(state, trg_ids, step):
				# 读取最后一步输出的单词,并读取其词向量
				trg_input = [trg_ids.read(step)]
				trg_emb= tf.nn.embedding_lookup(self.trg_embedding, trg_input)

				# 这里不调用dynamic_rnn, 而是直接调用dec_cell向前计算一步
				dec_outputs, next_state = self.dec_cell.call(state=state, inputs=trg_emb)

				# 计算每个可能的输出单词对应的logit,并选取logit值最大的单词作为这一步的输出
				output = tf.reshape(dec_output, [-1, HIDDEN_SIZE])
				logits = (tf.matmul(output, self.softmax_weight)+self.biases)
				next_id = tf.argmax(logits, axis=1, output_type=tf.int32)

				# 将这一步输出的单词写入循环状态的trg_ids中
				trg_ids = trg_ids.write(step+1, next_id[0])

				return next_state, trg_ids, step+1

			# 执行tf.while_loop,返回最终状态
			state, trg_ids, step = tf.while_loop(continue_loop_condition, loop_body, init_loop_var)
			
			return trg_ids.stack()
			
def main():
		with tf.variable_scope("nmt_model", reuse=None):
			model = NMTModel()

			# 定义一个测试例子
			test_sentence = [90, 13, 9, 689, 4, 2]

			# 简历解码所需的计算图
			output_op = model.inference(test_sentence)
			sess = tf.Session()
			saver = tf.train.Saver()
			saver.restore(sess, CHECKPOINT_PATH)
			output = sess.run(output_op)
			print output
			sess.close()

if __name__=="__main__":
	main() 

【注意力机制】
注意力机制允许解码器随时查阅输入句子中的部分单词或片段,因此不再需要在中间向量中存储所有信息——解码器在解码的每一步将隐藏状态作为查询的输入来“查询”编码器的隐藏状态,在每个输入的位置计算一个反应与查询输入相关程度的权重,再根据这个权重对各输入位置的隐藏状态求加权平均,加权平均后得到的向量称为“context”,表示它是与翻译当前单词最相关的原文信息。在解码下一个单词时,将context作为额外信息输入到循环神经网络中,这样循环神经网络可以时刻读取原文中最相关的信息,而不必完全依赖于上一时刻的隐藏状态。

# 下面self.enc_cell_fw和self.enc_cell_bw定义了编码器中的前向和后项循环网络, 它们取代了seq2seq样例代码中__init__函数里的self.enc_cell
self.enc_cell_fw = tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE)
self.enc_cell_bw = tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE)

# 下面的代码取代了seq2seq样例代码中forward函数里的相应部分
with tf.variable_scope("encoder"):
	# 构造编码器时,使用bidirectional_dynamic_rnn构造双向循环网络,双向循环网络的顶层输出enc_outputs是一个包含两个张量的tuuple,每个张量的维度都是[batch_size, max_time, HIDDEN_SIZE],代表两个lstm在每一步的输出
	enc_outputs, enc_state = tf.nn.bidirectional_dynamic_rnn(self.enc_cell_fw, self.enc_cell_bw, src_emb, src_size, dtype=tf.float32)
	# 将两个lstm的输出拼接为一个张量
	enc_outputs = tf.concat([enc_outputs[0], enc_outputs[1]], -1)

with tf.variable_scope("decoder"):
	# 选择注意力权重的计算模型,BahanauAttention是使用一个隐藏层的前馈神经网络,memory_sequence_length是一个维度为[batch]的张量,代表batch中每一个句子的长度,attention需要根据这个信息把填充位置的注意力权重设置为0
	attention_mechanism = tf.contrib.seq2seq.BahanauAttention(HIDDEN_SIZE, enc_outputs, memory_sequence_length=src_size)

	# 将解码器的循环神经网络self.dec_cell和注意力一起封装成更高层的循环神经网络
	attention_cell = tf.contrib.seq2seq.AttentionWrapper(self.dec_cel, attenion_mechanism, attention_layer_size=HIDDEN_SIZE)
	
	# 使用attention_cell和dynamic_rnn构造编码器,这里没有指定init_state,也就是没有使用编码器的输出来初始化输入,而完全依赖著利益作为信息来源
	dec_outputs, _ = tf.nn.dynamic_enn(attention_cell, trg_emb, trg_size, dtype=tf.float32)

参考资料:
TensorFlow:实战Google深度学习框架-中国工信出版集团

上一篇:论文阅读:《Attention Is All You Need》


下一篇:5.3 Transformer意境级讲解