def test_mnist_training_cpu(self):
devices = tf.config.list_physical_devices("CPU")
tf.config.set_logical_device_configuration(
devices[0],
[
tf.config.LogicalDeviceConfiguration(),
] * 8,
)
mesh = mesh_util.create_mesh(devices=["CPU:%d" % i for i in range(8)],
mesh_dims=[("batch", 8)])
backend.enable_tf_random_generator()
# Needed by keras initializers.
tf_utils.set_random_seed(1337)
model = integration_test_utils.get_model_with_layout_map(
integration_test_utils.get_all_replicated_layout_map(mesh))
optimizer = optimizer_lib.Adam(learning_rate=0.001, mesh=mesh)
optimizer.build(model.trainable_variables)
train_losses = integration_test_utils.train_mnist_model_batch_sharded(
model,
optimizer,
mesh,
num_epochs=3,
steps_per_epoch=100,
global_batch_size=64,
)
# Make sure the losses are decreasing
self.assertEqual(train_losses, sorted(train_losses, reverse=True))
def setUp(self):
super(LayersTest, self).setUp()
backend.enable_tf_random_generator()
tf_utils.set_random_seed(1337)
global_ids = test_util.create_device_ids_array((2, 2))
local_device_ids = np.ravel(global_ids).tolist()
mesh_dict = {
'CPU':
dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
test_util.create_device_list((2, 2), 'CPU'))
}
self.mesh = self.configTestMesh(mesh_dict)
def setUp(self):
super().setUp()
backend.enable_tf_random_generator()
tf_utils.set_random_seed(1337)
global_ids = test_util.create_device_ids_array((2, 2))
local_device_ids = np.ravel(global_ids).tolist()
mesh_dict = {
"CPU":
dtensor.Mesh(
["X", "Y"],
global_ids,
local_device_ids,
test_util.create_device_list((2, 2), "CPU"),
)
}
self.mesh = self.configTestMesh(mesh_dict)
def setUp(self):
super(LayoutMapTest, self).setUp()
backend.enable_tf_random_generator()
tf_utils.set_random_seed(1337)
global_ids = test_util.create_device_ids_array((2, 2))
local_device_ids = np.ravel(global_ids).tolist()
mesh_dict = {
'CPU':
dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
test_util.create_device_list((2, 2), 'CPU'))
}
self.mesh = self.configTestMesh(mesh_dict)
self.layout_2d = dtensor.Layout.replicated(self.mesh, rank=2)
self.layout_1d = dtensor.Layout.replicated(self.mesh, rank=1)
self.sharded_2d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=2)
self.sharded_1d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=1)