def shake_shake_block(x, conv_filters, stride, hparams):
"""A shake-shake block."""
with tf.variable_scope("branch_1"):
branch1 = shake_shake_block_branch(x, conv_filters, stride)
with tf.variable_scope("branch_2"):
branch2 = shake_shake_block_branch(x, conv_filters, stride)
if x.shape[-1] == conv_filters:
skip = tf.identity(x)
else:
skip = downsampling_residual_branch(x, conv_filters)
# TODO(rshin): Use different alpha for each image in batch.
if hparams.mode == tf.contrib.learn.ModeKeys.TRAIN:
if hparams.shakeshake_type == "batch":
shaken = common_layers.shakeshake2(branch1, branch2)
elif hparams.shakeshake_type == "image":
shaken = common_layers.shakeshake2_indiv(branch1, branch2)
elif hparams.shakeshake_type == "equal":
shaken = common_layers.shakeshake2_py(branch1, branch2, equal=True)
else:
raise ValueError("Invalid shakeshake_type: {!r}".format(shaken))
else:
shaken = common_layers.shakeshake2_py(branch1, branch2, equal=True)
shaken.set_shape(branch1.get_shape())
return skip + shaken
def shake_shake_block(x, conv_filters, stride, hparams):
"""A shake-shake block."""
with tf.variable_scope("branch_1"):
branch1 = shake_shake_block_branch(x, conv_filters, stride)
with tf.variable_scope("branch_2"):
branch2 = shake_shake_block_branch(x, conv_filters, stride)
if x.shape[-1] == conv_filters:
skip = tf.identity(x)
else:
skip = downsampling_residual_branch(x, conv_filters)
# TODO(rshin): Use different alpha for each image in batch.
if hparams.mode == tf.estimator.ModeKeys.TRAIN:
if hparams.shakeshake_type == "batch":
shaken = common_layers.shakeshake2(branch1, branch2)
elif hparams.shakeshake_type == "image":
shaken = common_layers.shakeshake2_indiv(branch1, branch2)
elif hparams.shakeshake_type == "equal":
shaken = common_layers.shakeshake2_py(branch1, branch2, equal=True)
else:
raise ValueError("Invalid shakeshake_type: {!r}".format(shaken))
else:
shaken = common_layers.shakeshake2_py(branch1, branch2, equal=True)
shaken.set_shape(branch1.get_shape())
return skip + shaken