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 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": 0, # We run on the first problem here. "input_space_id": hparams.problems[0].input_space_id, "target_space_id": 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 testSliceNet(self): x = np.random.random_integers(0, high=255, size=(3, 5, 5, 3)) y = np.random.random_integers(0, high=9, size=(3, 5, 1, 1)) hparams = slicenet.slicenet_params1_tiny() hparams.add_hparam("data_dir", "") problem = registry.problem("image_cifar10") p_hparams = problem.get_hparams(hparams) hparams.problems = [p_hparams] with self.test_session() as session: features = { "inputs": tf.constant(x, dtype=tf.int32), "targets": tf.constant(y, dtype=tf.int32), "target_space_id": tf.constant(1, dtype=tf.int32), } model = slicenet.SliceNet(hparams, tf.estimator.ModeKeys.TRAIN, p_hparams) sharded_logits, _ = model.model_fn(features) logits = tf.concat(sharded_logits, 0) session.run(tf.global_variables_initializer()) res = session.run(logits) self.assertEqual(res.shape, (3, 1, 1, 1, 10))
def problem(name): return registry.problem(name)
def setUpClass(cls): tf.set_random_seed(1) cls.problem = registry.problem("test_problem") cls.data_dir = tempfile.gettempdir() cls.filepatterns = generate_test_data(cls.problem, cls.data_dir)