在https://github.com/tensorflow/models/blob/master/tutorials/rnn/translate/translate.py
加入这两个函数
def self_test_predict():
"""Test the translation model."""
with tf.Session() as sess:
print("Self-test for neural translation model.")
# Create model with vocabularies of 10, 2 small buckets, 2 layers of 32.
model = seq2seq_model.Seq2SeqModel(10, 10, [(3, 3), (6, 6)], 32, 2,
5.0, 1, 0.3, 0.99, num_samples=8,forward_only=True)
sess.run(tf.global_variables_initializer())
ckpt = tf.train.get_checkpoint_state(".\\tmp")
if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
print("Reading model parameters from %s" % ckpt.model_checkpoint_path)
model.saver.restore(sess, ckpt.model_checkpoint_path)
# Fake data set for both the (3, 3) and (6, 6) bucket.
data_set = ([([1, 1], [2, 2]), ([3, 3], [4]), ([5], [6])],
[([1, 1, 1, 1, 1], [2, 2, 2, 2, 2]), ([3, 3, 3], [5, 6])])
for _ in xrange(5): # Train the fake model for 5 steps.
bucket_id = random.choice([0, 1])
encoder_inputs, decoder_inputs, target_weights = model.get_batch(
data_set, bucket_id)
print(encoder_inputs)
_, eval_loss, output_logits = model.step(sess, encoder_inputs, decoder_inputs, target_weights,
bucket_id, True)
outputs = [np.argmax(logit, axis=1) for logit in output_logits]
print(outputs)
def self_test_train():
"""Test the translation model."""
with tf.Session() as sess:
print("Self-test for neural translation model.")
# Create model with vocabularies of 10, 2 small buckets, 2 layers of 32.
model = seq2seq_model.Seq2SeqModel(10, 10, [(3, 3), (6, 6)], 32, 2,
5.0, 16, 0.3, 0.99, num_samples=8)
sess.run(tf.global_variables_initializer())
# Fake data set for both the (3, 3) and (6, 6) bucket.
data_set = ([([1, 1], [2, 2]), ([3, 3], [4]), ([5], [6])],
[([1, 1, 1, 1, 1], [2, 2, 2, 2, 2]), ([3, 3, 3], [5, 6])])
for step in xrange(501): # Train the fake model for 5 steps.
bucket_id = random.choice([0, 1])
encoder_inputs, decoder_inputs, target_weights = model.get_batch(
data_set, bucket_id)
loss,_,_ = model.step(sess, encoder_inputs, decoder_inputs, target_weights,
bucket_id, False)
if step % 100 == 0:
print(step,loss)
checkpoint_path = os.path.join(".\\tmp", "translate_example.ckpt")
model.saver.save(sess, checkpoint_path, global_step=model.global_step)
