Tensorflow --BeamSearch

github:https://github.com/zle1992/Seq2Seq-Chatbot

1、 注意在infer阶段，需要需要reuse，

2、If you are using the BeamSearchDecoder with a cell wrapped in AttentionWrapper, then you must ensure that:

• The encoder output has been tiled to beam_width via tf.contrib.seq2seq.tile_batch (NOT tf.tile).
• The batch_size argument passed to the zero_state method of this wrapper is equal to true_batch_size * beam_width.
• The initial state created with zero_state above contains a cell_state value containing properly tiled final state from the encoder.

 import tensorflow as tf
 from tensorflow.python.layers.core import Dense
 
 BEAM_WIDTH = 5
 BATCH_SIZE = 128
 
 # INPUTS
 X = tf.placeholder(tf.int32, [BATCH_SIZE, None])
 Y = tf.placeholder(tf.int32, [BATCH_SIZE, None])
 X_seq_len = tf.placeholder(tf.int32, [BATCH_SIZE])
 Y_seq_len = tf.placeholder(tf.int32, [BATCH_SIZE])
 
 # ENCODER
 encoder_out, encoder_state = tf.nn.dynamic_rnn(
     cell = tf.nn.rnn_cell.BasicLSTMCell(128),
     inputs = tf.contrib.layers.embed_sequence(X, 10000, 128),
     sequence_length = X_seq_len,
     dtype = tf.float32)
 
 # DECODER COMPONENTS
 Y_vocab_size = 10000
 decoder_embedding = tf.Variable(tf.random_uniform([Y_vocab_size, 128], -1.0, 1.0))
 projection_layer = Dense(Y_vocab_size)
 
 # ATTENTION (TRAINING)
 with tf.variable_scope('shared_attention_mechanism'):
     attention_mechanism = tf.contrib.seq2seq.LuongAttention(
         num_units = 128,
         memory = encoder_out,
         memory_sequence_length = X_seq_len)
 
 decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
     cell = tf.nn.rnn_cell.BasicLSTMCell(128),
     attention_mechanism = attention_mechanism,
     attention_layer_size = 128)
 
 # DECODER (TRAINING)
 training_helper = tf.contrib.seq2seq.TrainingHelper(
     inputs = tf.nn.embedding_lookup(decoder_embedding, Y),
     sequence_length = Y_seq_len,
     time_major = False)
 training_decoder = tf.contrib.seq2seq.BasicDecoder(
     cell = decoder_cell,
     helper = training_helper,
     initial_state = decoder_cell.zero_state(BATCH_SIZE,tf.float32).clone(cell_state=encoder_state),
     output_layer = projection_layer)
 with tf.variable_scope('decode_with_shared_attention'):
     training_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(
         decoder = training_decoder,
         impute_finished = True,
         maximum_iterations = tf.reduce_max(Y_seq_len))
 training_logits = training_decoder_output.rnn_output
 
 # BEAM SEARCH TILE
 encoder_out = tf.contrib.seq2seq.tile_batch(encoder_out, multiplier=BEAM_WIDTH)
 X_seq_len = tf.contrib.seq2seq.tile_batch(X_seq_len, multiplier=BEAM_WIDTH)
 encoder_state = tf.contrib.seq2seq.tile_batch(encoder_state, multiplier=BEAM_WIDTH)
 
 # ATTENTION (PREDICTING)
 with tf.variable_scope('shared_attention_mechanism', reuse=True):
     attention_mechanism = tf.contrib.seq2seq.LuongAttention(
         num_units = 128,
         memory = encoder_out,
         memory_sequence_length = X_seq_len)
 
 decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
     cell = tf.nn.rnn_cell.BasicLSTMCell(128),
     attention_mechanism = attention_mechanism,
     attention_layer_size = 128)
 
 # DECODER (PREDICTING)
 predicting_decoder = tf.contrib.seq2seq.BeamSearchDecoder(
     cell = decoder_cell,
     embedding = decoder_embedding,
     start_tokens = tf.tile(tf.constant([1], dtype=tf.int32), [BATCH_SIZE]),
     end_token = 2,
     initial_state = decoder_cell.zero_state(BATCH_SIZE * BEAM_WIDTH,tf.float32).clone(cell_state=encoder_state),
     beam_width = BEAM_WIDTH,
     output_layer = projection_layer,
     length_penalty_weight = 0.0)
 with tf.variable_scope('decode_with_shared_attention', reuse=True):
     predicting_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(
         decoder = predicting_decoder,
         impute_finished = False,
         maximum_iterations = 2 * tf.reduce_max(Y_seq_len))
 predicting_logits = predicting_decoder_output.predicted_ids[:, :, 0]
 
 print('successful')

参考：

https://gist.github.com/higepon/eb81ba0f6663a57ff1908442ce753084

https://www.tensorflow.org/api_docs/python/tf/contrib/seq2seq/BeamSearchDecoder

https://github.com/tensorflow/nmt#beam-search