def test_device_methods(self,
num_ps,
num_subnetworks,
want_variable_devices,
dedicate_parameter_servers=True):
with context.graph_mode():
x = tf.constant([[1., 0.]])
strategy = RoundRobinStrategy(
dedicate_parameter_servers=dedicate_parameter_servers)
strategy.config = ParameterServerConfig(num_ps)
variable_devices = []
for i in range(num_subnetworks):
with strategy.subnetwork_devices(num_subnetworks, i):
subnetwork = tf.keras.Sequential()
subnetwork.add(tf.keras.layers.Dense(4))
subnetwork.add(tf.keras.layers.Dense(3))
subnetwork(x)
variable_devices.append(
[w.op.device for w in subnetwork.weights])
self.assertEqual(want_variable_devices, variable_devices)
def test_methods(self, num_workers, num_subnetworks, drop_remainder,
want_should_build_ensemble, want_should_build_subnetwork,
want_should_train_subnetworks):
should_build_ensemble = []
should_build_subnetwork = []
should_train_subnetworks = []
for worker_index in range(num_workers):
strategy = RoundRobinStrategy(drop_remainder)
strategy.config = Config(num_workers, worker_index)
should_build_ensemble.append(
strategy.should_build_ensemble(num_subnetworks))
should_build_subnetwork.append([])
should_train_subnetworks.append(
strategy.should_train_subnetworks(num_subnetworks))
for subnetwork_index in range(num_subnetworks):
should_build_subnetwork[-1].append(
strategy.should_build_subnetwork(num_subnetworks,
subnetwork_index))
self.assertEqual(want_should_build_ensemble, should_build_ensemble)
self.assertEqual(want_should_build_subnetwork, should_build_subnetwork)
self.assertEqual(want_should_train_subnetworks,
should_train_subnetworks)