Skip to content

GitLab

Commit 82ba6174 authored by Adrian Kuegel's avatar Adrian Kuegel Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 265856020
parent 85956b16
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
official/transformer/v2/transformer_benchmark.py
View file @ 82ba6174
...@@ -43,6 +43,7 @@ class TransformerBenchmark(PerfZeroBenchmark): ...@@ -43,6 +43,7 @@ class TransformerBenchmark(PerfZeroBenchmark):
def __init__(self, output_dir=None, default_flags=None, root_data_dir=None, def __init__(self, output_dir=None, default_flags=None, root_data_dir=None,
flag_methods=None): flag_methods=None):
assert tf.version.VERSION.startswith('2.')
self.train_data_dir = os.path.join(root_data_dir, self.train_data_dir = os.path.join(root_data_dir,
TRANSFORMER_EN2DE_DATA_DIR_NAME) TRANSFORMER_EN2DE_DATA_DIR_NAME)
...@@ -628,7 +629,7 @@ class TransformerBaseKerasBenchmarkReal(TransformerKerasBenchmark): ...@@ -628,7 +629,7 @@ class TransformerBaseKerasBenchmarkReal(TransformerKerasBenchmark):
def __init__(self, output_dir=TMP_DIR, root_data_dir=None, **kwargs): def __init__(self, output_dir=TMP_DIR, root_data_dir=None, **kwargs):
def_flags = {} def_flags = {}
def_flags['param_set'] = 'base' def_flags['param_set'] = 'base'
def_flags['train_steps'] = 200 def_flags['train_steps'] = 50
def_flags['log_steps'] = 10 def_flags['log_steps'] = 10
super(TransformerBaseKerasBenchmarkReal, self).__init__( super(TransformerBaseKerasBenchmarkReal, self).__init__(
...@@ -642,7 +643,7 @@ class TransformerBigKerasBenchmarkReal(TransformerKerasBenchmark): ...@@ -642,7 +643,7 @@ class TransformerBigKerasBenchmarkReal(TransformerKerasBenchmark):
def __init__(self, output_dir=TMP_DIR, root_data_dir=None, **kwargs): def __init__(self, output_dir=TMP_DIR, root_data_dir=None, **kwargs):
def_flags = {} def_flags = {}
def_flags['param_set'] = 'big' def_flags['param_set'] = 'big'
def_flags['train_steps'] = 200 def_flags['train_steps'] = 50
def_flags['log_steps'] = 10 def_flags['log_steps'] = 10
super(TransformerBigKerasBenchmarkReal, self).__init__( super(TransformerBigKerasBenchmarkReal, self).__init__(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment