def mnist_model(image, labels, mesh):
"""The model.
Args:
image: tf.Tensor with shape [batch, 28*28]
labels: a tf.Tensor with shape [batch] and dtype tf.int32
mesh: a mtf.Mesh
Returns:
logits: a tf.Tensor with shape [batch, 10]
loss: a mtf.Tensor with shape []
"""
batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
rows_dim = mtf.Dimension("rows", 28)
cols_dim = mtf.Dimension("cols", 28)
classes_dim = mtf.Dimension("classes", 10)
hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
x = mtf.import_tf_tensor(mesh, tf.reshape(image, [-1, 28, 28]),
mtf.Shape([batch_dim, rows_dim, cols_dim]))
h1 = mtf_layers.dense(
x, hidden_dim1, reduced_dims=[rows_dim, cols_dim],
activation=mtf.relu, name="hidden1")
h2 = mtf_layers.dense(
h1, hidden_dim2, activation=mtf.relu, name="hidden2")
logits = mtf_layers.dense(h2, classes_dim, name="logits")
if labels is None:
loss = None
else:
labels = mtf.import_tf_tensor(mesh, labels, mtf.Shape([batch_dim]))
loss = mtf_layers.softmax_cross_entropy_with_logits(
logits, mtf.one_hot(labels, classes_dim), classes_dim)
loss = mtf.reduce_mean(loss)
return logits, loss
def mnist_model(image, labels, mesh):
"""The model.
Args:
image: tf.Tensor with shape [batch, 28*28]
labels: a tf.Tensor with shape [batch] and dtype tf.int32
mesh: a mtf.Mesh
Returns:
logits: a tf.Tensor with shape [batch, 10]
loss: a mtf.Tensor with shape []
"""
batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
rows_dim = mtf.Dimension("rows", 28)
cols_dim = mtf.Dimension("cols", 28)
classes_dim = mtf.Dimension("classes", 10)
one_channel_dim = mtf.Dimension("one_channel", 1)
x = mtf.import_tf_tensor(mesh, tf.reshape(image, [-1, 28, 28]),
mtf.Shape([batch_dim, rows_dim, cols_dim]))
x = mtf.reshape(x, [batch_dim, rows_dim, cols_dim, one_channel_dim])
# add some convolutional layers to demonstrate that convolution works.
# TODO(noam): get spatially-partitioned convolution working.
fh_dim = mtf.Dimension("fh", 3)
fw_dim = mtf.Dimension("fw", 3)
filters1_dim = mtf.Dimension("filters1", 32)
filters2_dim = mtf.Dimension("filters2", 32)
kernel1 = mtf.get_variable(mesh, "kernel1",
[fh_dim, fw_dim, one_channel_dim, filters1_dim])
kernel2 = mtf.get_variable(mesh, "kernel2",
[fh_dim, fw_dim, filters1_dim, filters2_dim])
f1 = mtf.relu(mtf.conv2d(x, kernel1))
f2 = mtf.relu(mtf.conv2d(f1, kernel2))
x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
# add some fully-connected dense layers.
hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
h1 = mtf_layers.dense(x,
hidden_dim1,
reduced_dims=[rows_dim, cols_dim],
activation=mtf.relu,
name="hidden1")
h2 = mtf_layers.dense(h1, hidden_dim2, activation=mtf.relu, name="hidden2")
logits = mtf_layers.dense(h2, classes_dim, name="logits")
if labels is None:
loss = None
else:
labels = mtf.import_tf_tensor(mesh, labels, mtf.Shape([batch_dim]))
loss = mtf_layers.softmax_cross_entropy_with_logits(
logits, mtf.one_hot(labels, classes_dim), classes_dim)
loss = mtf.reduce_mean(loss)
return logits, loss
def testGraph(self):
graph = mtf.Graph()
self.assertLen(graph.operations, 0)
self.assertLen(graph.tensors, 0)
self.assertLen(graph.trainable_variables, 0)
self.assertLen(graph.all_variables, 0)
mesh = mtf.Mesh(graph, "mesh_test")
_ = mtf.import_tf_tensor(mesh,
tf_tensor=tf.constant(0.),
shape=mtf.Shape([]))
self.assertLen(graph.operations, 1)
self.assertLen(graph.tensors, 1)
self.assertLen(graph.trainable_variables, 0)
self.assertLen(graph.all_variables, 0)
_ = mtf.get_variable(mesh, "variable_0", mtf.Shape([]), trainable=True)
self.assertLen(graph.operations, 2)
self.assertLen(graph.tensors, 2)
self.assertLen(graph.trainable_variables, 1)
self.assertLen(graph.all_variables, 1)
_ = mtf.get_variable(mesh,
"variable_1",
mtf.Shape([]),
trainable=False)
self.assertLen(graph.operations, 3)
self.assertLen(graph.tensors, 3)
self.assertLen(graph.trainable_variables, 1)
self.assertLen(graph.all_variables, 2)
def testLayerNorm(self):
batch = 2
channels = 3
inputs = tf.random_normal([batch, channels])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
channels_dim = mtf.Dimension("channels", channels)
mtf_inputs = mtf.import_tf_tensor(
mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
mtf_outputs = mtf_layers.layer_norm(mtf_inputs,
dim=channels_dim)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(
shape=[], layout={}, devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
expected_outputs = common_layers.layer_norm(inputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual, expected = sess.run([actual_outputs, expected_outputs])
self.assertEqual(actual.shape, expected.shape)
def testDense(self, units, use_bias):
batch = 2
channels = 3
inputs = tf.random_normal([batch, channels])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
channels_dim = mtf.Dimension("channels", channels)
depth_dim = mtf.Dimension("depth", units)
mtf_inputs = mtf.import_tf_tensor(
mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
mtf_outputs = mtf_layers.dense(mtf_inputs,
output_dim=depth_dim,
reduced_dims=[channels_dim],
activation=mtf.relu,
use_bias=use_bias)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(
shape=[], layout={}, devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
expected_outputs = tf.keras.layers.Dense(units=units,
activation=tf.nn.relu,
use_bias=use_bias)(inputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual, expected = sess.run([actual_outputs, expected_outputs])
self.assertEqual(actual.shape, expected.shape)
def testDenseReluDense(self):
batch = 2
channels = 3
hidden = 5
inputs = tf.random_normal([batch, channels])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
channels_dim = mtf.Dimension("channels", channels)
hidden_dim = mtf.Dimension("hidden", hidden)
mtf_inputs = mtf.import_tf_tensor(
mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
mtf_outputs = mtf_layers.dense_relu_dense(mtf_inputs,
hidden_channels=hidden_dim)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(
shape=[], layout={}, devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual = sess.run(actual_outputs)
self.assertEqual(actual.shape, inputs.shape)
def testWeightsNonzero(self):
inputs = tf.constant([[3, 1, 0], [1, 0, 0]])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", inputs.shape.as_list()[0])
channels_dim = mtf.Dimension("channels", inputs.shape.as_list()[1])
mtf_inputs = mtf.import_tf_tensor(mesh,
inputs,
shape=mtf.Shape(
[batch_dim, channels_dim]))
mtf_outputs = mtf_layers.weights_nonzero(mtf_inputs)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(shape=[],
layout={},
devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
expected_outputs = common_layers.weights_nonzero(inputs)
tf_group = lowering.copy_masters_to_slices()
self.evaluate(tf_group)
actual, expected = self.evaluate([actual_outputs, expected_outputs])
self.assertAllEqual(actual, expected)
def testDotProductAttention(self, batch, heads, length_q, length_kv,
depth_k, depth_v):
query = tf.random_normal([batch, heads, length_q, depth_k])
key = tf.random_normal([batch, heads, length_kv, depth_k])
value = tf.random_normal([batch, heads, length_kv, depth_v])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
heads_dim = mtf.Dimension("heads", heads)
length_q_dim = mtf.Dimension("length_q", length_q)
length_kv_dim = mtf.Dimension("length_kv", length_kv)
depth_k_dim = mtf.Dimension("depth_k", depth_k)
depth_v_dim = mtf.Dimension("depth_v", depth_v)
mtf_query = mtf.import_tf_tensor(
mesh,
query,
shape=mtf.Shape([batch_dim, heads_dim, length_q_dim, depth_k_dim]))
mtf_key = mtf.import_tf_tensor(
mesh,
key,
shape=mtf.Shape([batch_dim, heads_dim, length_kv_dim,
depth_k_dim]))
mtf_value = mtf.import_tf_tensor(
mesh,
value,
shape=mtf.Shape([batch_dim, heads_dim, length_kv_dim,
depth_v_dim]))
mtf_outputs = mtf_layers.dot_product_attention(mtf_query,
mtf_key,
mtf_value,
mask=None)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(shape=[],
layout={},
devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual = sess.run(actual_outputs)
self.assertEqual(actual.shape, (batch, heads, length_q, depth_v))
def testMaskedLocalAttention1D(self, batch, length, io_channels,
kv_channels, heads, block_length):
length_q = length
length_m = length
query = tf.random_normal([batch, length_q, io_channels])
memory = tf.random_normal([batch, length_m, io_channels])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
length_q_dim = mtf.Dimension("length_q", length_q)
length_m_dim = mtf.Dimension("length_m", length_m)
io_channels_dim = mtf.Dimension("io_channels", io_channels)
kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
heads_dim = mtf.Dimension("heads", heads)
mtf_query = mtf.import_tf_tensor(
mesh,
query,
shape=mtf.Shape([batch_dim, length_q_dim, io_channels_dim]))
mtf_memory = mtf.import_tf_tensor(
mesh,
memory,
shape=mtf.Shape([batch_dim, length_m_dim, io_channels_dim]))
mtf_outputs = mtf_layers.masked_local_attention_1d(
mtf_query,
mtf_memory,
kv_channels=kv_channels_dim,
heads=heads_dim,
block_length=block_length)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(shape=[],
layout={},
devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual = sess.run(actual_outputs)
self.assertEqual(actual.shape, (batch, length_q, io_channels))
def toy_model(features, mesh):
"""A toy model implemented by mesh tensorlfow."""
batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
hidden_dim = mtf.Dimension('hidden', FLAGS.hidden_size)
io_dim = mtf.Dimension('io', FLAGS.io_size)
x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
h = mtf_layers.dense(x, hidden_dim, name='layer1', use_bias=False)
y = mtf_layers.dense(h, io_dim, name='layer2', use_bias=False)
loss = mtf.reduce_sum(mtf.square(y - x))
return y, loss
def testLowering(self):
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
inputs = tf.constant(0.)
mtf_inputs = mtf.import_tf_tensor(mesh,
tf_tensor=inputs,
shape=mtf.Shape([]))
mesh_impl = placement_mesh_impl.PlacementMeshImpl(
shape=[], layout={}, devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
outputs = lowering.export_to_tf_tensor(mtf_inputs)
inputs_value, outputs_value = self.evaluate([inputs, outputs])
self.assertEqual(inputs_value, outputs_value)
# Check that methods run without error.
_ = lowering.copy_masters_to_slices()
_ = lowering.copy_slices_to_masters()
def testMultiheadAttention(self, kv_channels, heads):
batch = 2
length = 8
channels = 3
query = tf.random_normal([batch, length, channels])
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", batch)
length_dim = mtf.Dimension("length", length)
channels_dim = mtf.Dimension("channels", channels)
kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
heads_dim = mtf.Dimension("heads", heads)
mtf_query = mtf.import_tf_tensor(
mesh,
query,
shape=mtf.Shape([batch_dim, length_dim, channels_dim]))
mtf_outputs = mtf_layers.multihead_attention(
mtf_query,
memory_antecedent=None,
mask=None,
kv_channels=kv_channels_dim,
heads=heads_dim)
mesh_impl = placement_mesh_impl.PlacementMeshImpl(shape=[],
layout={},
devices=[""])
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
tf_group = lowering.copy_masters_to_slices()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
sess.run(tf_group)
actual = sess.run(actual_outputs)
self.assertEqual(actual.shape, query.shape)
def estimator_model_fn(cls,
hparams,
features,
labels,
mode,
config=None,
params=None,
decode_hparams=None,
use_tpu=False,
xla_compile=False):
del xla_compile
hparams = copy.deepcopy(hparams)
hparams.use_tpu = use_tpu
# merge decode_hparams into hparams if present
if mode == tf.estimator.ModeKeys.PREDICT and decode_hparams is not None:
for k, v in six.iteritems(decode_hparams.values()):
if hasattr(hparams, k) and getattr(hparams, k) != v:
tf.logging.warning(
"Overriding hparams.%s with %s from decode_hparams" %
(k, v))
setattr(hparams, k, v)
# Instantiate model
data_parallelism = None
if not use_tpu and config:
data_parallelism = config.data_parallelism
model = cls(hparams,
mode,
data_parallelism=data_parallelism,
decode_hparams=decode_hparams)
global_step = tf.train.get_global_step()
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
layout_rules = mtf.convert_to_layout_rules(hparams.layout)
if use_tpu:
mesh_devices = [""] * mesh_shape.size
mesh_impl = simd_mesh_impl.SimdMeshImpl(
mesh_shape, layout_rules, mesh_devices,
params["context"].device_assignment)
else:
if len(data_parallelism.ps_devices) == 1:
mesh_devices = [""] * mesh_shape.size
else:
assert len(data_parallelism.ps_devices) == mesh_shape.size
mesh_devices = data_parallelism.ps_devices
mesh_impl = placement_mesh_impl.PlacementMeshImpl(
mesh_shape, layout_rules, mesh_devices)
# PREDICT mode
if mode == tf.estimator.ModeKeys.PREDICT:
return model.estimator_spec_predict(features, mesh, mesh_impl,
use_tpu)
logits, loss = model.mtf_model_fn(features, mesh)
if use_tpu and logits is not None:
logits = mtf.anonymize(logits)
# TRAIN mode
if mode == tf.estimator.ModeKeys.TRAIN:
var_grads = mtf.gradients(
[loss], [v.outputs[0] for v in graph.trainable_variables])
lr = learning_rate.learning_rate_schedule(hparams)
mtf_lr = mtf.import_tf_tensor(
mesh, tf.convert_to_tensor(lr, dtype=tf.float32),
mtf.Shape([]))
optimizer = mtf_optimize.make_optimizer(hparams, mtf_lr)
update_ops = []
for grad, var in zip(var_grads, graph.trainable_variables):
update_ops.extend(optimizer.apply_grad(grad, var))
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
tf_loss = lowering.export_to_tf_tensor(loss)
tf_loss = tf.to_float(tf_loss)
if logits and mode != tf.estimator.ModeKeys.TRAIN:
tf_logits = lowering.export_to_tf_tensor(logits)
if mode == tf.estimator.ModeKeys.TRAIN:
tf_update_ops = [
lowering.lowered_operation(op) for op in update_ops
]
tf_update_ops.append(tf.assign_add(global_step, 1))
# tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
train_op = tf.group(tf_update_ops)
with mtf_utils.outside_all_rewrites():
# Copy master variables to slices. Must be called first.
restore_hook = mtf.MtfRestoreHook(lowering)
saver = tf.train.Saver(tf.global_variables(),
sharded=True,
max_to_keep=10,
keep_checkpoint_every_n_hours=2,
defer_build=False,
save_relative_paths=True)
tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
saver_listener = mtf.MtfCheckpointSaverListener(lowering)
saver_hook = tf.train.CheckpointSaverHook(
hparams.model_dir,
save_steps=1000,
saver=saver,
listeners=[saver_listener])
# EVAL mode
if mode == tf.estimator.ModeKeys.EVAL:
tf_logits = lowering.export_to_tf_tensor(logits)
return model.estimator_spec_eval(features, tf_logits, labels,
tf_loss, restore_hook, use_tpu)
if use_tpu:
_remove_summaries()
return tpu_estimator.TPUEstimatorSpec(
mode=tf.estimator.ModeKeys.TRAIN,
loss=tf_loss,
train_op=train_op,
training_hooks=[restore_hook, saver_hook])
else:
return tf.estimator.EstimatorSpec(
tf.estimator.ModeKeys.TRAIN,
loss=tf_loss,
train_op=train_op,
training_chief_hooks=[restore_hook, saver_hook])
def mtf_model_fn(self, features, mesh):
features = copy.copy(features)
tf.logging.info("features = %s" % features)
hparams = self._hparams
activation_dtype = self.set_activation_type()
is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
# Declare all the dimensions
batch_dim = mtf.Dimension("batch", hparams.batch_size)
hidden_dim = mtf.Dimension("hidden", hparams.hidden_size)
filter_h_dim = mtf.Dimension("filter_height", 7)
filter_w_dim = mtf.Dimension("filter_width", 7)
filters = mtf.Dimension("filters", hparams.filter_sizes[0])
rows_dim = mtf.Dimension("rows_size", 32)
cols_dim = mtf.Dimension("cols_size", 96)
row_blocks_dim = mtf.Dimension("row_blocks", hparams.row_blocks)
col_blocks_dim = mtf.Dimension("col_blocks", hparams.col_blocks)
classes_dim = mtf.Dimension("classes", 10)
one_channel_dim = mtf.Dimension("one_channel", 1)
inputs = features["inputs"]
x = mtf.import_tf_tensor(
mesh,
tf.reshape(inputs, [
hparams.batch_size, hparams.row_blocks,
hparams.rows_size // hparams.row_blocks, hparams.col_blocks,
hparams.num_channels * hparams.cols_size // hparams.col_blocks,
1
]),
mtf.Shape([
batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim,
one_channel_dim
]))
x = mtf.transpose(x, [
batch_dim, row_blocks_dim, col_blocks_dim, rows_dim, cols_dim,
one_channel_dim
])
x = mtf.to_float(x)
initial_filters = mtf.get_variable(
mesh, "init_filters",
mtf.Shape([filter_h_dim, filter_w_dim, one_channel_dim, filters]))
x = mtf.conv2d_with_blocks(x,
initial_filters,
strides=[1, 1, 1, 1],
padding="SAME",
h_blocks_dim=None,
w_blocks_dim=col_blocks_dim)
x = batch_norm_relu(x, is_training)
# Conv blocks
# [ self attention - ffn - residual + dropout] x n
for layer in range(hparams.num_layers):
layer_name = "block_layer_%d" % layer
with tf.variable_scope(layer_name):
# Residual block layer
x = block_layer(inputs=x,
filters=hparams.filter_sizes[0],
blocks=hparams.layer_sizes[0],
strides=[1, 1, 1, 1],
is_training=is_training,
name="block_layer1",
row_blocks_dim=None,
col_blocks_dim=None)
x = block_layer(inputs=x,
filters=hparams.filter_sizes[1],
blocks=hparams.layer_sizes[1],
strides=[1, 2, 2, 1],
is_training=is_training,
name="block_layer2",
row_blocks_dim=None,
col_blocks_dim=None)
x = block_layer(inputs=x,
filters=hparams.filter_sizes[2],
blocks=hparams.layer_sizes[2],
strides=[1, 2, 2, 1],
is_training=is_training,
name="block_layer3",
row_blocks_dim=None,
col_blocks_dim=None)
# Calculate the logits and loss.
out = x
outputs = mtf_layers.dense(out,
hidden_dim,
reduced_dims=out.shape.dims[-5:],
activation=mtf.relu,
name="dense")
# We assume fixed vocab size for targets
labels = tf.squeeze(tf.to_int32(features["targets"]), [2, 3])
labels = mtf.import_tf_tensor(mesh,
tf.reshape(labels, [hparams.batch_size]),
mtf.Shape([batch_dim]))
logits = mtf_layers.dense(outputs, classes_dim, name="logits")
soft_targets = mtf.one_hot(labels, classes_dim, dtype=activation_dtype)
loss = mtf_layers.softmax_cross_entropy_with_logits(
logits, soft_targets, classes_dim)
# Reshape logits so it doesn't break inside t2t.
logits = mtf.reshape(
logits, mtf.Shape([batch_dim, one_channel_dim, classes_dim]))
loss = mtf.reduce_mean(loss)
return logits, loss
def import_to_batch_by_length(x, name):
return mtf.import_tf_tensor(mesh,
x,
mtf.Shape([batch_dim, length_dim]),
name=name)
def mnist_model(image, labels, mesh):
"""The model.
Args:
image: tf.Tensor with shape [batch, 28*28]
labels: a tf.Tensor with shape [batch] and dtype tf.int32
mesh: a mtf.Mesh
Returns:
logits: a tf.Tensor with shape [batch, 10]
loss: a mtf.Tensor with shape []
"""
batch_dim = mtf.Dimension("batch", FLAGS.batch_size)
row_blocks_dim = mtf.Dimension("row_blocks", 4)
col_blocks_dim = mtf.Dimension("col_blocks", 4)
rows_dim = mtf.Dimension("rows_size", 7)
cols_dim = mtf.Dimension("cols_size", 7)
classes_dim = mtf.Dimension("classes", 10)
one_channel_dim = mtf.Dimension("one_channel", 1)
x = mtf.import_tf_tensor(
mesh, tf.reshape(image, [FLAGS.batch_size, 4, 7, 4, 7, 1]),
mtf.Shape([
batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim,
one_channel_dim
]))
x = mtf.transpose(x, [
batch_dim, row_blocks_dim, col_blocks_dim, rows_dim, cols_dim,
one_channel_dim
])
# add some convolutional layers to demonstrate that convolution works.
fh_dim = mtf.Dimension("fh", 9)
fw_dim = mtf.Dimension("fw", 9)
filters1_dim = mtf.Dimension("filters1", 16)
filters2_dim = mtf.Dimension("filters2", 16)
kernel1 = mtf.get_variable(mesh, "kernel1",
[fh_dim, fw_dim, one_channel_dim, filters1_dim])
kernel2 = mtf.get_variable(mesh, "kernel2",
[fh_dim, fw_dim, filters1_dim, filters2_dim])
f1 = mtf.relu(
mtf.conv2d_with_blocks(x,
kernel1,
strides=[1, 1, 1, 1],
padding="SAME",
h_blocks_dim=row_blocks_dim,
w_blocks_dim=col_blocks_dim))
f2 = mtf.relu(
mtf.conv2d_with_blocks(f1,
kernel2,
strides=[1, 1, 1, 1],
padding="SAME",
h_blocks_dim=row_blocks_dim,
w_blocks_dim=col_blocks_dim))
x = mtf.reduce_mean(f2, reduced_dim=filters2_dim)
# add some fully-connected dense layers.
hidden_dim1 = mtf.Dimension("hidden1", FLAGS.hidden_size)
hidden_dim2 = mtf.Dimension("hidden2", FLAGS.hidden_size)
h1 = mtf_layers.dense(x,
hidden_dim1,
reduced_dims=x.shape.dims[-4:],
activation=mtf.relu,
name="hidden1")
h2 = mtf_layers.dense(h1, hidden_dim2, activation=mtf.relu, name="hidden2")
logits = mtf_layers.dense(h2, classes_dim, name="logits")
if labels is None:
loss = None
else:
labels = mtf.import_tf_tensor(mesh,
tf.reshape(labels, [FLAGS.batch_size]),
mtf.Shape([batch_dim]))
loss = mtf_layers.softmax_cross_entropy_with_logits(
logits, mtf.one_hot(labels, classes_dim), classes_dim)
loss = mtf.reduce_mean(loss)
return logits, loss
