Internal change.

PiperOrigin-RevId: 522388337
google · Apr 6, 2023 · 93a5ec7 · 93a5ec7
1 parent 49a7693
commit 93a5ec7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 11 deletions.
diff --git a/struct2tensor/benchmarks/struct2tensor_benchmark.py b/struct2tensor/benchmarks/struct2tensor_benchmark.py
@@ -14,14 +14,6 @@
 r"""Benchmarks for struct2tensor.
 
 
-Usage:
-blaze run -c opt --dynamic_mode=off \
-  --run_under='perflab \
-  --constraints=arch=x86_64,platform_family=iota,platform_genus=sandybridge' \
-  //struct2tensor/benchmarks:struct2tensor_benchmark \
-  -- --notest_mode
-
-
 """
 
 from absl.testing import parameterized

diff --git a/struct2tensor/benchmarks/struct2tensor_benchmark_util.py b/struct2tensor/benchmarks/struct2tensor_benchmark_util.py
@@ -31,10 +31,10 @@
 
 flags.DEFINE_bool(
     "test_mode", False,
-    "if True, run all benchmarks with one iteration to make sure there is no error."
+    "if True, run all benchmarks with two iterations."
 )
 
-_BASE_DIR = "struct2tensor/benchmarks/testdata"
+_BASE_DIR = os.path.join(os.path.dirname(__file__), "testdata")
 
 
 class Struct2tensorBenchmarksBase(parameterized.TestCase):
@@ -65,7 +65,7 @@ def run_benchmarks(self, fn_name, get_benchmark_fn, fn_args, data_key):
           "Wall Time avg(ms)\tWall Time std\tUser CPU avg (ms)\t"
           "User CPU std\tSystem CPU avg (ms)\tSystem CPU std")
 
-    iterations = 10000
+    iterations = 1000
     # This is the number of iterations in a sample. We will benchmark the
     # total compute time per sample. And find std across all samples.
     sample_size = 100