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