Add in MLIR benchmarking for ResNet.

PiperOrigin-RevId: 321039601

Add in MLIR benchmarking for ResNet.
PiperOrigin-RevId: 321039601
be8c2556 · Allen Wang · A. Unique TensorFlower · 5df3d2e3 · be8c2556
Commit be8c2556 authored Jul 13, 2020 by Allen Wang Committed by A. Unique TensorFlower Jul 13, 2020
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 0 deletions

official/benchmark/resnet_ctl_imagenet_benchmark.py official/benchmark/resnet_ctl_imagenet_benchmark.py +9 -0

No files found.
--- a/official/benchmark/resnet_ctl_imagenet_benchmark.py
+++ b/official/benchmark/resnet_ctl_imagenet_benchmark.py
@@ -389,6 +389,15 @@ class Resnet50CtlBenchmarkBase(CtlBenchmark):
    FLAGS.dtype = 'bf16'
    self._run_and_report_benchmark()

+  @owner_utils.Owner('tf-graph-compiler')
+  def benchmark_2x2_tpu_bf16_mlir(self):
+    self._setup()
+    self._set_df_common()
+    FLAGS.batch_size = 1024
+    FLAGS.dtype = 'bf16'
+    tf.config.experimental.enable_mlir_bridge()
+    self._run_and_report_benchmark()
+
  def benchmark_4x4_tpu_bf16(self):
    self._setup()
    self._set_df_common()