class_docstr: Testing usage of quantization ranges inserted in graph.
class_docstr_tok: ["Testing","usage","of","quantization","ranges","inserted","in","graph","."]
docstr: Gets the frozen mnist GraphDef.
  
  Args:
    use_trt: whether use TF-TRT to convert the graph.
    max_batch_size: the max batch size to apply during TF-TRT conversion.
    model_dir: the model directory to load the checkpoints.
  
  Returns:
    The frozen mnist GraphDef.
code: def _GetGraphDef(self, use_trt, max_batch_size, model_dir):
      """Gets the frozen mnist GraphDef.
  
      Args:
        use_trt: whether use TF-TRT to convert the graph.
        max_batch_size: the max batch size to apply during TF-TRT conversion.
        model_dir: the model directory to load the checkpoints.
  
      Returns:
        The frozen mnist GraphDef.
      """
      graph = ops.Graph()
      with self.session(graph=graph) as sess:
          with graph.device("/GPU:0"):
              x = array_ops.placeholder(
                  shape=(None, 28, 28, 1),
                  dtype=dtypes.float32,
                  name=INPUT_NODE_NAME,
              )
              self._BuildGraph(x)
          self._LoadWeights(model_dir, sess)
          # Freeze
          graph_def = (
              convert_to_constants.convert_variables_to_constants(
                  sess,
                  sess.graph_def,
                  output_node_names=[OUTPUT_NODE_NAME],
              )
          )
      # Convert with TF-TRT
      if use_trt:
          logging.info(
              "Number of nodes before TF-TRT conversion: %d",
              len(graph_def.node),
          )
          converter = trt_convert.TrtGraphConverter(
              input_graph_def=graph_def,
              nodes_denylist=[OUTPUT_NODE_NAME],
              max_batch_size=max_batch_size,
              precision_mode="INT8",
              max_workspace_size_bytes=(
                  trt_convert.DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES
              ),
              minimum_segment_size=2,
              use_calibration=False,
          )
          graph_def = converter.convert()
          logging.info(
              "Number of nodes after TF-TRT conversion: %d",
              len(graph_def.node),
          )
          num_engines = len(
              [1 for n in graph_def.node if str(n.op) == "TRTEngineOp"]
          )
          self.assertEqual(1, num_engines)
      return graph_def
  
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 3 for calling the QuantizationAwareTrainingMNISTTest obj's underlying member method code with example usage: obj._GetGraphDef(use_trt, max_batch_size, model_dir) and returns: graph_def