--no_tf32 added to the benchmark code

Author: DEKHTIARJonathan

tftrt/examples/benchmark_args.py

@@ -133,6 +133,13 @@ def __init__(self):
             "least as large as the number of samples in the dataset."
         )
 
+        self._add_bool_argument(
+            name="no_tf32",
+            default=False,
+            required=False,
+            help="If set to True, the benchmark will force not using TF32."
+        )
+
         self._add_bool_argument(
             name="use_xla",
             default=False,

tftrt/examples/benchmark_runner.py

@@ -61,6 +61,10 @@ def __init__(self, args):
             print("[Benchmark] - Activating XLA JIT Auto Clustering")
             os.environ["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"
 
+        if args.no_tf32:
+            print("[Benchmark] - Deactivating the use of TF32 format")
+            os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
+
         logging.getLogger("tensorflow").setLevel(logging.INFO)
         logging.disable(logging.WARNING)
 

tftrt/examples/image_classification/base_run_inference.sh

@@ -8,8 +8,6 @@ DATA_DIR=""
 MODEL_DIR=""
 
 # Default Argument Values
-NVIDIA_TF32_OVERRIDE=""
-
 BYPASS_ARGUMENTS=""
 
 # Loop through arguments and process them
@@ -20,10 +18,6 @@ do
         MODEL_NAME="${arg#*=}"
         shift # Remove --model_name from processing
         ;;
-        --no_tf32)
-        NVIDIA_TF32_OVERRIDE="NVIDIA_TF32_OVERRIDE=0"
-        shift # Remove --no_tf32 from processing
-        ;;
         --data_dir=*)
         DATA_DIR="${arg#*=}"
         shift # Remove --data_dir= from processing
@@ -100,8 +94,6 @@ echo ""
 echo "[*] DATA_DIR: ${DATA_DIR}"
 echo "[*] MODEL_DIR: ${MODEL_DIR}"
 echo ""
-echo "[*] NVIDIA_TF32_OVERRIDE: ${NVIDIA_TF32_OVERRIDE}"
-echo ""
 # Custom Image Classification Task Flags
 echo "[*] INPUT_SIZE: ${INPUT_SIZE}"
 echo "[*] PREPROCESS_METHOD: ${PREPROCESS_METHOD}"
@@ -151,10 +143,7 @@ BENCH_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 cd ${BENCH_DIR}
 
 # Execute the example
-
-PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
-
-COMMAND="${PREPEND_COMMAND} python image_classification.py \
+COMMAND="python image_classification.py \
     --data_dir ${DATA_DIR} \
     --calib_data_dir ${DATA_DIR} \
     --input_saved_model_dir ${INPUT_SAVED_MODEL_DIR} \

tftrt/examples/object_detection/base_run_inference.sh

@@ -8,8 +8,6 @@ DATA_DIR=""
 MODEL_DIR=""
 
 # Default Argument Values
-NVIDIA_TF32_OVERRIDE=""
-
 BYPASS_ARGUMENTS=""
 BATCH_SIZE=8
 
@@ -21,10 +19,6 @@ do
         MODEL_NAME="${arg#*=}"
         shift # Remove --model_name from processing
         ;;
-        --no_tf32)
-        NVIDIA_TF32_OVERRIDE="NVIDIA_TF32_OVERRIDE=0"
-        shift # Remove --no_tf32 from processing
-        ;;
         --batch_size=*)
         BATCH_SIZE="${arg#*=}"
         shift # Remove --batch_size= from processing
@@ -70,8 +64,6 @@ echo ""
 echo "[*] DATA_DIR: ${DATA_DIR}"
 echo "[*] MODEL_DIR: ${MODEL_DIR}"
 echo ""
-echo "[*] NVIDIA_TF32_OVERRIDE: ${NVIDIA_TF32_OVERRIDE}"
-echo ""
 # Custom Object Detection Task Flags
 echo "[*] BATCH_SIZE: ${BATCH_SIZE}"
 echo "[*] INPUT_SIZE: ${INPUT_SIZE}"
@@ -142,10 +134,7 @@ if [[ ${DEPENDENCIES_STATUS} != 0 ]]; then
 fi
 
 # Step 2: Execute the example
-
-PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
-
-COMMAND="${PREPEND_COMMAND} python object_detection.py \
+COMMAND="python object_detection.py \
     --data_dir ${VAL_DATA_DIR} \
     --calib_data_dir ${VAL_DATA_DIR} \
     --annotation_path ${ANNOTATIONS_DATA_FILE} \

tftrt/examples/transformers/base_run_inference.sh

@@ -7,14 +7,11 @@ MODEL_NAME=""
 MODEL_DIR=""
 
 # Default Argument Values
-NVIDIA_TF32_OVERRIDE=""
-
-# TODO: remove when real dataloader is implemented
-DATA_DIR="/tmp"
-
 BYPASS_ARGUMENTS=""
 BATCH_SIZE=32
 SEQ_LEN=128
+# TODO: remove when real dataloader is implemented
+DATA_DIR="/tmp"
 
 # Loop through arguments and process them
 for arg in "$@"
@@ -24,10 +21,6 @@ do
         MODEL_NAME="${arg#*=}"
         shift # Remove --model_name from processing
         ;;
-        --no_tf32)
-        NVIDIA_TF32_OVERRIDE="NVIDIA_TF32_OVERRIDE=0"
-        shift # Remove --no_tf32 from processing
-        ;;
         --batch_size=*)
         BATCH_SIZE="${arg#*=}"
         shift # Remove --batch_size= from processing
@@ -90,8 +83,6 @@ echo ""
 echo "[*] DATA_DIR: ${DATA_DIR}"
 echo "[*] MODEL_DIR: ${MODEL_DIR}"
 echo ""
-echo "[*] NVIDIA_TF32_OVERRIDE: ${NVIDIA_TF32_OVERRIDE}"
-echo ""
 # Custom Transormer Task Flags
 echo "[*] VOCAB_SIZE: ${VOCAB_SIZE}"
 echo "[*] SEQ_LEN: ${SEQ_LEN}"
@@ -142,10 +133,7 @@ BENCH_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )/" >/dev/null 2>&1 && pwd )"
 cd ${BENCH_DIR}
 
 # Execute the example
-
-PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
-
-COMMAND="${PREPEND_COMMAND} python transformers.py \
+COMMAND="python transformers.py \
     --data_dir ${DATA_DIR} \
     --calib_data_dir ${DATA_DIR} \
     --input_saved_model_dir ${INPUT_SAVED_MODEL_DIR} \