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