def test_transformer2d_single_step_e2e(self):
    """Minimal end-to-end test of training and eval on allen_brain_image2image.

    Notes:

      * Runs problem generate_data

      * Runs a single step of training

      * Runs model in eval mode to obtain a prediction and confirms the
        resulting shape.

        * TODO: Running this in predict mode crashes in my environment.
          Separately have seen predict mode not produce the right shape
          output tensors, as if .infer is still a wip.

    """

    problem_object = allen_brain.Img2imgAllenBrainDim8to32()

    with TemporaryDirectory() as tmp_dir:

      mock_raw_data(tmp_dir, raw_dim=256, num_images=100)

      with TemporaryDirectory() as data_dir:

        problem_object.generate_data(data_dir, tmp_dir)

        input_xy_dim = problem_object.input_dim
        target_xy_dim = problem_object.output_dim
        num_channels = problem_object.num_channels

        hparams = image_transformer_2d.img2img_transformer2d_tiny()
        hparams.data_dir = data_dir

        p_hparams = problem_object.get_hparams(hparams)

        model = image_transformer_2d.Img2imgTransformer(
            hparams, tf_estimator.ModeKeys.TRAIN, p_hparams
        )

        @tfe.implicit_value_and_gradients
        def loss_fn(features):
          _, losses = model(features)
          return losses["training"]

        batch_size = 1
        train_dataset = problem_object.dataset(Modes.TRAIN, data_dir)
        train_dataset = train_dataset.repeat(None).batch(batch_size)

        optimizer = tf.train.AdamOptimizer()

        example = tfe.Iterator(train_dataset).next()
        example["targets"] = tf.reshape(example["targets"],
                                        [batch_size,
                                         target_xy_dim,
                                         target_xy_dim,
                                         num_channels])
        _, gv = loss_fn(example)
        optimizer.apply_gradients(gv)

        model.set_mode(Modes.EVAL)
        dataset = problem_object.dataset(Modes.EVAL, data_dir)

        example = tfe.Iterator(dataset).next()
        example["inputs"] = tf.reshape(example["inputs"],
                                       [1,
                                        input_xy_dim,
                                        input_xy_dim,
                                        num_channels])
        example["targets"] = tf.reshape(example["targets"],
                                        [1,
                                         target_xy_dim,
                                         target_xy_dim,
                                         num_channels])

        predictions, _ = model(example)

        self.assertEqual(predictions.numpy().shape,
                         (1,
                          target_xy_dim,
                          target_xy_dim,
                          num_channels,
                          256))
Minimal end-to-end test of training and eval on allen_brain_image2image.

    Notes:

      * Runs problem generate_data

      * Runs a single step of training

      * Runs model in eval mode to obtain a prediction and confirms the
        resulting shape.

        * TODO: Running this in predict mode crashes in my environment.
          Separately have seen predict mode not produce the right shape
          output tensors, as if .infer is still a wip.

    
tensor2tensor/data_generators/allen_brain_test.py