def upsample_and_concat(inputs, inputs2, strides=(2, 2, 2)):
"""Upsampling and concatenation layer according to [1].
[1] O. Ronneberger et al. U-Net: Convolutional Networks for Biomedical Image
Segmentation. MICCAI 2015.
Args:
inputs (TYPE): Input features to be upsampled.
inputs2 (TYPE): Higher resolution features from the encoder to
concatenate.
strides (tuple, optional): Upsampling factor for a strided transpose
convolution.
Returns:
tf.Tensor: Upsampled feature tensor
"""
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
# assert len(inputs.get_shape().as_list()) == len(inputs2.get_shape().as_list()), \
# 'Ranks of input and input2 differ'
# Upsample inputs
inputs = linear_upsample_3d(inputs, strides)
return tf.concat(axis=-1, values=[inputs2, inputs])
def upsample(inputs, strides=(2, 2, 2)):
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
# Upsample inputs
inputs = linear_upsample_3d(inputs, strides)
return inputs
def _upsample(x, factor):
if isinstance(factor, int):
factor = [factor] * (len(x.get_shape().as_list()) - 2)
# TODO: build repeat upsampling
x = linear_upsample_3d(x, strides=factor)
return x
def _upsample(x, factor):
if isinstance(factor, int):
factor = [factor] * (len(x.get_shape().as_list()) - 2)
# TODO: build repeat upsampling
x = linear_upsample_3d(x, strides=factor)
return x
def upsample_and_concat(inputs, inputs2, strides=(2, 2, 2)):
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
assert len(inputs.get_shape().as_list()) == len(inputs2.get_shape().as_list()), \
'Ranks of input and input2 differ'
# Upsample inputs
inputs1 = linear_upsample_3d(inputs, strides)
return tf.concat(axis=-1, values=[inputs2, inputs1])
def upsample_and_concat_with_conv(inputs,
inputs2,
strides=(2, 2, 2),
mode=tf.estimator.ModeKeys.EVAL,
conv_params=global_conv_params):
"""Upsampling and concatenation layer according to [1].
[1] O. Ronneberger et al. U-Net: Convolutional Networks for Biomedical Image
Segmentation. MICCAI 2015.
Args:
inputs (TYPE): Input features to be upsampled.
inputs2 (TYPE): Higher resolution features from the encoder to
concatenate.
strides (tuple, optional): Upsampling factor for a strided transpose
convolution.
Returns:
tf.Tensor: Upsampled feature tensor
"""
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
assert len(inputs.get_shape().as_list()) == len(inputs2.get_shape().as_list()), \
'Ranks of input and input2 differ'
# Upsample inputs
inputs1 = linear_upsample_3d(inputs, strides)
with tf.variable_scope('reduce_channel_unit'):
result = tf.concat(axis=-1, values=[inputs2, inputs1])
result = tf.layers.batch_normalization(
result, training=mode == tf.estimator.ModeKeys.TRAIN)
result = leaky_relu(result)
result = tf.layers.conv3d(inputs=result,
filters=inputs2.shape[-1],
kernel_size=(3, 3, 3),
strides=(1, 1, 1),
**conv_params)
return result
def upsample_and_conv(inputs,
strides=(2, 2, 2),
mode=tf.estimator.ModeKeys.EVAL,
conv_params=global_conv_params,
name=None):
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
# Upsample inputs
with tf.variable_scope('reduce_channel_unit_{}'.format(name)):
inputs1 = linear_upsample_3d(inputs, strides)
result = tf.layers.batch_normalization(
inputs1, training=mode == tf.estimator.ModeKeys.TRAIN)
result = leaky_relu(result)
result = tf.layers.conv3d(inputs=result,
filters=inputs.get_shape().as_list()[-1] / 2,
kernel_size=(3, 3, 3),
strides=(1, 1, 1),
**conv_params)
return result
def upscore_layer_3d(inputs,
inputs2,
out_filters,
in_filters=None,
strides=(2, 2, 2),
mode=tf.estimator.ModeKeys.EVAL, use_bias=False,
kernel_initializer=tf.initializers.variance_scaling(distribution='uniform'),
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None):
"""Upscore layer according to [1].
[1] J. Long et al. Fully convolutional networks for semantic segmentation.
CVPR 2015.
Args:
inputs (tf.Tensor): Input features to be upscored.
inputs2 (tf.Tensor): Higher resolution features from the encoder to add.
out_filters (int): Number of output filters (typically, number of
segmentation classes)
in_filters (None, optional): None or number of input filters.
strides (tuple, optional): Upsampling factor for a strided transpose
convolution.
mode (TYPE, optional): One of the tf.estimator.ModeKeys strings: TRAIN,
EVAL or PREDICT
use_bias (bool, optional): Boolean, whether the layer uses a bias.
kernel_initializer (TYPE, optional): An initializer for the convolution
kernel.
bias_initializer (TYPE, optional): An initializer for the bias vector.
If None, no bias will be applied.
kernel_regularizer (None, optional): Optional regularizer for the
convolution kernel.
bias_regularizer (None, optional): Optional regularizer for the bias
vector.
Returns:
tf.Tensor: Upscore tensor
"""
conv_params = {'use_bias': use_bias,
'kernel_initializer': kernel_initializer,
'bias_initializer': bias_initializer,
'kernel_regularizer': kernel_regularizer,
'bias_regularizer': bias_regularizer}
# Compute an upsampling shape dynamically from the input tensor. Input
# filters are required to be static.
if in_filters is None:
in_filters = inputs.get_shape().as_list()[-1]
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
assert len(inputs.get_shape().as_list()) == len(inputs2.get_shape().as_list()), \
'Ranks of input and input2 differ'
# Account for differences in the number of input and output filters
if in_filters != out_filters:
x = tf.layers.conv3d(inputs=inputs,
filters=out_filters,
kernel_size=(1, 1, 1),
strides=(1, 1, 1),
padding='same',
name='filter_conversion',
**conv_params)
else:
x = inputs
# Upsample inputs
x = linear_upsample_3d(inputs=x, strides=strides)
# Skip connection
x2 = tf.layers.conv3d(inputs=inputs2,
filters=out_filters,
kernel_size=(1, 1, 1),
strides=(1, 1, 1),
padding='same',
**conv_params)
x2 = tf.layers.batch_normalization(
x2, training=mode == tf.estimator.ModeKeys.TRAIN)
# Return the element-wise sum
return tf.add(x, x2)
def model_fn(features, labels, mode, params):
"""Model function to construct a tf.estimator.EstimatorSpec. It creates a
network given input features (e.g. from a dltk.io.abstract_reader) and
training targets (labels). Further, loss, optimiser, evaluation ops and
custom tensorboard summary ops can be added. For additional information,
please refer to https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator#model_fn.
Args:
features (tf.Tensor): Tensor of input features to train from. Required
rank and dimensions are determined by the subsequent ops
(i.e. the network).
labels (tf.Tensor): Tensor of training targets or labels. Required rank
and dimensions are determined by the network output.
mode (str): One of the tf.estimator.ModeKeys: TRAIN, EVAL or PREDICT
params (dict, optional): A dictionary to parameterise the model_fn
(e.g. learning_rate)
Returns:
tf.estimator.EstimatorSpec: A custom EstimatorSpec for this experiment
"""
assert len(params['upsampling_factor']) == 3
# During training downsample the original input to create a lower resolution
# image
if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
lo_res = tf.layers.average_pooling3d(
inputs=features['x'],
pool_size=[
2 * s if s > 1 else 1 for s in params['upsampling_factor']
],
strides=params['upsampling_factor'],
padding='same',
name='training_downsample')
# Compute some linear upsampling for comparison:
lin_up = linear_upsample_3d(inputs=lo_res,
strides=params['upsampling_factor'])
else:
lo_res = features['x']
# 1. create a model and its outputs
net_output_ops = simple_super_resolution_3d(
inputs=lo_res,
num_convolutions=2,
upsampling_factor=params['upsampling_factor'],
filters=(16, 32, 64, 128),
mode=mode)
# 1.1 Generate predictions only (for `ModeKeys.PREDICT`)
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=net_output_ops,
export_outputs={
'out': tf.estimator.export.PredictOutput(net_output_ops)
})
# 2. set up a loss function
loss = tf.losses.mean_squared_error(labels=features['x'],
predictions=net_output_ops['x_'])
# 3. define a training op and ops for updating moving averages
# (i.e. for batch normalisation)
global_step = tf.train.get_global_step()
optimiser = tf.train.AdamOptimizer(learning_rate=params['learning_rate'],
epsilon=1e-5)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimiser.minimize(loss, global_step=global_step)
# 4.1 (optional) create custom image summaries for tensorboard
tf.summary.image('feat_lo_res',
tf.reshape(lo_res[0, 4, :, :, 0:1], [1, 32, 32, 1]))
my_image_summaries = {}
my_image_summaries['feat_hi_res'] = features['x'][0, 16, :, :, 0:1]
my_image_summaries['linear_up_hi_res'] = tf.cast(lin_up,
tf.float32)[0, 16, :, :,
0:1]
my_image_summaries['pred_hi_res'] = tf.cast(net_output_ops['x_'],
tf.float32)[0, 16, :, :, 0:1]
expected_output_size = [1, 128, 128, 1] # [B, W, H, C]
[
tf.summary.image(name, tf.reshape(image, expected_output_size))
for name, image in my_image_summaries.items()
]
# 5. Return EstimatorSpec object
return tf.estimator.EstimatorSpec(mode=mode,
predictions=net_output_ops,
loss=loss,
train_op=train_op,
eval_metric_ops=None)
def model_fn(features, labels, mode, params):
"""Model function to construct a tf.estimator.EstimatorSpec. It creates a
network given input features (e.g. from a dltk.io.abstract_reader) and
training targets (labels). Further, loss, optimiser, evaluation ops and
custom tensorboard summary ops can be added. For additional information,
please refer to https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator#model_fn.
Args:
features (tf.Tensor): Tensor of input features to train from. Required
rank and dimensions are determined by the subsequent ops
(i.e. the network).
labels (tf.Tensor): Tensor of training targets or labels. Required rank
and dimensions are determined by the network output.
mode (str): One of the tf.estimator.ModeKeys: TRAIN, EVAL or PREDICT
params (dict, optional): A dictionary to parameterise the model_fn
(e.g. learning_rate)
Returns:
tf.estimator.EstimatorSpec: A custom EstimatorSpec for this experiment
"""
assert len(params['upsampling_factor']) == 3
# During training downsample the original input to create a lower resolution
# image
if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
lo_res = tf.layers.average_pooling3d(
inputs=features['x'],
pool_size=[2 * s if s > 1 else 1 for s in params['upsampling_factor']],
strides=params['upsampling_factor'],
padding='same',
name='training_downsample')
# Compute some linear upsampling for comparison:
lin_up = linear_upsample_3d(inputs=lo_res,
strides=params['upsampling_factor'])
else:
lo_res = features['x']
# 1. create a model and its outputs
net_output_ops = simple_super_resolution_3d(
inputs=lo_res,
num_convolutions=2,
upsampling_factor=params['upsampling_factor'],
filters=(16, 32, 64, 128),
mode=mode)
# 1.1 Generate predictions only (for `ModeKeys.PREDICT`)
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=net_output_ops,
export_outputs={'out': tf.estimator.export.PredictOutput(net_output_ops)})
# 2. set up a loss function
loss = tf.losses.mean_squared_error(
labels=features['x'],
predictions=net_output_ops['x_'])
# 3. define a training op and ops for updating moving averages
# (i.e. for batch normalisation)
global_step = tf.train.get_global_step()
optimiser = tf.train.AdamOptimizer(
learning_rate=params['learning_rate'],
epsilon=1e-5)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimiser.minimize(loss, global_step=global_step)
# 4.1 (optional) create custom image summaries for tensorboard
tf.summary.image('feat_lo_res', tf.reshape(lo_res[0, 4, :, :, 0:1], [1, 32, 32, 1]))
my_image_summaries = {}
my_image_summaries['feat_hi_res'] = features['x'][0, 16, :, :, 0:1]
my_image_summaries['linear_up_hi_res'] = tf.cast(lin_up, tf.float32)[0, 16, :, :, 0:1]
my_image_summaries['pred_hi_res'] = tf.cast(net_output_ops['x_'], tf.float32)[0, 16, :, :, 0:1]
expected_output_size = [1, 128, 128, 1] # [B, W, H, C]
[tf.summary.image(name, tf.reshape(image, expected_output_size))
for name, image in my_image_summaries.items()]
# 5. Return EstimatorSpec object
return tf.estimator.EstimatorSpec(mode=mode,
predictions=net_output_ops,
loss=loss,
train_op=train_op,
eval_metric_ops=None)
def dcgan_generator_3d(inputs,
filters=(256, 128, 64, 32, 1),
kernel_size=((4, 4, 4), (3, 3, 3), (3, 3, 3), (3, 3, 3),
(4, 4, 4)),
strides=((4, 4, 4), (1, 2, 2), (1, 2, 2), (1, 2, 2),
(1, 2, 2)),
mode=tf.estimator.ModeKeys.TRAIN,
use_bias=False):
"""
Deep convolutional generative adversial network (DCGAN) generator
network. with num_convolutions on len(filters) resolution scales. The
upsampling of features is done via strided transpose convolutions. On
each resolution scale s are num_convolutions with filter size = filters[
s]. strides[s] determine the upsampling factor at each resolution scale.
Args:
inputs (tf.Tensor): Input noise tensor to the network.
out_filters (int): Number of output filters.
num_convolutions (int, optional): Number of convolutions per resolution
scale.
filters (tuple, optional): Number of filters for all convolutions at
each resolution scale.
strides (tuple, optional): Stride of the first convolution on a
resolution scale.
mode (TYPE, optional): One of the tf.estimator.ModeKeys strings: TRAIN,
EVAL or PREDICT
use_bias (bool, optional): Boolean, whether the layer uses a bias.
Returns:
dict: dictionary of output tensors
"""
outputs = {}
assert len(strides) == len(filters)
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
conv_op = tf.layers.conv3d
conv_params = {
'padding': 'same',
'use_bias': use_bias,
'kernel_initializer': tf.uniform_unit_scaling_initializer(),
'bias_initializer': tf.zeros_initializer(),
'kernel_regularizer': None,
'bias_regularizer': None
}
x = inputs
tf.logging.info('Input tensor shape {}'.format(x.get_shape()))
for res_scale in range(0, len(filters)):
with tf.variable_scope('gen_unit_{}'.format(res_scale)):
tf.logging.info('Generator at res_scale before up {} tensor '
'shape: {}'.format(res_scale, x.get_shape()))
x = linear_upsample_3d(x, strides[res_scale], trainable=True)
x = conv_op(inputs=x,
filters=filters[res_scale],
kernel_size=kernel_size[res_scale],
**conv_params)
tf.logging.info('Generator at res_scale after up {} tensor '
'shape: {}'.format(res_scale, x.get_shape()))
x = tf.layers.batch_normalization(
x, training=mode == tf.estimator.ModeKeys.TRAIN)
x = leaky_relu(x, 0.2)
tf.logging.info('Generator at res_scale {} tensor shape: '
'{}'.format(res_scale, x.get_shape()))
outputs['gen'] = x
return outputs
def upscore_layer_3d(inputs,
inputs2,
out_filters,
in_filters=None,
strides=(2, 2, 2),
mode=tf.estimator.ModeKeys.EVAL,
use_bias=False,
kernel_initializer=tf.initializers.variance_scaling(distribution='uniform'),
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None):
"""Upscore layer according to [1].
[1] J. Long et al. Fully convolutional networks for semantic segmentation.
CVPR 2015.
Args:
inputs (tf.Tensor): Input features to be upscored.
inputs2 (tf.Tensor): Higher resolution features from the encoder to add.
out_filters (int): Number of output filters (typically, number of
segmentation classes)
in_filters (None, optional): None or number of input filters.
strides (tuple, optional): Upsampling factor for a strided transpose
convolution.
mode (TYPE, optional): One of the tf.estimator.ModeKeys strings: TRAIN,
EVAL or PREDICT
use_bias (bool, optional): Boolean, whether the layer uses a bias.
kernel_initializer (TYPE, optional): An initializer for the convolution
kernel.
bias_initializer (TYPE, optional): An initializer for the bias vector.
If None, no bias will be applied.
kernel_regularizer (None, optional): Optional regularizer for the
convolution kernel.
bias_regularizer (None, optional): Optional regularizer for the bias
vector.
Returns:
tf.Tensor: Upscore tensor
"""
conv_params = {'use_bias': use_bias,
'kernel_initializer': kernel_initializer,
'bias_initializer': bias_initializer,
'kernel_regularizer': kernel_regularizer,
'bias_regularizer': bias_regularizer}
# Compute an upsampling shape dynamically from the input tensor. Input
# filters are required to be static.
if in_filters is None:
in_filters = inputs.get_shape().as_list()[-1]
assert len(inputs.get_shape().as_list()) == 5, \
'inputs are required to have a rank of 5.'
assert len(inputs.get_shape().as_list()) == len(inputs2.get_shape().as_list()), \
'Ranks of input and input2 differ'
# Account for differences in the number of input and output filters
if in_filters != out_filters:
x = tf.layers.conv3d(inputs=inputs,
filters=out_filters,
kernel_size=(1, 1, 1),
strides=(1, 1, 1),
padding='same',
name='filter_conversion',
**conv_params)
else:
x = inputs
# Upsample inputs
x = linear_upsample_3d(inputs=x, strides=strides)
# Skip connection
x2 = tf.layers.conv3d(inputs=inputs2,
filters=out_filters,
kernel_size=(1, 1, 1),
strides=(1, 1, 1),
padding='same',
**conv_params)
x2 = tf.layers.batch_normalization(
x2, training=mode == tf.estimator.ModeKeys.TRAIN)
# Return the element-wise sum
return tf.add(x, x2)