Beispiel #1
0
 def __init__(self, output_dir=None, default_flags=None, tpu=None):
     self.default_flags = default_flags or {}
     flag_methods = trainer.define_flags()
     super(NHNetBenchmark, self).__init__(output_dir=output_dir,
                                          default_flags=default_flags,
                                          flag_methods=flag_methods,
                                          tpu=tpu)
import os

from absl import flags
from absl.testing import parameterized
import tensorflow as tf

# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import strategy_combinations
# pylint: enable=g-direct-tensorflow-import
from official.nlp.nhnet import trainer
from official.nlp.nhnet import utils

FLAGS = flags.FLAGS
trainer.define_flags()


def all_strategy_combinations():
    return combinations.combine(
        distribution=[
            strategy_combinations.one_device_strategy,
            strategy_combinations.one_device_strategy_gpu,
            strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
            strategy_combinations.cloud_tpu_strategy,
        ],
        mode="eager",
    )


def get_trivial_data(config) -> tf.data.Dataset: