class_docstr: Testing usage of quantization ranges inserted in graph.
class_docstr_tok: ["Testing","usage","of","quantization","ranges","inserted","in","graph","."]
docstr: Trains or evaluates the model.
  
  Args:
    is_training: whether to train or evaluate the model. In training mode,
      quantization will be simulated where the quantize_and_dequantize_v2 are
      placed.
    use_trt: if true, use TRT INT8 mode for evaluation, which will perform
      real quantization. Otherwise use native TensorFlow which will perform
      simulated quantization. Ignored if is_training is True.
    batch_size: batch size.
    num_epochs: how many epochs to train. Ignored if is_training is False.
    model_dir: where to save or load checkpoint.
  
  Returns:
    The Estimator evaluation result.
code: def _Run(
      self, is_training, use_trt, batch_size, num_epochs, model_dir
  ):
      """Trains or evaluates the model.
  
      Args:
        is_training: whether to train or evaluate the model. In training mode,
          quantization will be simulated where the quantize_and_dequantize_v2 are
          placed.
        use_trt: if true, use TRT INT8 mode for evaluation, which will perform
          real quantization. Otherwise use native TensorFlow which will perform
          simulated quantization. Ignored if is_training is True.
        batch_size: batch size.
        num_epochs: how many epochs to train. Ignored if is_training is False.
        model_dir: where to save or load checkpoint.
  
      Returns:
        The Estimator evaluation result.
      """
  
      def _EvalInputFn():
          dataset = _GetDataSet(batch_size)
          iterator = dataset_ops.make_one_shot_iterator(dataset)
          features, labels = iterator.get_next()
          return features, labels
  
      def _TrainInputFn():
          dataset = tfds.load("mnist", split="train")
          dataset = dataset.shuffle(60000)
          dataset = dataset.map(
              map_func=_PreprocessFn, num_parallel_calls=8
          ).batch(batch_size=batch_size)
          dataset = dataset.repeat(count=num_epochs)
          iterator = dataset_ops.make_one_shot_iterator(dataset)
          features, labels = iterator.get_next()
          return features, labels
  
      def _ModelFn(features, labels, mode):
          if is_training:
              logits_out = self._BuildGraph(features)
          else:
              graph_def = self._GetGraphDef(
                  use_trt, batch_size, model_dir
              )
              logits_out = importer.import_graph_def(
                  graph_def,
                  input_map={INPUT_NODE_NAME: features},
                  return_elements=[OUTPUT_NODE_NAME + ":0"],
                  name="",
              )[0]
  
          loss = losses.sparse_softmax_cross_entropy(
              labels=labels, logits=logits_out
          )
          summary.scalar("loss", loss)
  
          classes_out = math_ops.argmax(
              logits_out, axis=1, name="classes_out"
          )
          accuracy = metrics.accuracy(
              labels=labels, predictions=classes_out, name="acc_op"
          )
          summary.scalar("accuracy", accuracy[1])
  
          if mode == ModeKeys.EVAL:
              return EstimatorSpec(
                  mode,
                  loss=loss,
                  eval_metric_ops={"accuracy": accuracy},
              )
          if mode == ModeKeys.TRAIN:
              optimizer = AdamOptimizer(learning_rate=1e-2)
              train_op = optimizer.minimize(
                  loss, global_step=get_global_step()
              )
              return EstimatorSpec(mode, loss=loss, train_op=train_op)
  
      config_proto = config_pb2.ConfigProto()
      config_proto.gpu_options.allow_growth = True
      estimator = Estimator(
          model_fn=_ModelFn,
          model_dir=model_dir if is_training else None,
          config=RunConfig(session_config=config_proto),
      )
  
      if is_training:
          estimator.train(_TrainInputFn)
      results = estimator.evaluate(_EvalInputFn)
      logging.info("accuracy: %s", str(results["accuracy"]))
      return results
  
imports: import os.path
  import tempfile
  import tensorflow_datasets
  from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.client import session
  from tensorflow.python.compiler.tensorrt import trt_convert
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.estimator.estimator import Estimator
  from tensorflow.python.estimator.model_fn import EstimatorSpec
  from tensorflow.python.estimator.model_fn import ModeKeys
  from tensorflow.python.estimator.run_config import RunConfig
  from tensorflow.python.framework import convert_to_constants
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.keras.metrics import Accuracy
  from tensorflow.python.layers import layers
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import metrics
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.saved_model import builder
  from tensorflow.python.saved_model import signature_constants
  from tensorflow.python.saved_model import tag_constants
  from tensorflow.python.saved_model import utils
  from tensorflow.python.saved_model import signature_def_utils
  from tensorflow.python.saved_model.load import load
  from tensorflow.python.summary import summary
  from tensorflow.python.training import saver
  from tensorflow.python.training.adam import AdamOptimizer
  from tensorflow.python.training.checkpoint_management import latest_checkpoint
  from tensorflow.python.training.training_util import get_global_step
label_desc: Use image node_id 4 for calling the QuantizationAwareTrainingMNISTTest obj's underlying member method code with example usage: obj._Run(is_training, use_trt, batch_size, num_epochs, model_dir) and returns: results, features,  labels, features,  labels, EstimatorSpec, EstimatorSpec