def layer_op(self, images, is_training=True, layer_id=-1, **unused_kwargs):
# image_size should be divisible by 4
assert layer_util.check_spatial_dims(images, lambda x: x % 4 == 0)
assert layer_util.check_spatial_dims(images, lambda x: x >= 21)
block_layer = UNetBlock('DOWNSAMPLE',
(self.n_features[0], self.n_features[1]),
(3, 3),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d0')
pool_1, conv_1 = block_layer(images, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE',
(self.n_features[1], self.n_features[2]),
(3, 3),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d1')
up_1, _ = block_layer(pool_1, is_training)
print(block_layer)
block_layer = UNetBlock(
'NONE', (self.n_features[1], self.n_features[1], self.num_classes),
(3, 3),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='u0')
crop_layer = CropLayer(border=4, name='crop-8')
concat_1 = ElementwiseLayer('CONCAT')(crop_layer(conv_1), up_1)
print(block_layer)
# for the last layer, upsampling path is not used
_, output_tensor = block_layer(concat_1, is_training)
output_conv_op = ConvolutionalLayer(
n_output_chns=self.num_classes,
kernel_size=1,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=None,
name='{}'.format(self.num_classes),
padding='VALID',
with_bn=False,
with_bias=True)
final_output_tensor = output_conv_op(output_tensor, is_training)
print(output_conv_op)
return final_output_tensor
def test_3d_shape(self):
input_shape = (2, 16, 16, 16, 8)
test_border = 3
x = tf.ones(input_shape)
crop_layer = CropLayer(border=test_border)
out_crop = crop_layer(x)
print(crop_layer)
input_shape = (2, 7, 7, 7, 8)
test_border = 3
x = tf.ones(input_shape)
crop_layer = CropLayer(border=test_border)
out_crop_1 = crop_layer(x)
print(crop_layer)
with self.cached_session() as sess:
out = sess.run(out_crop)
out_1 = sess.run(out_crop_1)
self.assertAllClose((2, 10, 10, 10, 8), out.shape)
self.assertAllClose((2, 1, 1, 1, 8), out_1.shape)
def layer_op(self, images, is_training, layer_id=-1):
# image_size should be divisible by 8
assert layer_util.check_spatial_dims(images, lambda x: x % 8 == 0)
assert layer_util.check_spatial_dims(images, lambda x: x >= 89)
block_layer = UNetBlock('DOWNSAMPLE',
(self.n_features[0], self.n_features[1]),
(3, 3),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='L1')
pool_1, conv_1 = block_layer(images, is_training)
print(block_layer)
block_layer = UNetBlock('DOWNSAMPLE',
(self.n_features[1], self.n_features[2]),
(3, 3),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='L2')
pool_2, conv_2 = block_layer(pool_1, is_training)
print(block_layer)
block_layer = UNetBlock('DOWNSAMPLE',
(self.n_features[2], self.n_features[3]),
(3, 3),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='L3')
pool_3, conv_3 = block_layer(pool_2, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE',
(self.n_features[3], self.n_features[4]),
(3, 3),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='L4')
up_3, _ = block_layer(pool_3, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE',
(self.n_features[3], self.n_features[3]),
(3, 3),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='R3')
concat_3 = ElementwiseLayer('CONCAT')(conv_3, up_3)
up_2, _ = block_layer(concat_3, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE',
(self.n_features[2], self.n_features[2]),
(3, 3),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='R2')
concat_2 = ElementwiseLayer('CONCAT')(conv_2, up_2)
up_1, _ = block_layer(concat_2, is_training)
print(block_layer)
block_layer = UNetBlock(
'NONE', (self.n_features[1], self.n_features[1], self.num_classes),
(3, 3, 1),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='R1_FC')
concat_1 = ElementwiseLayer('CONCAT')(conv_1, up_1)
# for the last layer, upsampling path is not used
_, output_tensor = block_layer(concat_1, is_training)
crop_layer = CropLayer(border=44, name='crop-88')
output_tensor = crop_layer(output_tensor)
print(block_layer)
return output_tensor
def connect_data_and_network(self,
outputs_collector=None,
gradients_collector=None):
def switch_sampler(for_training):
with tf.name_scope('train' if for_training else 'validation'):
sampler = self.get_sampler()[0][0 if for_training else -1]
return sampler.pop_batch_op()
if self.is_training:
if self.action_param.validation_every_n > 0:
data_dict = tf.cond(tf.logical_not(self.is_validation),
lambda: switch_sampler(True),
lambda: switch_sampler(False))
else:
data_dict = switch_sampler(for_training=True)
image = tf.cast(data_dict['image'], tf.float32)
net_args = {
'is_training': self.is_training,
'keep_prob': self.net_param.keep_prob
}
net_out = self.net(image, **net_args)
with tf.name_scope('Optimiser'):
optimiser_class = OptimiserFactory.create(
name=self.action_param.optimiser)
self.optimiser = optimiser_class.get_instance(
learning_rate=self.action_param.lr)
loss_func = LossFunction(loss_type=self.action_param.loss_type)
crop_layer = CropLayer(border=self.regression_param.loss_border,
name='crop-88')
prediction = crop_layer(net_out)
ground_truth = crop_layer(data_dict.get('output', None))
weight_map = None if data_dict.get('weight', None) is None \
else crop_layer(data_dict.get('weight', None))
data_loss = loss_func(prediction=prediction,
ground_truth=ground_truth,
weight_map=weight_map)
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
if self.net_param.decay > 0.0 and reg_losses:
reg_loss = tf.reduce_mean(
[tf.reduce_mean(reg_loss) for reg_loss in reg_losses])
loss = data_loss + reg_loss
else:
loss = data_loss
grads = self.optimiser.compute_gradients(loss)
# Gradient Clipping associated with VDSR3D
# Gradients are clipped by value, instead of clipping by global norm.
# The authors of VDSR do not specify a threshold for the clipping process.
# grads2, vars2 = zip(*grads)
# grads2, _ = tf.clip_by_global_norm(grads2, 5.0)
# grads = zip(grads2, vars2)
grads = [(tf.clip_by_value(grad, -0.00001 / self.action_param.lr,
+0.00001 / self.action_param.lr), val)
for grad, val in grads if grad is not None]
# collecting gradients variables
gradients_collector.add_to_collection([grads])
# collecting output variables
outputs_collector.add_to_collection(var=data_loss,
name='Loss',
average_over_devices=False,
collection=CONSOLE)
outputs_collector.add_to_collection(var=data_loss,
name='Loss',
average_over_devices=True,
summary_type='scalar',
collection=TF_SUMMARIES)
elif self.is_inference:
data_dict = switch_sampler(for_training=False)
image = tf.cast(data_dict['image'], tf.float32)
net_args = {
'is_training': self.is_training,
'keep_prob': self.net_param.keep_prob
}
net_out = self.net(image, **net_args)
crop_layer = CropLayer(border=0, name='crop-88')
post_process_layer = PostProcessingLayer('IDENTITY')
net_out = post_process_layer(crop_layer(net_out))
outputs_collector.add_to_collection(var=net_out,
name='window',
average_over_devices=False,
collection=NETWORK_OUTPUT)
outputs_collector.add_to_collection(
var=data_dict['image_location'],
name='location',
average_over_devices=False,
collection=NETWORK_OUTPUT)
self.initialise_aggregator()
def connect_data_and_network(self,
outputs_collector=None,
gradients_collector=None):
def switch_sampler(for_training):
with tf.name_scope('train' if for_training else 'validation'):
sampler = self.get_sampler()[0][0 if for_training else -1]
return sampler.pop_batch_op()
if self.is_training:
if self.action_param.validation_every_n > 0:
data_dict = tf.cond(tf.logical_not(self.is_validation),
lambda: switch_sampler(True),
lambda: switch_sampler(False))
else:
data_dict = switch_sampler(for_training=True)
image = tf.cast(data_dict['image'], tf.float32)
net_out = self.net(image, is_training=self.is_training)
with tf.name_scope('Optimiser'):
optimiser_class = OptimiserFactory.create(
name=self.action_param.optimiser)
self.optimiser = optimiser_class.get_instance(
learning_rate=self.action_param.lr)
loss_func = LossFunction(
loss_type=self.action_param.loss_type)
crop_layer = CropLayer(
border=self.regression_param.loss_border, name='crop-88')
prediction = crop_layer(net_out)
ground_truth = crop_layer(data_dict.get('output', None))
weight_map = None if data_dict.get('weight', None) is None \
else crop_layer(data_dict.get('weight', None))
data_loss = loss_func(prediction=prediction,
ground_truth=ground_truth,
weight_map=weight_map)
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
if self.net_param.decay > 0.0 and reg_losses:
reg_loss = tf.reduce_mean(
[tf.reduce_mean(reg_loss) for reg_loss in reg_losses])
loss = data_loss + reg_loss
else:
loss = data_loss
grads = self.optimiser.compute_gradients(loss)
# collecting gradients variables
gradients_collector.add_to_collection([grads])
# collecting output variables
outputs_collector.add_to_collection(
var=data_loss, name='Loss',
average_over_devices=False, collection=CONSOLE)
outputs_collector.add_to_collection(
var=data_loss, name='Loss',
average_over_devices=True, summary_type='scalar',
collection=TF_SUMMARIES)
else:
data_dict = switch_sampler(for_training=False)
image = tf.cast(data_dict['image'], tf.float32)
net_out = self.net(image, is_training=self.is_training)
crop_layer = CropLayer(border=0, name='crop-88')
post_process_layer = PostProcessingLayer('IDENTITY')
net_out = post_process_layer(crop_layer(net_out))
outputs_collector.add_to_collection(
var=net_out, name='window',
average_over_devices=False, collection=NETWORK_OUTPUT)
outputs_collector.add_to_collection(
var=data_dict['image_location'], name='location',
average_over_devices=False, collection=NETWORK_OUTPUT)
init_aggregator = \
self.SUPPORTED_SAMPLING[self.net_param.window_sampling][2]
init_aggregator()
def connect_data_and_network(self,
outputs_collector=None,
gradients_collector=None):
def switch_sampler(for_training):
with tf.name_scope('train' if for_training else 'validation'):
sampler = self.get_sampler()[0][0 if for_training else -1]
return sampler.pop_batch_op()
if self.is_training:
if self.action_param.validation_every_n > 0:
data_dict = tf.cond(tf.logical_not(self.is_validation),
lambda: switch_sampler(for_training=True),
lambda: switch_sampler(for_training=False))
else:
data_dict = switch_sampler(for_training=True)
image = tf.cast(data_dict['image'], tf.float32)
net_args = {
'is_training': self.is_training,
'keep_prob': self.net_param.keep_prob
}
net_out = self.net(image, **net_args)
with tf.name_scope('Optimiser'):
optimiser_class = OptimiserFactory.create(
name=self.action_param.optimiser)
self.optimiser = optimiser_class.get_instance(
learning_rate=self.action_param.lr)
loss_func = LossFunction(loss_type=self.action_param.loss_type)
crop_layer = CropLayer(border=self.regression_param.loss_border)
weight_map = data_dict.get('weight', None)
weight_map = None if weight_map is None else crop_layer(weight_map)
data_loss = loss_func(prediction=crop_layer(net_out),
ground_truth=crop_layer(data_dict['output']),
weight_map=weight_map)
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
if self.net_param.decay > 0.0 and reg_losses:
reg_loss = tf.reduce_mean(
[tf.reduce_mean(reg_loss) for reg_loss in reg_losses])
loss = data_loss + reg_loss
else:
loss = data_loss
# Get all vars
to_optimise = tf.trainable_variables()
vars_to_freeze = \
self.action_param.vars_to_freeze or \
self.action_param.vars_to_restore
if vars_to_freeze:
import re
var_regex = re.compile(vars_to_freeze)
# Only optimise vars that are not frozen
to_optimise = \
[v for v in to_optimise if not var_regex.search(v.name)]
tf.logging.info(
"Optimizing %d out of %d trainable variables, "
"the other variables are fixed (--vars_to_freeze %s)",
len(to_optimise), len(tf.trainable_variables()),
vars_to_freeze)
grads = self.optimiser.compute_gradients(
loss, var_list=to_optimise, colocate_gradients_with_ops=True)
# collecting gradients variables
gradients_collector.add_to_collection([grads])
# collecting output variables
outputs_collector.add_to_collection(var=data_loss,
name='loss',
average_over_devices=False,
collection=CONSOLE)
outputs_collector.add_to_collection(var=data_loss,
name='loss',
average_over_devices=True,
summary_type='scalar',
collection=TF_SUMMARIES)
elif self.is_inference:
data_dict = switch_sampler(for_training=False)
image = tf.cast(data_dict['image'], tf.float32)
net_args = {
'is_training': self.is_training,
'keep_prob': self.net_param.keep_prob
}
net_out = self.net(image, **net_args)
net_out = PostProcessingLayer('IDENTITY')(net_out)
outputs_collector.add_to_collection(var=net_out,
name='window',
average_over_devices=False,
collection=NETWORK_OUTPUT)
outputs_collector.add_to_collection(
var=data_dict['image_location'],
name='location',
average_over_devices=False,
collection=NETWORK_OUTPUT)
self.initialise_aggregator()
def layer_op(self, images, is_training, layer_id=-1, **unused_kwargs):
"""
:param images: tensor, input to the network, size should be divisible by d_factor
:param is_training: boolean, True if network is in training mode
:param layer_id: not in use
:param unused_kwargs:
:return: tensor, network output
"""
# image_size is defined as the largest context, then:
# downsampled path size: image_size / d_factor
# downsampled path output: image_size / d_factor - 16
# to make sure same size of feature maps from both pathways:
# normal path size: (image_size / d_factor - 16) * d_factor + 16
# normal path output: (image_size / d_factor - 16) * d_factor
# where 16 is fixed by the receptive field of conv layers
# TODO: make sure label_size = image_size/d_factor - 16
# image_size has to be an odd number and divisible by 3 and
# smaller than the smallest image size of the input volumes
# label_size should be (image_size/d_factor - 16) * d_factor
assert self.d_factor % 2 == 1 # to make the downsampling centered
assert (layer_util.check_spatial_dims(
images, lambda x: x % self.d_factor == 0))
assert (layer_util.check_spatial_dims(images, lambda x: x % 2 == 1)
) # to make the crop centered
assert (layer_util.check_spatial_dims(images,
lambda x: x > self.d_factor * 16)
) # required by receptive field
# crop 25x25x25 from 57x57x57
crop_op = CropLayer(border=self.crop_diff, name='cropping_input')
normal_path = crop_op(images)
print(crop_op)
# downsample 19x19x19 from 57x57x57
downsample_op = DownSampleLayer(func='CONSTANT',
kernel_size=self.d_factor,
stride=self.d_factor,
padding='VALID',
name='downsample_input')
downsample_path = downsample_op(images)
print(downsample_op)
# convolutions for both pathways
for n_features in self.conv_features:
# normal pathway convolutions
conv_path_1 = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=3,
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='normal_conv')
normal_path = conv_path_1(normal_path, is_training)
print(conv_path_1)
# downsampled pathway convolutions
conv_path_2 = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=3,
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='downsample_conv')
downsample_path = conv_path_2(downsample_path, is_training)
print(conv_path_2)
# upsampling the downsampled pathway
downsample_path = UpSampleLayer('REPLICATE',
kernel_size=self.d_factor,
stride=self.d_factor)(downsample_path)
# concatenate both pathways
output_tensor = ElementwiseLayer('CONCAT')(normal_path,
downsample_path)
# 1x1x1 convolution layer
for n_features in self.fc_features:
conv_fc = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=1,
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_1x1x1_{}'.format(n_features))
output_tensor = conv_fc(output_tensor, is_training)
print(conv_fc)
return output_tensor
def layer_op(self,
images,
is_training,
keep_prob=0.5,
layer_id=-1,
**unused_kwargs):
# crop 27x27x27 from 59x59x59
crop_op = CropLayer(border=self.crop_diff, name='cropping_input')
normal_path = crop_op(images)
dilated_path = images
print(crop_op)
# dilated pathway
# dilation rate: 1,2,4,2,8,2,4,2,1
for n_features, rate in zip(self.conv2_features, self.dilation_rates):
dilated_block = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=3,
padding='VALID',
dilation=rate,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='dilated_conv_{}'.format(n_features))
dilated_path = dilated_block(dilated_path, is_training)
print(dilated_block)
# normal pathway
for n_features in self.conv1_features:
# normal pathway convolutions
conv_path_1 = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=3,
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='normal_conv_{}'.format(n_features))
normal_path = conv_path_1(normal_path, is_training)
print(conv_path_1)
# concatenate both pathways
output_tensor = ElementwiseLayer('CONCAT')(normal_path, dilated_path)
# 1x1x1 convolution layer
for n_features in self.fc_features:
conv_fc = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=1,
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_1x1x1_{}'.format(n_features))
output_tensor = conv_fc(output_tensor,
is_training,
keep_prob=keep_prob)
print('#----------------------------------- keep_prob: ',
keep_prob)
print(conv_fc)
# classification layer
for n_features in self.conv_classification:
conv_classification = ConvolutionalLayer(
n_output_chns=n_features,
kernel_size=1,
acti_func=None,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_1x1x1_{}'.format(n_features))
output_tensor = conv_classification(output_tensor, is_training)
print(conv_classification)
return output_tensor
def layer_op(self, images, is_training=True, layer_id=-1, **unused_kwargs):
# image_size-4 should be divisible by 8
#assert layer_util.check_spatial_dims(images, lambda x: x % 16 == 4)
#assert layer_util.check_spatial_dims(images, lambda x: x >= 89)
block_layer = UNetBlock('DOWNSAMPLE_ANISOTROPIC',
(self.n_features[0], self.n_features[1]),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d0')
pool_1, conv_1 = block_layer(images, is_training)
print(block_layer)
block_layer = UNetBlock('DOWNSAMPLE_ANISOTROPIC',
(self.n_features[1], self.n_features[2]),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d1')
pool_2, conv_2 = block_layer(pool_1, is_training)
print(block_layer)
block_layer = UNetBlock('DOWNSAMPLE_ANISOTROPIC',
(self.n_features[2], self.n_features[3]),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d2')
pool_3, conv_3 = block_layer(pool_2, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE_ANISOTROPIC',
(self.n_features[3], self.n_features[4]),
([3, 3, 3], [3, 3, 3]),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='d3')
up_3, _ = block_layer(pool_3, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE_ANISOTROPIC',
(self.n_features[3], self.n_features[3]),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='u2')
crop_layer = CropLayer(border=[4, 4, 2], name='crop-8x8x0')
concat_3 = ElementwiseLayer('CONCAT')(crop_layer(conv_3), up_3)
up_2, _ = block_layer(concat_3, is_training)
print(block_layer)
block_layer = UNetBlock('UPSAMPLE_ANISOTROPIC',
(self.n_features[2], self.n_features[2]),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=False,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='u1')
crop_layer = CropLayer(border=[16, 16, 2], name='crop-32x32x0')
concat_2 = ElementwiseLayer('CONCAT')(crop_layer(conv_2), up_2)
up_1, _ = block_layer(concat_2, is_training)
print(block_layer)
block_layer = UNetBlock(
'NONE', (self.n_features[1], self.n_features[1], self.num_classes),
([3, 3, 1], [3, 3, 1]),
with_downsample_branch=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='u0')
crop_layer = CropLayer(border=[40, 40, 2], name='crop-80x80x0')
concat_1 = ElementwiseLayer('CONCAT')(crop_layer(conv_1), up_1)
print(block_layer)
# for the last layer, upsampling path is not used
_, output_tensor = block_layer(concat_1, is_training)
output_conv_op = ConvolutionalLayer(
n_output_chns=self.num_classes,
kernel_size=1,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=None,
name='{}'.format(self.num_classes),
padding='VALID',
with_bn=False,
with_bias=True)
final_output_tensor = output_conv_op(output_tensor, is_training)
print(output_conv_op)
return final_output_tensor
