def layer_op(self, main_flow, bypass_flow):
"""
:param main_flow: tensor, input to the VNet block
:param bypass_flow: tensor, input from skip connection
:return: res_flow is tensor before final block operation (for residual connections),
main_flow is final output tensor
"""
for i in range(self.n_conv):
main_flow = ConvLayer(name='conv_{}'.format(i),
n_output_chns=self.n_feature_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=5)(main_flow)
if i < self.n_conv - 1: # no activation for the last conv layer
main_flow = ActiLayer(
func=self.acti_func,
regularizer=self.regularizers['w'])(main_flow)
res_flow = ElementwiseLayer('SUM')(main_flow, bypass_flow)
if self.func == 'DOWNSAMPLE':
main_flow = ConvLayer(name='downsample',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=2,
stride=2,
with_bias=True)(res_flow)
elif self.func == 'UPSAMPLE':
main_flow = DeconvLayer(name='upsample',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=2,
stride=2,
with_bias=True)(res_flow)
elif self.func == 'SAME':
main_flow = ConvLayer(name='conv_1x1x1',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
kernel_size=1,
with_bias=True)(res_flow)
main_flow = ActiLayer(self.acti_func)(main_flow)
print(self)
return res_flow, main_flow
def layer_op(self, roots):
"""
Performs channel-wise merging of input tensors
:param roots: tensors to be merged
:return: fused tensor
"""
if self.func == 'MAXOUT':
return tf.reduce_max(tf.stack(roots, axis=-1), axis=-1)
elif self.func == 'AVERAGE':
return tf.reduce_mean(tf.stack(roots, axis=-1), axis=-1)
elif self.func == 'WEIGHTED_AVERAGE':
input_tensor = tf.stack(roots, axis=-1)
rank = input_tensor.shape.ndims
perm = [i for i in range(rank)]
perm[-2], perm[-1] = perm[-1], perm[-2]
output_tensor = input_tensor
output_tensor = tf.transpose(output_tensor, perm=perm)
output_tensor = tf.unstack(output_tensor, axis=-1)
roots_merged = []
for f in range(len(output_tensor)):
conv_layer = ConvLayer(n_output_chns=1,
kernel_size=1,
stride=1)
roots_merged_f = conv_layer(output_tensor[f])
roots_merged.append(roots_merged_f)
return tf.concat(roots_merged, axis=-1)
def layer_op(self, input_tensor, is_training):
output_tensor = input_tensor
for i in range(len(self.kernels)):
# create parameterised layers
bn_op = BNLayer(
regularizer=self.
regularizers['w'], # Add regulizer for samplicity
name='bn_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=self.n_output_chns,
kernel_size=self.kernels[i],
stride=self.strides[i],
dilation=self.dilation_rates[i],
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
output_tensor = conv_op(output_tensor)
output_tensor = acti_op(output_tensor)
output_tensor = bn_op(
output_tensor, is_training
) # Construct operation first and then connect them.
# make residual connections
if self.with_res:
# The input is directly added to the output.
output_tensor = ElementwiseLayer('SUM')(output_tensor,
input_tensor)
return output_tensor
def layer_op(self, input_tensor):
output_tensor = input_tensor
for (i, k) in enumerate(self.kernels):
# create parameterised layers
if self.encoding:
if i == 0:
nb_channels = self.n_output_chns
elif i == 1:
nb_channels = self.n_output_chns
else:
if self.double_n:
if i == 0:
nb_channels = self.n_output_chns
elif i == 1:
nb_channels = int(self.n_output_chns / 2)
else:
nb_channels = int(self.n_output_chns / 2)
in_op = InstanceNormLayer(name='in_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=nb_channels,
kernel_size=k,
stride=self.stride,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
# connect layers
output_tensor = in_op(output_tensor)
output_tensor = acti_op(output_tensor)
output_tensor = conv_op(output_tensor)
return output_tensor
def layer_op(self, param_flow, bypass_flow):
n_param_flow = param_flow.get_shape()[-1]
n_bypass_flow = bypass_flow.get_shape()[-1]
spatial_rank = layer_util.infer_spatial_rank(param_flow)
output_tensor = param_flow
if self.func == 'SUM':
if n_param_flow > n_bypass_flow: # pad the channel dim
pad_1 = np.int((n_param_flow - n_bypass_flow) // 2)
pad_2 = np.int(n_param_flow - n_bypass_flow - pad_1)
padding_dims = np.vstack(
([[0, 0]], [[0, 0]] * spatial_rank, [[pad_1, pad_2]]))
bypass_flow = tf.pad(tensor=bypass_flow,
paddings=padding_dims.tolist(),
mode='CONSTANT')
elif n_param_flow < n_bypass_flow: # make a projection
projector = ConvLayer(n_output_chns=n_param_flow,
kernel_size=1,
stride=1,
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='proj')
bypass_flow = projector(bypass_flow)
# element-wise sum of both paths
output_tensor = param_flow + bypass_flow
elif self.func == 'CONCAT':
output_tensor = tf.concat([param_flow, bypass_flow], axis=-1)
return output_tensor
def layer_op_selu(self, input_tensor, is_training):
output_tensor = input_tensor
for i in range(len(self.kernels)):
# create parameterised layers
# bn_op = BNLayer(regularizer=self.regularizers['w'],
# name='bn_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=self.n_output_chns,
kernel_size=self.kernels[i],
stride=self.strides[i],
dilation=self.dilation_rates[i],
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
# connect layers
# output_tensor = bn_op(output_tensor, is_training)
output_tensor = conv_op(output_tensor)
output_tensor = acti_op(output_tensor)
# make residual connections
if self.with_res:
output_tensor = ElementwiseLayer('SUM')(output_tensor,
input_tensor)
return output_tensor
def layer_op(self, input_tensor, is_training):
"""
:param input_tensor: tensor, input to the network
:param is_training: boolean, True if network is in training mode
:return: tensor, output of the residual block
"""
output_tensor = input_tensor
for (i, k) in enumerate(self.kernels):
# create parameterised layers
bn_op = BNLayer(regularizer=self.regularizers['w'],
name='bn_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=self.n_output_chns,
kernel_size=k,
stride=1,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
# connect layers
output_tensor = bn_op(output_tensor, is_training)
output_tensor = acti_op(output_tensor)
output_tensor = conv_op(output_tensor)
# make residual connections
if self.with_res:
output_tensor = ElementwiseLayer('SUM')(output_tensor,
input_tensor)
return output_tensor
def layer_op(self, input_tensor, is_training):
output_tensor = input_tensor
for (i, k) in enumerate(self.kernels):
# Batch Normalization is removed from the residual blocks.
# create parameterised layers
# bn_op = BNLayer(regularizer=self.regularizers['w'],
# name='bn_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=self.n_output_chns,
kernel_size=k,
stride=1,
padding='SAME',
with_bias=True,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
# connect layers
# output_tensor = bn_op(output_tensor, is_training)
output_tensor = acti_op(output_tensor)
output_tensor = conv_op(output_tensor)
# make residual connections
if self.with_res:
output_tensor = ElementwiseLayer('SUM')(output_tensor, input_tensor)
return output_tensor
def layer_op(self, main_flow, bypass_flow):
for i in range(self.n_conv):
main_flow = ConvLayer(name='conv_{}'.format(i),
n_output_chns=self.n_feature_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=5)(main_flow)
if i < self.n_conv - 1: # no activation for the last conv layer
main_flow = ActiLayer(
func=self.acti_func,
regularizer=self.regularizers['w'])(main_flow)
res_flow = ElementwiseLayer('SUM')(main_flow, bypass_flow)
if self.func == 'DOWNSAMPLE':
main_flow = ConvLayer(name='downsample',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=2,
stride=2,
with_bias=True)(res_flow)
elif self.func == 'UPSAMPLE':
main_flow = DeconvLayer(name='upsample',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
kernel_size=2,
stride=2,
with_bias=True)(res_flow)
elif self.func == 'SAME':
main_flow = ConvLayer(name='conv_1x1x1',
n_output_chns=self.n_output_chns,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
kernel_size=1,
with_bias=True)(res_flow)
main_flow = ActiLayer(self.acti_func)(main_flow)
print(self)
print('VNet is running')
return res_flow, main_flow
def _test_conv_output_shape(self, rank, param_dict, output_shape):
if rank == 2:
input_data = self.get_2d_input()
elif rank == 3:
input_data = self.get_3d_input()
conv_layer = ConvLayer(**param_dict)
output_data = conv_layer(input_data)
print(conv_layer)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
output_value = sess.run(output_data)
self.assertAllClose(output_shape, output_value.shape)
def layer_op(self, roots):
if self.func == 'MAXOUT':
return tf.reduce_max(tf.stack(roots, axis=-1), axis=-1)
elif self.func == 'AVERAGE':
return tf.reduce_mean(tf.stack(roots, axis=-1), axis=-1)
elif self.func == 'WEIGHTED_AVERAGE':
input_tensor = tf.stack(roots, axis=-1)
rank = input_tensor.get_shape().ndims
perm = [i for i in range(rank)]
perm[-2], perm[-1] = perm[-1], perm[-2]
output_tensor = input_tensor
output_tensor = tf.transpose(output_tensor, perm=perm)
output_tensor = tf.unstack(output_tensor, axis=-1)
roots_merged = []
for f in range(len(output_tensor)):
conv_layer = ConvLayer(
n_output_chns=1, kernel_size=1, stride=1)
roots_merged_f = conv_layer(output_tensor[f])
roots_merged.append(roots_merged_f)
return tf.concat(roots_merged, axis=-1)
def _res_bottleneck(self,
inputs,
filters,
kernel,
t,
s,
r=False,
is_training=True):
tchannel = tf.keras.backend.int_shape(inputs)[-1] * t
#conv layer only (no activation or dropout)
conv_block = ConvLayer(kernel_size=1,
n_output_chns=tchannel,
stride=1,
padding='same',
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
x = conv_block(inputs)
x = ConvolutionalLayer(kernel,
stride=s,
padding='same',
conv_type='DEPTHWISE_2D',
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)(
x, is_training=is_training)
conv_block_2 = ConvolutionalLayer(filters,
kernel_size=1,
stride=1,
padding='same',
acti_func=None,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
x = conv_block_2(x, is_training=is_training)
if r:
x = tf.math.add(x, inputs)
return x
def layer_op(self, images, is_training=True, layer_id=-1, **unused_kwargs):
concat_list = [images]
w, h = images.shape[1], images.shape[2]
for bin_size in self.bin_sizes:
x = tf.layers.AveragePooling2D(pool_size=(w // bin_size,
h // bin_size),
strides=(w // bin_size,
h // bin_size))(images)
x = ConvLayer(128,
kernel_size=3,
stride=2,
padding='same',
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)(x)
x = tf.image.resize_images(x, (w, h))
concat_list.append(x)
x = tf.concat(concat_list, axis=-1)
return x
def layer_op(self, images, is_training):
block1_1 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_1')
block1_2 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_2')
block2_1 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_1')
block2_2 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_2')
block3_1 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 1, 1], [1, 1, 1]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_1')
block3_2 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_2')
block4_1 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_1')
block4_2 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_2')
fuse1 = ConvolutionalLayer(
self.base_chns[0],
kernel_size=[3, 1, 1], # Convolution on intra layers
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='fuse1')
downsample1 = ConvolutionalLayer(self.base_chns[0],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='downsample1')
fuse2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='fuse2')
downsample2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='downsample2')
fuse3 = ConvolutionalLayer(self.base_chns[2],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='fuse3')
fuse4 = ConvolutionalLayer(self.base_chns[3],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='fuse4')
feature_expand1 = ConvolutionalLayer(
self.base_chns[1], # Output channels
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='feature_expand1')
feature_expand2 = ConvolutionalLayer(
self.base_chns[2],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='feature_expand2')
feature_expand3 = ConvolutionalLayer(
self.base_chns[3],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='feature_expand3')
centra_slice1 = TensorSliceLayer(margin=2)
centra_slice2 = TensorSliceLayer(margin=1)
image_resize1 = ImageResize()
image_resize2 = ImageResize()
pred_up1 = DeconvolutionalLayer(self.num_classes,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='pred_up1')
pred_up2_1 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='pred_up2_1')
pred_up2_2 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='pred_up2_2')
pred_up3_1 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='pred_up3_1')
pred_up3_2 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
name='pred_up3_2')
final_pred = ConvLayer(
self.num_classes, # Output two class: target and background
kernel_size=[1, 3, 3],
padding=
'SAME', # Same: keep shape; Valid: only get pixels with valid calculation.
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
name='final_pred')
f1 = images
f1 = block1_1(f1, is_training)
f1 = block1_2(f1, is_training)
f1 = fuse1(f1, is_training)
f1 = downsample1(f1, is_training)
img_resize1 = image_resize1(images, f1)
img_resize1 = centra_slice2(img_resize1)
f1 = tf.concat([img_resize1, f1], axis=4, name='concate')
if (self.base_chns[0] != self.base_chns[1]):
f1 = feature_expand1(
f1, is_training
) # To keep same channel number in cascaded netblocks
f1 = block2_1(f1, is_training)
f1 = block2_2(f1, is_training)
f1 = fuse2(f1, is_training)
f2 = downsample2(f1, is_training)
img_resize2 = image_resize2(images, f2)
img_resize2 = centra_slice1(img_resize2)
f2 = tf.concat([img_resize2, f2], axis=4, name='concate')
if (self.base_chns[1] != self.base_chns[2]):
f2 = feature_expand2(f2, is_training)
f2 = block3_1(f2, is_training)
f2 = block3_2(f2, is_training)
f2 = fuse3(f2, is_training)
f3 = f2
if (self.base_chns[2] != self.base_chns[3]):
f3 = feature_expand3(f3, is_training)
f3 = block4_1(f3, is_training)
f3 = block4_2(f3, is_training)
f3 = fuse4(f3, is_training)
p1 = centra_slice1(f1)
p1 = pred_up1(p1, is_training)
p2 = centra_slice2(f2)
p2 = pred_up2_1(p2, is_training)
p2 = pred_up2_2(p2, is_training)
p3 = pred_up3_1(f3, is_training)
p3 = pred_up3_2(p3, is_training)
cat = tf.concat([p1, p2, p3], axis=4, name='concate')
pred = final_pred(cat)
return pred
def layer_op(self, images, is_training):
block1_1 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_1')
block1_2 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_2')
block2_1 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_1')
block2_2 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_2')
block3_1 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 1, 1], [1, 1, 1]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_1')
block3_2 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_2')
block3_3 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_3')
block4_1 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_1')
block4_2 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_2')
block4_3 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 1, 1], [1, 1, 1]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_3')
fuse1 = ConvolutionalLayer(self.base_chns[0],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='fuse1')
downsample1 = ConvolutionalLayer(self.base_chns[0],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='downsample1')
fuse2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='fuse2')
downsample2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='downsample2')
fuse3 = ConvolutionalLayer(self.base_chns[2],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='fuse3')
fuse4 = ConvolutionalLayer(self.base_chns[3],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='fuse4')
feature_expand1 = ConvolutionalLayer(
self.base_chns[1],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='feature_expand1')
feature_expand2 = ConvolutionalLayer(
self.base_chns[2],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='feature_expand2')
feature_expand3 = ConvolutionalLayer(
self.base_chns[3],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='feature_expand3')
centra_slice1 = TensorSliceLayer(margin=2)
centra_slice2 = TensorSliceLayer(margin=1)
pred_up1 = DeconvolutionalLayer(self.num_classes,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='pred_up1')
pred_up2_1 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='pred_up2_1')
pred_up2_2 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='pred_up2_2')
pred_up3_1 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='pred_up3_1')
pred_up3_2 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=self.acti_func != 'selu',
name='pred_up3_2')
final_pred = ConvLayer(self.num_classes,
kernel_size=[1, 3, 3],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
name='final_pred')
f1 = images
f1 = block1_1(f1, is_training=is_training)
f1 = block1_2(f1, is_training=is_training)
f1 = fuse1(f1, is_training=is_training)
f1 = downsample1(f1, is_training=is_training)
if self.base_chns[0] != self.base_chns[1]:
f1 = feature_expand1(f1, is_training=is_training)
f1 = block2_1(f1, is_training=is_training)
f1 = block2_2(f1, is_training=is_training)
f1 = fuse2(f1, is_training=is_training)
f2 = downsample2(f1, is_training=is_training)
if self.base_chns[1] != self.base_chns[2]:
f2 = feature_expand2(f2, is_training=is_training)
f2 = block3_1(f2, is_training=is_training)
f2 = block3_2(f2, is_training=is_training)
f2 = block3_3(f2, is_training=is_training)
f2 = fuse3(f2, is_training=is_training)
f3 = f2
if self.base_chns[2] != self.base_chns[3]:
f3 = feature_expand3(f3, is_training)
f3 = block4_1(f3, is_training=is_training)
f3 = block4_2(f3, is_training=is_training)
f3 = block4_3(f3, is_training=is_training)
f3 = fuse4(f3, is_training=is_training)
p1 = centra_slice1(f1)
p1 = pred_up1(p1, is_training=is_training)
p2 = centra_slice2(f2)
p2 = pred_up2_1(p2, is_training=is_training)
p2 = pred_up2_2(p2, is_training=is_training)
p3 = pred_up3_1(f3, is_training=is_training)
p3 = pred_up3_2(p3, is_training=is_training)
cat = tf.concat([p1, p2, p3], axis=4, name='concate')
pred = final_pred(cat)
return pred
def layer_op(self,
images,
is_training=True,
layer_id=-1,
keep_prob=0.7,
**unused_kwargs):
print('learning downsample ...')
# >>>>>>>>>>>>>>>> learning down sample
lds1 = ConvolutionalLayer(32,
conv_type='REGULAR',
kernel_size=3,
stride=2,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
lds2 = ConvolutionalLayer(48,
conv_type='SEPARABLE_2D',
kernel_size=3,
stride=2,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
lds3 = ConvolutionalLayer(64,
conv_type='SEPARABLE_2D',
kernel_size=3,
stride=2)
flow = lds1(images, is_training=is_training)
flow = lds2(flow, is_training=is_training)
flow = lds3(flow, is_training=is_training)
lds = flow
# >>>>>>>>>>>>>>>> global feature extraction
print('global feature extractor ...')
bottle1 = SCCNBottleneckBlock(64,
3,
t=6,
stride=2,
n=3,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
bottle2 = SCCNBottleneckBlock(96,
3,
t=6,
stride=2,
n=3,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
bottle3 = SCCNBottleneckBlock(128,
3,
t=6,
stride=1,
n=3,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
pyramid = SCNNPyramidBlock([2, 4, 6, 8],
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
flow = bottle1(flow)
flow = bottle2(flow)
flow = bottle3(flow)
flow = pyramid(flow)
gfe = flow
# >>>>>>>>>>>>>>>> feature fusion
print('Feature fusion ...')
conv1 = ConvolutionalLayer(128,
conv_type='REGULAR',
kernel_size=1,
padding='same',
stride=1,
acti_func=None,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
upsample1 = tf.keras.layers.UpSampling2D((4, 4),
interpolation='bilinear')
dwconv = ConvolutionalLayer(1,
conv_type='DEPTHWISE_2D',
kernel_size=3,
stride=1,
padding='same',
acti_func=self.acti_func,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
conv2 = ConvLayer(128,
conv_type='REGULAR',
kernel_size=1,
padding='same',
stride=1,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
bn = BNLayer()
acti = ActiLayer(func=self.acti_func,
regularizer=self.w_regularizer,
name='acti_')
flow1 = conv1(lds, is_training=is_training)
flow2 = upsample1(gfe)
flow2 = dwconv(flow2, is_training=is_training)
flow2 = conv2(flow2)
flow = tf.math.add(flow1, flow2)
flow = bn(flow, is_training=is_training)
flow = acti(flow)
# ff = flow
# >>>>>>>>>>>>>>>> classifier
sep_conv1 = ConvolutionalLayer(128,
conv_type='SEPARABLE_2D',
kernel_size=3,
padding='same',
stride=1,
name='DSConv1_classifier',
acti_func=self.acti_func,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
sep_conv2 = ConvolutionalLayer(128,
conv_type='SEPARABLE_2D',
kernel_size=3,
padding='same',
stride=1,
name='DSConv2_classifier',
acti_func=self.acti_func,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
flow = sep_conv1(flow, is_training=is_training)
flow = sep_conv2(flow, is_training=is_training)
conv = ConvolutionalLayer(self.num_classes,
conv_type='REGULAR',
kernel_size=1,
padding='same',
stride=1,
w_initializer=self.w_initializer,
w_regularizer=self.w_regularizer)
dropout = ActiLayer(func='dropout',
regularizer=self.w_regularizer,
name='dropout_')
# tf.keras.layers.Dropout(0.3)
upsample = tf.keras.layers.UpSampling2D((8, 8),
interpolation='bilinear')
flow = conv(flow, is_training=is_training)
flow = dropout(flow, keep_prob=keep_prob)
flow = upsample(flow)
flow = tf.nn.softmax(flow)
return flow
def layer_op(self, images, is_training, bn_momentum=0.9, layer_id=-1):
block1 = VGGBlock((self.n_features[0], self.n_features[0]),
(self.dilations[0], self.dilations[0]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B1')
block2 = VGGBlock((self.n_features[1], self.n_features[1]),
(self.dilations[1], self.dilations[1]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B2')
block3 = VGGBlock((self.n_features[2], self.n_features[2]),
(self.dilations[2], self.dilations[2]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B3')
block4 = VGGBlock((self.n_features[3], self.n_features[3],
self.n_features[3], self.n_features[3]),
(self.dilations[3], self.dilations[3],
self.dilations[3], self.dilations[3]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B4')
block5 = VGGBlock((self.n_features[4], self.n_features[4],
self.n_features[4], self.n_features[4]),
(self.dilations[4], self.dilations[4],
self.dilations[4], self.dilations[4]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B5')
block6 = VGGBlock((self.n_features[5], self.n_features[5],
self.n_features[5], self.n_features[5]),
(self.dilations[5], self.dilations[5],
self.dilations[5], self.dilations[5]),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='B6')
fc1 = ConvolutionalLayer(n_output_chns=self.n_features[6],
kernel_size=[1, 6, 6],
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='fc_1')
fc2 = ConvolutionalLayer(n_output_chns=self.n_features[6],
kernel_size=[1, 1, 1],
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
moving_decay=self.moving_decay,
acti_func=self.acti_func,
name='fc_2')
fc3 = ConvLayer(n_output_chns=self.num_classes,
kernel_size=[1, 1, 1],
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
with_bias=True,
name='fc_3')
out = block1(images, is_training, bn_momentum)
out = block2(out, is_training, bn_momentum)
out = block3(out, is_training, bn_momentum)
out = block4(out, is_training, bn_momentum)
out = block5(out, is_training, bn_momentum)
out = block6(out, is_training, bn_momentum)
out = fc1(out, is_training, bn_momentum)
out = fc2(out, is_training, bn_momentum)
out = fc3(out)
return out
def layer_op(self, images, is_training, bn_momentum=0.9, layer_id=-1):
# image_size should be divisible by 8
# spatial_dims = images.get_shape()[1:-1].as_list()
# assert (spatial_dims[-2] % 16 == 0 )
# assert (spatial_dims[-1] % 16 == 0 )
block1 = UNetBlock((self.n_features[0], self.n_features[0]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B1')
block2 = UNetBlock((self.n_features[1], self.n_features[1]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B2')
block3 = UNetBlock((self.n_features[2], self.n_features[2]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B3')
block4 = UNetBlock((self.n_features[3], self.n_features[3]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B4')
block5 = UNetBlock((self.n_features[4], self.n_features[4]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B5')
block6 = UNetBlock((self.n_features[3], self.n_features[3]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B6')
block7 = UNetBlock((self.n_features[2], self.n_features[2]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B7')
block8 = UNetBlock((self.n_features[1], self.n_features[1]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B8')
block9 = UNetBlock((self.n_features[0], self.n_features[0]),
((1, 3, 3), (1, 3, 3)),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
acti_func=self.acti_func,
name='B9')
conv = ConvLayer(n_output_chns=self.num_classes,
kernel_size=(1, 1, 1),
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
with_bias=True,
name='conv')
down1 = DownSampleLayer('MAX',
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='down1')
down2 = DownSampleLayer('MAX',
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='down2')
down3 = DownSampleLayer('MAX',
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='down3')
down4 = DownSampleLayer('MAX',
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='down4')
up1 = DeconvolutionalLayer(n_output_chns=self.n_features[3],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='up1')
up2 = DeconvolutionalLayer(n_output_chns=self.n_features[2],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='up2')
up3 = DeconvolutionalLayer(n_output_chns=self.n_features[1],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='up3')
up4 = DeconvolutionalLayer(n_output_chns=self.n_features[0],
kernel_size=(1, 2, 2),
stride=(1, 2, 2),
name='up4')
f1 = block1(images, is_training, bn_momentum)
d1 = down1(f1)
f2 = block2(d1, is_training, bn_momentum)
d2 = down2(f2)
f3 = block3(d2, is_training, bn_momentum)
d3 = down3(f3)
f4 = block4(d3, is_training, bn_momentum)
d4 = down4(f4)
f5 = block5(d4, is_training, bn_momentum)
# add dropout to the original version
f5 = tf.nn.dropout(f5, self.dropout)
f5up = up1(f5, is_training, bn_momentum)
f4cat = tf.concat((f4, f5up), axis=-1)
f6 = block6(f4cat, is_training, bn_momentum)
# add dropout to the original version
f6 = tf.nn.dropout(f6, self.dropout)
f6up = up2(f6, is_training, bn_momentum)
f3cat = tf.concat((f3, f6up), axis=-1)
f7 = block7(f3cat, is_training, bn_momentum)
# add dropout to the original version
f7 = tf.nn.dropout(f7, self.dropout)
f7up = up3(f7, is_training, bn_momentum)
f2cat = tf.concat((f2, f7up), axis=-1)
f8 = block8(f2cat, is_training, bn_momentum)
# add dropout to the original version
f8 = tf.nn.dropout(f8, self.dropout)
f8up = up4(f8, is_training, bn_momentum)
f1cat = tf.concat((f1, f8up), axis=-1)
f9 = block9(f1cat, is_training, bn_momentum)
# add dropout to the original version
f9 = tf.nn.dropout(f9, self.dropout)
output = conv(f9)
return output
def layer_op(self, images, is_training, input_image_shape):
"""
images: a 5D tensor of [B, D, H, W, N]
is_training: a bool value
input_image_shape: a tensor of [D, H, W]
"""
d = tf.slice(input_image_shape, [0], [1])
h = tf.slice(input_image_shape, [1], [1])
w = tf.slice(input_image_shape, [2], [1])
up_shape1 = tf.concat([d - 8, h / 4, w / 4], axis=0)
up_shape2 = tf.concat([d - 8, h / 2, w / 2], axis=0)
block1_1 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_1')
block1_2 = ResBlock(self.base_chns[0],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block1_2')
block2_1 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_1')
block2_2 = ResBlock(self.base_chns[1],
kernels=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block2_2')
block3_1 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 1, 1], [1, 1, 1]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_1')
block3_2 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_2')
block3_3 = ResBlock(self.base_chns[2],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block3_3')
block4_1 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 3, 3], [1, 3, 3]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_1')
block4_2 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 2, 2], [1, 2, 2]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_2')
block4_3 = ResBlock(self.base_chns[3],
kernels=[[1, 3, 3], [1, 3, 3]],
dilation_rates=[[1, 1, 1], [1, 1, 1]],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='block4_3')
fuse1 = ConvolutionalLayer(self.base_chns[0],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='fuse1')
downsample1 = ConvolutionalLayer(self.base_chns[0],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='downsample1')
fuse2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='fuse2')
downsample2 = ConvolutionalLayer(self.base_chns[1],
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='downsample2')
fuse3 = ConvolutionalLayer(self.base_chns[2],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='fuse3')
fuse4 = ConvolutionalLayer(self.base_chns[3],
kernel_size=[3, 1, 1],
padding='VALID',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='fuse4')
feature_expand1 = ConvolutionalLayer(
self.base_chns[1],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='feature_expand1')
feature_expand2 = ConvolutionalLayer(
self.base_chns[2],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='feature_expand2')
feature_expand3 = ConvolutionalLayer(
self.base_chns[3],
kernel_size=[1, 1, 1],
stride=[1, 1, 1],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='feature_expand3')
centra_slice1 = TensorSliceLayer(margin=2)
centra_slice2 = TensorSliceLayer(margin=1)
pred1 = ConvLayer(self.num_classes,
kernel_size=[1, 3, 3],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
name='pred1')
pred_up1 = DeconvolutionalLayer(self.num_classes,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='pred_up1')
pred_up2_1 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='pred_up2_1')
pred_up2_2 = DeconvolutionalLayer(self.num_classes * 2,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='pred_up2_2')
pred_up3_1 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='pred_up3_1')
pred_up3_2 = DeconvolutionalLayer(self.num_classes * 4,
kernel_size=[1, 3, 3],
stride=[1, 2, 2],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
acti_func=self.acti_func,
with_bn=True,
name='pred_up3_2')
final_pred = ConvLayer(self.num_classes,
kernel_size=[1, 3, 3],
padding='SAME',
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
b_initializer=self.initializers['b'],
b_regularizer=self.regularizers['b'],
name='final_pred')
f1 = images
f1 = block1_1(f1, is_training)
f1 = block1_2(f1, is_training)
f1 = fuse1(f1, is_training)
if (self.downsample_twice):
f1 = downsample1(f1, is_training)
if (self.base_chns[0] != self.base_chns[1]):
f1 = feature_expand1(f1, is_training)
f1 = block2_1(f1, is_training)
f1 = block2_2(f1, is_training)
f1 = fuse2(f1, is_training)
f2 = downsample2(f1, is_training)
if (self.base_chns[1] != self.base_chns[2]):
f2 = feature_expand2(f2, is_training)
f2 = block3_1(f2, is_training)
f2 = block3_2(f2, is_training)
f2 = block3_3(f2, is_training)
f2 = fuse3(f2, is_training)
f3 = f2
if (self.base_chns[2] != self.base_chns[3]):
f3 = feature_expand3(f3, is_training)
f3 = block4_1(f3, is_training)
f3 = block4_2(f3, is_training)
f3 = block4_3(f3, is_training)
f3 = fuse4(f3, is_training)
p1 = centra_slice1(f1)
if (self.downsample_twice):
p1 = pred_up1(p1, is_training, input_image_shape=up_shape2)
else:
p1 = pred1(p1)
p2 = centra_slice2(f2)
if (self.downsample_twice):
p2 = pred_up2_1(p2, is_training, input_image_shape=up_shape1)
p2 = pred_up2_2(p2, is_training, input_image_shape=up_shape2)
else:
p2 = pred_up2_1(p2, is_training, input_image_shape=up_shape2)
p3 = f3
if (self.downsample_twice):
p3 = pred_up3_1(p3, is_training, input_image_shape=up_shape1)
p3 = pred_up3_2(p3, is_training, input_image_shape=up_shape2)
else:
p3 = pred_up3_1(p3, is_training, input_image_shape=up_shape2)
cat = tf.concat([p1, p2, p3], axis=4, name='concate')
pred = final_pred(cat)
return pred
def _test_extended_conv(self, orig_input, init_dict):
"""
Tests the extended padding options of ConvLayer
"""
def _w_init(shape, dtype=tf.float32, **kwargs):
data = np.arange(ft.reduce(lambda prod, x: prod*x, shape, 1))\
.astype(np.float32)
data *= 2.374/data.mean()
data -= data.mean()
return tf.constant(data.reshape(shape), dtype=dtype)
def _b_init(shape, dtype=tf.float32, **kwargs):
data = np.arange(shape[0]).astype(np.float32)
data *= 0.273/data.mean()
data -= data.mean()
return tf.constant(data.reshape(shape), dtype=dtype)
init_dict['w_initializer'] = _w_init
init_dict['b_initializer'] = _b_init
conv_layer = ConvLayer(**init_dict)
small_output = conv_layer(tf.constant(orig_input))
input_shape = orig_input.shape
multiplier = init_dict['kernel_size'] + init_dict['dilation'] \
+ init_dict['stride']
pad = [d*multiplier for d in input_shape[1:-1]]
paddings = [(0, 0)] + [(p, p) for p in pad] + [(0, 0)]
if init_dict['padding'] == 'CONSTANT':
opts = {'constant_values': init_dict.get('padding_constant', 0)}
else:
opts = {}
enlarged_input = np.pad(orig_input,
paddings,
init_dict['padding'].lower(),
**opts)
conv_layer.padding = 'SAME'
large_output = conv_layer(tf.constant(enlarged_input))
def _extract_valid_region(output_tensor, target_tensor):
output_shape = output_tensor.shape
target_shape = target_tensor.shape
extr_slices = []
for d in range(len(target_shape)):
opad = (output_shape[d] - target_shape[d])//2
extr_slices.append(slice(
opad, opad + target_shape[d]))
return output_tensor[tuple(extr_slices)]
assert np.square(
_extract_valid_region(enlarged_input, orig_input) - orig_input).sum() \
<= 1e-6*np.square(orig_input).sum()
with self.cached_session() as sess:
sess.run(tf.global_variables_initializer())
small_value = sess.run(small_output)
large_value = sess.run(large_output)
extr_value = _extract_valid_region(large_value, small_value)
print(np.square(small_value - extr_value).sum()/np.square(extr_value).sum())
self.assertAllClose(small_value, extr_value, rtol=1e-3)
def layer_op(self, input_tensor, is_training):
"""
:param input_tensor: tensor, input to the network
:param is_training: boolean, True if network is in training mode
:return: tensor, output of the autofocus block
"""
output_tensor = input_tensor
########################################################################
# 1: Create first of two autofocus layer of autofocus block.
########################################################################
# A convolution without feature norm and activation.
conv_1 = ConvLayer(n_output_chns = self.n_output_chns[0],
kernel_size = self.kernel_size[0],
padding='SAME',
dilation = 1,
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_1')
# Create two conv layers for the attention model. The output of the
# attention model will be needed for the K parallel conv layers.
# First convolutional layer of the attention model (conv l,1).
conv_att_11 = ConvLayer(n_output_chns = int(self.n_input_chns[0]/2),
kernel_size = self.kernel_size[0],
padding = 'SAME',
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_att_11')
# Second convolutional layer of the attention model (conv l,2).
conv_att_12 = ConvLayer(n_output_chns = self.num_branches,
kernel_size = [1, 1, 1],
padding = 'SAME',
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_att_12')
# Batch norm (BN) layer for each of the K parallel convolutions
bn_layer_1 = []
for i in range(self.num_branches):
bn_layer_1.append(BNLayer(regularizer = self.regularizers['w'],
name = 'bn_layer_1_{}'.format(i)))
# Activation function used in the first attention model
acti_op_1 = ActiLayer(func = self.acti_func,
regularizer = self.regularizers['w'],
name = 'acti_op_1')
########################################################################
# 2: Create second of two autofocus layer of autofocus block.
########################################################################
# A convolution without feature norm and activation.
conv_2 = ConvLayer(n_output_chns = self.n_output_chns[1],
kernel_size = self.kernel_size[1],
padding='SAME',
dilation = 1,
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_2')
# Create two conv layers for the attention model. The output of the
# attention model will be needed for the K parallel conv layers.
# First convolutional layer of the attention model (conv l,1).
conv_att_21 = ConvLayer(n_output_chns = int(self.n_input_chns[1]/2),
kernel_size = self.kernel_size[1],
padding = 'SAME',
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_att_21')
# Second convolutional layer of the attention model (conv l,2).
conv_att_22 = ConvLayer(n_output_chns = self.num_branches,
kernel_size = [1, 1, 1],
padding = 'SAME',
w_initializer = self.initializers['w'],
w_regularizer = self.regularizers['w'],
name = 'conv_att_22')
# Batch norm (BN) layer for each of the K parallel convolutions
bn_layer_2 = []
for i in range(self.num_branches):
bn_layer_2.append(BNLayer(regularizer = self.regularizers['w'],
name = 'bn_layer_2_{}'.format(i)))
# Activation function used in the second attention model
acti_op_2 = ActiLayer(func = self.acti_func,
regularizer = self.regularizers['w'],
name = 'acti_op_2')
########################################################################
# 3: Create other parameterised layers
########################################################################
acti_op = ActiLayer(func = self.acti_func,
regularizer = self.regularizers['w'],
name = 'acti_op')
########################################################################
# 4: Connect layers
########################################################################
# compute attention weights for the K parallel conv layers in the first
# autofocus convolutional layer
feature_1 = output_tensor
att_1 = acti_op_1(conv_att_11(feature_1))
att_1 = conv_att_12(att_1)
att_1 = tf.nn.softmax(att_1, axis=1)
# Create K dilated tensors as input to the autofocus layer. This
# simulates the K parallel convolutions with different dilation
# rates. Doing it this way ensures the required weight sharing.
dilated_tensor_1 = []
for i in range(self.num_branches):
dilated_1 = output_tensor
with DilatedTensor(dilated_1, dilation_factor = self.dilation_list[i]) as dilated:
dilated.tensor = conv_1(dilated.tensor)
dilated.tensor = bn_layer_1[i](dilated.tensor, is_training)
dilated.tensor = dilated.tensor * att_1[:,:,:,:,i:(i+1)]
dilated_tensor_1.append(dilated.tensor)
output_tensor = tf.add_n(dilated_tensor_1)
output_tensor = acti_op(output_tensor)
# compute attention weights for the K parallel conv layers in the second
# autofocus convolutional layer
feature_2 = output_tensor
att_2 = acti_op_2(conv_att_21(feature_2))
att_2 = conv_att_22(att_2)
att_2 = tf.nn.softmax(att_2, axis=1)
# Create K dilated tensors as input to the autofocus layer. This
# simulates the K parallel convolutions with different dilation
# rates. Doing it this way ensures the required weight sharing.
dilated_tensor_2 = []
for i in range(self.num_branches):
dilated_2 = output_tensor
with DilatedTensor(dilated_2, dilation_factor = self.dilation_list[i]) as dilated:
dilated.tensor = conv_2(dilated.tensor)
dilated.tensor = bn_layer_2[i](dilated.tensor, is_training)
dilated.tensor = dilated.tensor * att_2[:,:,:,:,i:(i+1)]
dilated_tensor_2.append(dilated.tensor)
output_tensor = tf.add_n(dilated_tensor_2)
# make residual connection using ElementwiseLayer with SUM
if self.with_res:
output_tensor = ElementwiseLayer('SUM')(output_tensor, input_tensor)
# apply the last ReLU activation
output_tensor = acti_op(output_tensor)
print("output_tensor:", output_tensor)
return output_tensor
