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
"""
# 1. create a model and its outputs
net_output_ops = residual_unet_3d(
inputs=features['x'],
num_classes=NUM_CLASSES,
num_res_units=2,
filters=(16, 32, 64, 128),
strides=((1, 1, 1), (1, 2, 2), (1, 2, 2), (1, 2, 2)),
mode=mode,
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4))
# 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
ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=net_output_ops['logits'],
labels=labels['y'])
loss = tf.reduce_mean(ce)
# 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.MomentumOptimizer(
learning_rate=params["learning_rate"],
momentum=0.9)
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
my_image_summaries = {}
my_image_summaries['feat_t1'] = features['x'][0, 0, :, :, 0]
my_image_summaries['feat_t1_ir'] = features['x'][0, 0, :, :, 1]
my_image_summaries['feat_t2_flair'] = features['x'][0, 0, :, :, 2]
my_image_summaries['labels'] = tf.cast(labels['y'], tf.float32)[0, 0, :, :]
my_image_summaries['predictions'] = tf.cast(net_output_ops['y_'], tf.float32)[0, 0, :, :]
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()]
# 4.2 (optional) create custom metric summaries for tensorboard
dice_tensor = tf.py_func(dice, [net_output_ops['y_'],
labels['y'],
tf.constant(NUM_CLASSES)], tf.float32)
[tf.summary.scalar('dsc_l{}'.format(i), dice_tensor[i])
for i in range(NUM_CLASSES)]
# 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
"""
print("Setting up U-Net")
# 1. create a model and its outputs
net_output_ops = residual_unet_3d(
inputs=features['x'],
num_classes=NUM_CLASSES,
num_res_units=2,
filters=(16, 32, 64, 128),
strides=((1, 1, 1), (1, 2, 2), (1, 2, 2), (1, 2, 2)),
mode=mode,
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4))
# 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
# print(labels['y'])
ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=net_output_ops['logits'], labels=labels['y'])
loss = tf.reduce_mean(ce)
# 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.MomentumOptimizer(
learning_rate=params["learning_rate"], momentum=0.9)
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
my_image_summaries = {
'feat_t2': features['x'][0, 0, :, :, 0],
'labels': tf.cast(labels['y'], tf.float32)[0, 0, :, :],
'predictions': tf.cast(net_output_ops['y_'], tf.float32)[0, 0, :, :]
}
expected_output_size = [1, 64, 64, 1] # [B, W, H, C]
[
tf.summary.image(name, tf.reshape(image, expected_output_size))
for name, image in my_image_summaries.items()
]
# 4.2 (optional) create custom metric summaries for tensorboard
dice_tensor = tf.py_func(
dice, [net_output_ops['y_'], labels['y'],
tf.constant(NUM_CLASSES)], tf.float32)
[
tf.summary.scalar('dsc_l{}'.format(i), dice_tensor[i])
for i in range(NUM_CLASSES)
]
# 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(self, features, labels, mode, params):
"""Build architecture of network as an instance of tf.estimator.EstimatorSpec according
to HPs from top-level optimiser.
Args:
features (TYPE): Description
labels (TYPE): Description
mode (TYPE): Description
params (TYPE): Description
Returns:
TYPE: tf.estimator.EstimatorSpec
"""
# 1. create a model and its outputs
from dltk.core.metrics import dice
from dltk.core.losses import sparse_balanced_crossentropy
from dltk.networks.segmentation.unet import residual_unet_3d
from dltk.networks.segmentation.unet import asymmetric_residual_unet_3d
from dltk.networks.segmentation.fcn import residual_fcn_3d
from dltk.core.activations import leaky_relu
filters = params["filters"]
strides = params["strides"]
num_residual_units = params["num_residual_units"]
loss_type = params["loss"]
net = params["net"]
def lrelu(x):
return leaky_relu(x, 0.1)
if net == 'fcn':
net_output_ops = residual_fcn_3d(features['x'],
self.NUM_CLASSES,
num_res_units=num_residual_units,
filters=filters,
strides=strides,
activation=lrelu,
mode=mode)
elif net == 'unet':
net_output_ops = residual_unet_3d(features['x'],
self.NUM_CLASSES,
num_res_units=num_residual_units,
filters=filters,
strides=strides,
activation=lrelu,
mode=mode)
elif net == 'asym_unet':
net_output_ops = asymmetric_residual_unet_3d(
features['x'],
self.NUM_CLASSES,
num_res_units=num_residual_units,
filters=filters,
strides=strides,
activation=lrelu,
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
if loss_type == 'ce':
ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=net_output_ops['logits'], labels=labels['y'])
loss = tf.reduce_mean(ce)
elif loss_type == 'balce':
loss = sparse_balanced_crossentropy(net_output_ops['logits'],
labels['y'])
# 3. define a training op and ops for updating
# moving averages (i.e. for batch normalisation)
global_step = tf.train.get_or_create_global_step()
if params["opt"] == 'adam':
optimiser = tf.train.AdamOptimizer(
learning_rate=params["learning_rate"], epsilon=1e-5)
elif params["opt"] == 'momentum':
optimiser = tf.train.MomentumOptimizer(
learning_rate=params["learning_rate"], momentum=0.9)
elif params["opt"] == 'rmsprop':
optimiser = tf.train.RMSPropOptimizer(
learning_rate=params["learning_rate"], momentum=0.9)
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
my_image_summaries = {}
my_image_summaries['feat_t1'] = tf.expand_dims(
features['x'][:, 0, :, :, 0], 3)
my_image_summaries['labels'] = tf.expand_dims(
tf.cast(labels['y'], tf.float32)[:, 0, :, :], 3)
my_image_summaries['predictions'] = tf.expand_dims(
tf.cast(net_output_ops['y_'], tf.float32)[:, 0, :, :], 3)
[
tf.summary.image(name, image)
for name, image in my_image_summaries.items()
]
# 4.2 (optional) create custom metric summaries for tensorboard
dice_tensor = tf.py_func(
dice,
[net_output_ops['y_'], labels['y'],
tf.constant(self.NUM_CLASSES)], tf.float32)
[
tf.summary.scalar('dsc_l{}'.format(i), dice_tensor[i])
for i in range(self.NUM_CLASSES)
]
# average dice over all classes
dice_metric = tf.metrics.mean(dice_tensor)
metrics = {'dice': dice_metric}
# 5. Return EstimatorSpec object
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=net_output_ops,
loss=loss,
train_op=train_op,
training_hooks=[self.training_metrics],
eval_metric_ops=metrics)