Skip to content

GitLab

Commit ee016fb0 authored by David Chen's avatar David Chen Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 265122309
parent 1d610ef9
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 2 deletions
official/transformer/v2/transformer_benchmark.py
View file @ ee016fb0
...@@ -58,8 +58,8 @@ class TransformerBenchmark(PerfZeroBenchmark): ...@@ -58,8 +58,8 @@ class TransformerBenchmark(PerfZeroBenchmark):
EN2DE_2014_BLEU_DATA_DIR_NAME, EN2DE_2014_BLEU_DATA_DIR_NAME,
'newstest2014.de') 'newstest2014.de')
default_flags['train_steps'] = 200 if default_flags is None:
default_flags['log_steps'] = 10 default_flags = {}
default_flags['data_dir'] = self.train_data_dir default_flags['data_dir'] = self.train_data_dir
default_flags['vocab_file'] = self.vocab_file default_flags['vocab_file'] = self.vocab_file
...@@ -628,6 +628,8 @@ class TransformerBaseKerasBenchmarkReal(TransformerKerasBenchmark): ...@@ -628,6 +628,8 @@ 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['log_steps'] = 10
super(TransformerBaseKerasBenchmarkReal, self).__init__( super(TransformerBaseKerasBenchmarkReal, self).__init__(
output_dir=output_dir, default_flags=def_flags, output_dir=output_dir, default_flags=def_flags,
...@@ -640,6 +642,8 @@ class TransformerBigKerasBenchmarkReal(TransformerKerasBenchmark): ...@@ -640,6 +642,8 @@ 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['log_steps'] = 10
super(TransformerBigKerasBenchmarkReal, self).__init__( super(TransformerBigKerasBenchmarkReal, self).__init__(
output_dir=output_dir, default_flags=def_flags, output_dir=output_dir, default_flags=def_flags,
......
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