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()
