class_docstr: Tests that are common to all Allen Brain Atlas problems.
class_docstr_tok: ["Tests","that","are","common","to","all","Allen","Brain","Atlas","problems","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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),
              )
  
imports: import os
  import shutil
  import tempfile
  import numpy
  from tensor2tensor.data_generators import allen_brain
  from tensor2tensor.models import image_transformer_2d
  from tensor2tensor.utils import contrib
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 4 for calling the TestAllenBrain obj's underlying member method code with example usage: obj.test_transformer2d_single_step_e2e() and returns: losses