# coding=utf-8

# Copyright 2017 The Tensor2Tensor Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Tests for trainer_utils."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import os

import shutil

# Dependency imports

from tensor2tensor.data_generators import algorithmic

from tensor2tensor.data_generators import generator_utils

from tensor2tensor.models import transformer

from tensor2tensor.utils import model_builder

from tensor2tensor.utils import registry

from tensor2tensor.utils import trainer_utils

import tensorflow as tf

flags = tf.flags

FLAGS = tf.flags.FLAGS

flags.DEFINE_string("schedule", "train_and_evaluate", "")

flags.DEFINE_integer("eval_steps", 10, "Number of steps in evaluation.")

flags.DEFINE_string("master", "", "Address of TensorFlow master.")

flags.DEFINE_string("output_dir", "", "Base output directory for run.")

@registry.register_problem

class TinyAlgo(algorithmic.AlgorithmicIdentityBinary40):

def generate_data(self, data_dir, _):

identity_problem = algorithmic.AlgorithmicIdentityBinary40()

generator_utils.generate_files(

identity_problem.generator(self.num_symbols, 40, 100000),

self.training_filepaths(data_dir, 1, shuffled=True), 100)

generator_utils.generate_files(

identity_problem.generator(self.num_symbols, 400, 10000),

self.dev_filepaths(data_dir, 1, shuffled=True), 100)

@registry.register_hparams

def transformer_test():

hparams = transformer.transformer_base()

hparams.batch_size = 10

hparams.hidden_size = 10

hparams.num_hidden_layers = 1

hparams.num_heads = 2

hparams.max_length = 16

return hparams

class TrainerUtilsTest(tf.test.TestCase):

@classmethod

def setUpClass(cls):

tmp_dir = tf.test.get_temp_dir()

shutil.rmtree(tmp_dir)

os.mkdir(tmp_dir)

# Generate a small test dataset

FLAGS.problems = "tiny_algo"

TrainerUtilsTest.data_dir = tmp_dir

registry.problem(FLAGS.problems).generate_data(TrainerUtilsTest.data_dir,

None)

def testModelsImported(self):

models = registry.list_models()

self.assertTrue("lstm_seq2seq" in models)

def testHParamsImported(self):

hparams = registry.list_hparams()

self.assertTrue("transformer_base" in hparams)

def testSingleStep(self):

model_name = "transformer"

data_dir = TrainerUtilsTest.data_dir

hparams = trainer_utils.create_hparams("transformer_test", data_dir)

trainer_utils.add_problem_hparams(hparams, FLAGS.problems)

exp = trainer_utils.create_experiment(

data_dir=data_dir,

model_name=model_name,

train_steps=1,

eval_steps=1,

hparams=hparams,

run_config=trainer_utils.create_run_config(

output_dir=tf.test.get_temp_dir()))

exp.test()

def testSingleEvalStepRawSession(self):

"""Illustrate how to run a T2T model in a raw session."""

# Set model name, hparams, problems as would be set on command line.

model_name = "transformer"

FLAGS.hparams_set = "transformer_test"

FLAGS.problems = "tiny_algo"

data_dir = "/tmp" # Used only when a vocab file or such like is needed.

# Create the problem object, hparams, placeholders, features dict.

encoders = registry.problem(FLAGS.problems).feature_encoders(data_dir)

hparams = trainer_utils.create_hparams(FLAGS.hparams_set, data_dir)

trainer_utils.add_problem_hparams(hparams, FLAGS.problems)

inputs_ph = tf.placeholder(dtype=tf.int32) # Just length dimension.

batch_inputs = tf.reshape(inputs_ph, [1, -1, 1, 1]) # Make it 4D.

# In INFER mode targets can be None.

targets_ph = tf.placeholder(dtype=tf.int32) # Just length dimension.

batch_targets = tf.reshape(targets_ph, [1, -1, 1, 1]) # Make it 4D.

features = {

"inputs": batch_inputs,

"targets": batch_targets,

"problem_choice": tf.constant(0), # We run on the first problem here.

"input_space_id": tf.constant(hparams.problems[0].input_space_id),

"target_space_id": tf.constant(hparams.problems[0].target_space_id)

}

# Now set a mode and create the graph by invoking model_fn.

mode = tf.estimator.ModeKeys.EVAL

estimator_spec = model_builder.model_fn(

model_name, features, mode, hparams, problem_names=[FLAGS.problems])

predictions_dict = estimator_spec.predictions

predictions = tf.squeeze( # These are not images, axis=2,3 are not needed.

predictions_dict["predictions"],

axis=[2, 3])

# Having the graph, let's run it on some data.

with self.test_session() as sess:

sess.run(tf.global_variables_initializer())

inputs = "0 1 0"

targets = "0 1 0"

# Encode from raw string to numpy input array using problem encoders.

inputs_numpy = encoders["inputs"].encode(inputs)

targets_numpy = encoders["targets"].encode(targets)

# Feed the encoded inputs and targets and run session.

feed = {inputs_ph: inputs_numpy, targets_ph: targets_numpy}

np_predictions = sess.run(predictions, feed)

# Check that the result has the correct shape: batch x length x vocab_size

# where, for us, batch = 1, length = 3, vocab_size = 4.

self.assertEqual(np_predictions.shape, (1, 3, 4))

def testSingleTrainStepCall(self):

"""Illustrate how to run a T2T model in a raw session."""

# Set model name, hparams, problems as would be set on command line.

model_name = "transformer"

FLAGS.hparams_set = "transformer_test"

FLAGS.problems = "tiny_algo"

data_dir = "/tmp" # Used only when a vocab file or such like is needed.

# Create the problem object, hparams, placeholders, features dict.

encoders = registry.problem(FLAGS.problems).feature_encoders(data_dir)

hparams = trainer_utils.create_hparams(FLAGS.hparams_set, data_dir)

trainer_utils.add_problem_hparams(hparams, FLAGS.problems)

# Now set a mode and create the model.

mode = tf.estimator.ModeKeys.TRAIN

model = registry.model(model_name)(hparams, mode)

# Create placeholder for features and make them batch-sized.

inputs_ph = tf.placeholder(dtype=tf.int32) # Just length dimension.

batch_inputs = tf.reshape(inputs_ph, [1, -1, 1, 1]) # Make it 4D.

targets_ph = tf.placeholder(dtype=tf.int32) # Just length dimension.

batch_targets = tf.reshape(targets_ph, [1, -1, 1, 1]) # Make it 4D.

features = {

"inputs": batch_inputs,

"targets": batch_targets,

"target_space_id": tf.constant(hparams.problems[0].target_space_id)

}

# Call the model.

predictions, _ = model(features)

nvars = len(tf.trainable_variables())

model(features) # Call again and check that reuse works.

self.assertEqual(nvars, len(tf.trainable_variables()))

# Having the graph, let's run it on some data.

with self.test_session() as sess:

sess.run(tf.global_variables_initializer())

inputs = "0 1 0"

targets = "0 1 0"

# Encode from raw string to numpy input array using problem encoders.

inputs_numpy = encoders["inputs"].encode(inputs)

targets_numpy = encoders["targets"].encode(targets)

# Feed the encoded inputs and targets and run session.

feed = {inputs_ph: inputs_numpy, targets_ph: targets_numpy}

np_predictions = sess.run(predictions, feed)

# Check that the result has the correct shape: batch x length x vocab_size

# where, for us, batch = 1, length = 3, vocab_size = 4.

self.assertEqual(np_predictions.shape, (1, 3, 1, 1, 4))

if __name__ == "__main__":

tf.test.main()