def setUp(self): super(AggregatePredictResultsTest, self).setUp() strategy_combinations.set_virtual_cpus_to_at_least(3) self.num_replica = 3 self.batch_size = 16 self.dense_shape = (2, 3) self.total_sample = 2 * self.batch_size mock_model = collections.namedtuple('Model', ['outputs']) self.mock_model = mock_model([1]) strategy = mirrored_strategy.MirroredStrategy( ['/cpu:0', '/cpu:1', '/cpu:2']) execution_function = lambda *inp: inp @def_function.function def predict_loop(batch): batch_result = strategy.experimental_run_v2( execution_function, batch) batch_result = dist_utils.unwrap_output_dict( strategy, batch_result, ModeKeys.PREDICT) # swap the order of replica 1 and 2, to mimic random order. batch_result[2], batch_result[1] = batch_result[1], batch_result[2] batch_result[5], batch_result[4] = batch_result[4], batch_result[5] return batch_result self.strategy = strategy self.predict_loop = predict_loop
def setUp(self): super(TestEstimatorDistributionStrategy, self).setUp() strategy_combinations.set_virtual_cpus_to_at_least(3) self._base_dir = os.path.join(self.get_temp_dir(), 'keras_to_estimator_strategy_test') gfile.MakeDirs(self._base_dir) self._config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir)
def testSetVirtualCPUsErrors(self): with self.assertRaises(ValueError): strategy_combinations.set_virtual_cpus_to_at_least(0) with self.assertRaisesRegexp(RuntimeError, "with 3 < 5 virtual CPUs"): strategy_combinations.set_virtual_cpus_to_at_least(5)
def testSetVirtualCPUsAgain(self): strategy_combinations.set_virtual_cpus_to_at_least(2) cpu_device = config.list_physical_devices("CPU")[0] self.assertLen(config.get_logical_device_configuration(cpu_device), 3)
def setUp(self): context._reset_context() # pylint: disable=protected-access # Need to call set_virtual_cpus_to_at_least() in setUp with the maximum # value needed in any test. strategy_combinations.set_virtual_cpus_to_at_least(3) super(VirtualDevicesTest, self).setUp()
def setUp(self): strategy_combinations.set_virtual_cpus_to_at_least(3) super(AutoCastVariableTest, self).setUp()
def setUp(self): super(MirroredFunctionStrategyTest, self).setUp() strategy_combinations.set_virtual_cpus_to_at_least(3) self._strategy = mirrored_function_strategy.MirroredFunctionStrategy( devices=("/cpu:1", "/cpu:2"))
def setUp(self): strategy_combinations.set_virtual_cpus_to_at_least(3) super(LossUtilitiesTest, self).setUp()
def setUp(self): # Need to call set_virtual_cpus_to_at_least() in setUp with the maximum # value needed in any test. strategy_combinations.set_virtual_cpus_to_at_least(3) super(StrategyCombinationsTest, self).setUp()
import pytest import tensorflow as tf from packaging import version from tensorflow.python.distribute import strategy_combinations from tensorflow.python.eager import context from larq import context as lq_context if version.parse(tf.__version__) >= version.parse("1.15"): strategy_combinations.set_virtual_cpus_to_at_least(3) distributed_devices = ["/cpu:1", "/cpu:2"] else: distributed_devices = ["/cpu:0"] @pytest.fixture def eager_mode(): """pytest fixture for running test in eager mode""" with context.eager_mode(): yield @pytest.fixture def graph_mode(): """pytest fixture for running test in graph mode""" with context.graph_mode(): with tf.compat.v1.Session().as_default(): yield tf.keras.backend.clear_session()