def projection_shortcut(inputs, kernel):
"""Project identity branch."""
inputs = mtf.conv2d_with_blocks(inputs,
kernel,
strides=strides,
padding="SAME",
h_blocks_dim=None,
w_blocks_dim=col_blocks_dim)
return batch_norm_relu(inputs, is_training, relu=False)
def bottleneck_block(inputs,
filters,
is_training,
strides,
projection_shortcut=None,
row_blocks_dim=None,
col_blocks_dim=None):
"""Bottleneck block variant for residual networks with BN after convolutions.
Args:
inputs: a `mtf.Tensor` of shape
`[batch_dim, row_blocks, col_blocks, rows, cols, in_channels]`.
filters: `int` number of filters for the first two convolutions. Note
that the third and final convolution will use 4 times as many filters.
is_training: `bool` for whether the model is in training mode.
strides: `int` block stride. If greater than 1, this block will ultimately
downsample the input.
projection_shortcut: `function` to use for projection shortcuts (typically
a 1x1 convolution to match the filter dimensions). If None, no
projection is used and the input is passed as unchanged through the
shortcut connection.
row_blocks_dim: a mtf.Dimension, row dimension which is
spatially partitioned along mesh axis
col_blocks_dim: a mtf.Dimension, row dimension which is
spatially partitioned along mesh axis
Returns:
The output `Tensor` of the block.
"""
shortcut = inputs
filter_h_dim = mtf.Dimension("filter_height", 3)
filter_w_dim = mtf.Dimension("filter_width", 3)
one_h_dim = mtf.Dimension("filter_height", 1)
one_w_dim = mtf.Dimension("filter_width", 1)
if projection_shortcut is not None:
filters_dim = mtf.Dimension("filtersp", filters)
kernel = mtf.get_variable(
inputs.mesh, "kernel",
mtf.Shape(
[one_h_dim, one_w_dim, inputs.shape.dims[-1], filters_dim]))
shortcut = projection_shortcut(inputs, kernel)
# First conv block
filters1_dim = mtf.Dimension("filters1", filters)
kernel1 = mtf.get_variable(
inputs.mesh, "kernel1",
mtf.Shape([one_h_dim, one_w_dim, inputs.shape.dims[-1], filters1_dim]))
inputs = mtf.conv2d_with_blocks(inputs,
kernel1,
strides=[1, 1, 1, 1],
padding="SAME",
h_blocks_dim=None,
w_blocks_dim=col_blocks_dim)
# TODO(nikip): Add Dropout?
inputs = batch_norm_relu(inputs, is_training)
# Second conv block
filters2_dim = mtf.Dimension("filters2", filters)
kernel2 = mtf.get_variable(
inputs.mesh, "kernel2",
mtf.Shape([filter_h_dim, filter_w_dim, filters1_dim, filters2_dim]))
inputs = mtf.conv2d_with_blocks(inputs,
kernel2,
strides=[1, 1, 1, 1],
padding="SAME",
h_blocks_dim=row_blocks_dim,
w_blocks_dim=col_blocks_dim)
inputs = batch_norm_relu(inputs, is_training)
# Third wide conv filter block
filters3_dim = mtf.Dimension("filters3", filters)
filters3_kernel = mtf.get_variable(
inputs.mesh, "wide_kernel",
mtf.Shape([one_h_dim, one_w_dim, filters2_dim, filters3_dim]))
inputs = mtf.conv2d_with_blocks(inputs,
filters3_kernel,
strides,
padding="SAME",
h_blocks_dim=None,
w_blocks_dim=col_blocks_dim)
inputs = batch_norm_relu(inputs, is_training, relu=False)
# TODO(nikip): Maybe add residual with a projection?
return mtf.relu(inputs + mtf.rename_dimension(
shortcut, shortcut.shape.dims[-1].name, inputs.shape.dims[-1].name))
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 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