def bottom(self, inputs):
"""Transform input from data space to model space.
Perform the Xception "Entry flow", which consists of two convolutional
filter upscalings followed by three residually connected separable
convolution blocks.
Args:
inputs: A Tensor with shape [batch, ...]
Returns:
body_input: A Tensor with shape [batch, ?, ?, body_input_depth].
"""
with tf.variable_scope(self.name):
def xnet_resblock(x, filters, res_relu, name):
with tf.variable_scope(name):
y = common_layers.separable_conv_block(
x,
filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
first_relu=True,
padding="SAME",
force2d=True,
name="sep_conv_block")
y = common_layers.pool(y, (3, 3),
"MAX",
"SAME",
strides=(2, 2))
return y + common_layers.conv_block(x,
filters, [((1, 1),
(1, 1))],
padding="SAME",
strides=(2, 2),
first_relu=res_relu,
force2d=True,
name="res_conv0")
inputs = common_layers.standardize_images(inputs)
# TODO(lukaszkaiser): summaries here don't work in multi-problem case yet.
# tf.summary.image("inputs", inputs, max_outputs=2)
x = common_layers.conv_block(inputs,
32, [((1, 1), (3, 3))],
first_relu=False,
padding="SAME",
strides=(2, 2),
force2d=True,
name="conv0")
x = common_layers.conv_block(x,
64, [((1, 1), (3, 3))],
padding="SAME",
force2d=True,
name="conv1")
x = xnet_resblock(x, min(128, self._body_input_depth), True,
"block0")
x = xnet_resblock(x, min(256, self._body_input_depth), False,
"block1")
return xnet_resblock(x, self._body_input_depth, False, "block2")
def bytenet_internal(inputs, targets, hparams):
"""ByteNet, main step used for training."""
with tf.variable_scope("bytenet"):
# Flatten inputs and extend length by 50%.
inputs = tf.expand_dims(common_layers.flatten4d3d(inputs), axis=2)
extend_length = tf.to_int32(0.5 * tf.to_float(tf.shape(inputs)[1]))
inputs_shape = inputs.shape.as_list()
inputs = tf.pad(inputs, [[0, 0], [0, extend_length], [0, 0], [0, 0]])
inputs_shape[1] = None
inputs.set_shape(
inputs_shape) # Don't lose the other shapes when padding.
# Pad inputs and targets to be the same length, divisible by 50.
inputs, targets = common_layers.pad_to_same_length(
inputs, targets, final_length_divisible_by=50)
final_encoder = residual_dilated_conv(inputs, hparams.num_block_repeat,
"SAME", "encoder", hparams)
shifted_targets = common_layers.shift_right(targets)
kernel = (hparams.kernel_height, hparams.kernel_width)
decoder_start = common_layers.conv_block(
tf.concat([final_encoder, shifted_targets], axis=3),
hparams.hidden_size, [((1, 1), kernel)],
padding="LEFT")
return residual_dilated_conv(decoder_start, hparams.num_block_repeat,
"LEFT", "decoder", hparams)
def ae_compress(x, is_2d, hparams, name, reuse=None):
"""Compress, then AE."""
with tf.variable_scope(name, reuse=reuse):
cur = compress(x, None, is_2d, hparams, "compress")
# Convolve and ReLu to get state.
cur = common_layers.conv_block(cur,
hparams.hidden_size, [((1, 1), (1, 1))],
name="mid_conv")
# To put a standard VAE use the line below.
# cur, vae_kl, _, _ = vae(cur, hparams, "kmeans_vae")
means = tf.get_variable("z_to_dense",
[hparams.v_size, hparams.hidden_size])
if hparams.use_gumbel_softmax:
_, hot, loss = dae(cur, hparams, "dae")
return cur, hot, loss
# Using k-means part. L2-normalizing to use fast cosine distance.
cur = mix(tf.nn.l2_normalize(cur, dim=3),
cur,
hparams.startup_steps // 3,
mode="exp",
simple=True)
cur_n = hparams.kmeans_lr_factor * cur
cur_n += (1.0 - hparams.kmeans_lr_factor) * tf.stop_gradient(cur)
hot, loss = kmeans(cur_n, means, hparams, name="kmeans")
# We need a linear layer to undo the l2-normalization.
cur = tf.layers.dense(cur, hparams.hidden_size, name="unnormalize")
return cur, hot, loss
def bottom(self, inputs):
with tf.variable_scope(self.name):
inputs = common_layers.standardize_images(inputs)
tf.summary.image("inputs", inputs, max_outputs=2)
return common_layers.conv_block(inputs,
self._body_input_depth,
[((1, 1), (3, 3))],
first_relu=False,
padding="SAME",
force2d=True,
name="small_image_conv")
def residual_conv(x, repeat, k, hparams, name, reuse=None):
"""A stack of convolution blocks with residual connections."""
with tf.variable_scope(name, reuse=reuse):
dilations_and_kernels = [((1, 1), k) for _ in xrange(3)]
for i in xrange(repeat):
with tf.variable_scope("repeat_%d" % i):
y = common_layers.conv_block(common_layers.layer_norm(
x, hparams.hidden_size, name="lnorm"),
hparams.hidden_size,
dilations_and_kernels,
padding="SAME",
name="residual_conv")
y = tf.nn.dropout(y, 1.0 - hparams.dropout)
x += y
return x
def compress(x, c, is_2d, hparams, name):
"""Compress."""
with tf.variable_scope(name):
# Run compression by strided convs.
cur = x
k1 = (3, 3) if is_2d else (3, 1)
k2 = (2, 2) if is_2d else (2, 1)
for i in xrange(hparams.num_compress_steps):
if c is not None:
cur = attend(cur, c, hparams, "compress_attend_%d" % i)
cur = residual_conv(cur, 1, k1, hparams, "compress_rc_%d" % i)
cur = common_layers.conv_block(cur,
hparams.hidden_size, [((1, 1), k2)],
strides=k2,
name="compress_%d" % i)
return cur
def residual_dilated_conv(x, repeat, padding, name, hparams):
"""A stack of convolution blocks with residual connections."""
with tf.variable_scope(name):
k = (hparams.kernel_height, hparams.kernel_width)
dilations_and_kernels = [((2**i, 1), k)
for i in xrange(hparams.num_hidden_layers)]
for i in xrange(repeat):
with tf.variable_scope("repeat_%d" % i):
y = common_layers.conv_block(common_layers.layer_norm(
x, hparams.hidden_size, name="lnorm"),
hparams.hidden_size,
dilations_and_kernels,
padding=padding,
name="residual_conv")
y = tf.nn.dropout(y, 1.0 - hparams.dropout)
x += y
return x
def decompress_step(source, c, hparams, first_relu, is_2d, name):
"""Decompression function."""
with tf.variable_scope(name):
shape = tf.shape(source)
if c is not None:
source = attend(source, c, hparams, "decompress_attend")
multiplier = 4 if is_2d else 2
kernel = (1, 1) if is_2d else (1, 1)
thicker = common_layers.conv_block(source,
hparams.hidden_size * multiplier,
[((1, 1), kernel)],
first_relu=first_relu,
name="decompress_conv")
if is_2d:
return tf.depth_to_space(thicker, 2)
return tf.reshape(thicker,
[shape[0], shape[1] * 2, 1, hparams.hidden_size])
def xnet_resblock(x, filters, res_relu, name):
with tf.variable_scope(name):
y = common_layers.separable_conv_block(
x,
filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
first_relu=True,
padding="SAME",
force2d=True,
name="sep_conv_block")
y = common_layers.pool(y, (3, 3),
"MAX",
"SAME",
strides=(2, 2))
return y + common_layers.conv_block(x,
filters, [((1, 1),
(1, 1))],
padding="SAME",
strides=(2, 2),
first_relu=res_relu,
force2d=True,
name="res_conv0")
def xnet_resblock(x, filters, res_relu, name):
with tf.variable_scope(name):
# We only stride along the length dimension to preserve the spectral
# bins (which are tiny in dimensionality relative to length)
y = common_layers.separable_conv_block(
x,
filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
first_relu=True,
padding="SAME",
force2d=True,
name="sep_conv_block")
y = common_layers.pool(y, (3, 3),
"MAX",
"SAME",
strides=(2, 1))
return y + common_layers.conv_block(x,
filters, [((1, 1),
(1, 1))],
padding="SAME",
strides=(2, 1),
first_relu=res_relu,
force2d=True,
name="res_conv0")
def xnet_resblock(x, filters, res_relu, name):
with tf.variable_scope(name):
# Typically audio samples are >100k samples in length and have a width
# of 2 or 4. Mono audio has a single channel while stereo has 2.
y = common_layers.separable_conv_block(
x,
filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
first_relu=True,
padding="SAME",
force2d=True,
name="sep_conv_block")
y = common_layers.pool(y, (3, 3),
"MAX",
"SAME",
strides=(2, 2))
return y + common_layers.conv_block(x,
filters, [((1, 1),
(1, 1))],
padding="SAME",
strides=(2, 2),
first_relu=res_relu,
force2d=True,
name="res_conv0")