code: def main(_):
      tf.logging.set_verbosity(tf.logging.INFO)
      trainer_lib.set_random_seed(FLAGS.random_seed)
      usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
  
      # Create hparams
      hparams = create_hparams()
      hparams.force_full_predict = True
      batch_size = hparams.batch_size
  
      # Iterating over dev/test partition of the data.
      # Change the data partition if necessary.
      dataset = registry.problem(FLAGS.problem).dataset(
          tf_estimator.ModeKeys.PREDICT,
          shuffle_files=False,
          hparams=hparams,
      )
  
      dataset = dataset.batch(batch_size, drop_remainder=True)
      data = dataset.make_one_shot_iterator().get_next()
      input_data = dict(
          (k, data[k]) for k in data.keys() if k.startswith("input")
      )
  
      # Creat model
      model_cls = registry.model(FLAGS.model)
      model = model_cls(hparams, tf_estimator.ModeKeys.PREDICT)
      prediction_ops = model.infer(input_data)
  
      # Confusion Matrix
      nr = hparams.problem.num_rewards
      cm_per_frame = np.zeros((nr, nr), dtype=np.uint64)
      cm_next_frame = np.zeros((nr, nr), dtype=np.uint64)
  
      saver = tf.train.Saver()
      with tf.train.SingularMonitoredSession() as sess:
          # Load latest checkpoint
          ckpt = tf.train.get_checkpoint_state(
              FLAGS.output_dir
          ).model_checkpoint_path
          saver.restore(sess.raw_session(), ckpt)
  
          counter = 0
          while not sess.should_stop():
              counter += 1
              if counter % 1 == 0:
                  print(counter)
  
              # Predict next frames
              rew_pd, rew_gt = sess.run(
                  [
                      prediction_ops["target_reward"],
                      data["target_reward"],
                  ]
              )
  
              for i in range(batch_size):
                  cm_next_frame[rew_gt[i, 0, 0], rew_pd[i, 0, 0]] += 1
                  for gt, pd in zip(rew_gt[i], rew_pd[i]):
                      cm_per_frame[gt, pd] += 1
  
      print_confusion_matrix("Per-frame Confusion Matrix", cm_per_frame)
      print_confusion_matrix(
          "Next-frame Confusion Matrix", cm_next_frame
      )
  
file_path: tensor2tensor/tensor2tensor/utils/video/reward_confusion.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.bin.t2t_decoder import create_hparams
  from tensor2tensor.data_generators import problem
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import trainer_lib
  from tensor2tensor.utils import usr_dir
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) without return types