diff --git a/Nets/PSMNet.py b/Nets/PSMNet.py new file mode 100644 index 0000000..2df0448 --- /dev/null +++ b/Nets/PSMNet.py @@ -0,0 +1,304 @@ +import tensorflow as tf +import tensorflow.keras as keras +from .utils import * +import numpy as np + +from Nets import Stereo_net +from Nets import sharedLayers +from Data_utils import preprocessing + +class PSMNet(Stereo_net.StereoNet): + _valid_args = [ + ("left_img", "meta op for left image batch"), + ("right_img", "meta op for right image batch"), + ("max_disp", "max_disp for the model") + ] + Stereo_net.StereoNet._valid_args + _netName = "PSMNet" + + def __init__(self, **kwargs): + """ + Creation of a PSMNet for stereo prediction + """ + super(PSMNet, self).__init__(**kwargs) + + def _validate_args(self, args): + """ + Check that args contains everything that is needed + Valid Keys for args: + left_img: left image op + right_img: right image op + """ + super(PSMNet, self)._validate_args(args) + if ("left_img" not in args) or ("right_img" not in args): + raise Exception('Missing input op for left and right images') + if 'max_disp' not in args: + print('WARNING: max_disp not setted, setting default value 192') + args['max_disp'] = 192 + return args + + def _preprocess_inputs(self, args): + self._left_input_batch = args['left_img'] + self._restore_shape = tf.shape(args['left_img'])[1:3] + self._left_input_batch = tf.cast(self._left_input_batch, tf.float32) + self._left_input_batch = preprocessing.pad_image( + self._left_input_batch, 64) + + self._right_input_batch = args['right_img'] + self._right_input_batch = tf.cast(self._right_input_batch, tf.float32) + self._right_input_batch = preprocessing.pad_image( + self._right_input_batch, 64) + self.max_disp = args['max_disp'] + self.image_size_tf = tf.shape(self._left_input_batch)[1:3] + + def _build_network(self, args): + conv4_left = self.CNN(self._left_input_batch) + fusion_left = self.SPP(conv4_left) + + conv4_right = self.CNN(self._right_input_batch, reuse=tf.compat.v1.AUTO_REUSE) + fusion_right = self.SPP(conv4_right, reuse=tf.compat.v1.AUTO_REUSE) + + cost_vol = self.cost_vol(fusion_left, fusion_right, self.max_disp) + outputs = self.CNN3D(cost_vol) + disps = self.output(outputs)#size of (B, H, W),3out + self._disparities.append(disps[2]) + + def CNN(self, bottom, reuse=False): + with tf.compat.v1.variable_scope('CNN',reuse=reuse): + with tf.compat.v1.variable_scope('conv0',reuse=reuse): + if reuse == tf.compat.v1.AUTO_REUSE: + tf.print(bottom, ['left_batch_size:',bottom]) + self._add_to_layers('CNN/conv0/conv0_1', sharedLayers.conv2d(bottom, [3, 3, 3, 32], \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,strides=2, name='conv0_1',reuse=reuse)) + for i in range(1, 3): + bottom = self._get_layer_as_input('CNN/conv0/conv0_'+str(i)) + self._add_to_layers('CNN/conv0/conv0_'+str(i+1), sharedLayers.conv2d(bottom,\ + [3, 3, 32, 32], activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True, strides=1, name='conv0_'+str(i+1),reuse=reuse)) + with tf.compat.v1.variable_scope('conv1',reuse=reuse): + bottom = self._get_layer_as_input('CNN/conv0/conv0_3') + for i in range(3): + bottom = self.res_block(bottom, 32, 3, layer_prefix='CNN/conv1/conv1_%d' % (i+1), reuse=reuse) + with tf.compat.v1.variable_scope('conv2',reuse=reuse): + bottom = self.res_block(bottom, 64, 3, strides=2, layer_prefix='CNN/conv2/conv2_1', reuse=reuse, projection=True) + for i in range(1, 16): + bottom = self.res_block(bottom, 64, 3, layer_prefix='CNN/conv2/conv2_%d' % (i+1), reuse=reuse) + with tf.compat.v1.variable_scope('conv3',reuse=reuse): + bottom = self.res_block(bottom, 128, 3, dilation_rate=1, layer_prefix='CNN/conv3/conv3_1', reuse=reuse,projection=True) + for i in range(1, 3): + bottom = self.res_block(bottom, 128, 3, dilation_rate=1, layer_prefix='CNN/conv3/conv3_%d' % (i+1), reuse=reuse) + with tf.compat.v1.variable_scope('conv4',reuse=reuse): + for i in range(3): + bottom = self.res_block(bottom, 128, 3, dilation_rate=2, layer_prefix='CNN/conv4/conv4_%d' % (i+1), reuse=reuse) + return bottom + + def res_block(self, bottom, filters, kernel_size, strides=1, dilation_rate=1, layer_prefix=None, reuse=None, projection=False): + with tf.compat.v1.variable_scope(layer_prefix,reuse=reuse): + names = [] + names.append('{}/conv1'.format(layer_prefix)) + + short_cut = bottom + input_layer = bottom + self._add_to_layers(names[-1], sharedLayers.conv2d(input_layer, [kernel_size, kernel_size, input_layer.get_shape()[-1].value, filters], \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,strides=strides, name='conv1', dilation_rate=dilation_rate, reuse=reuse)) + + names.append('{}/conv2'.format(layer_prefix)) + input_layer = self._get_layer_as_input(names[-2]) + self._add_to_layers(names[-1], sharedLayers.conv2d(input_layer, [kernel_size, kernel_size, input_layer.get_shape()[-1].value, filters], \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False,strides=1, name='conv2', dilation_rate=dilation_rate, reuse=reuse)) + + if projection: + names.append('{}/projection'.format(layer_prefix)) + self._add_to_layers(names[-1], sharedLayers.conv2d(short_cut, [1, 1, short_cut.get_shape()[-1].value, filters], \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True, strides=strides, name='projection', reuse=reuse)) + short_cut = self._get_layer_as_input(names[-1]) + + if projection: + input_layer = self._get_layer_as_input(names[-2]) + else: + input_layer = self._get_layer_as_input(names[-1]) + bottom = tf.add(input_layer, short_cut, 'add') + names.append('{}/leaky_relu'.format(layer_prefix)) + self._add_to_layers(names[-1],tf.nn.leaky_relu(bottom, name='leaky_relu')) + return self._get_layer_as_input(names[-1]) + + def SPP(self, bottom, reuse=False): + with tf.compat.v1.variable_scope('SPP',reuse=reuse): + branches = [] + for i, p in enumerate([64, 32, 16, 8]): + branches.append(self.SPP_branch(bottom, p, 32, 3, name='branch_%d' % (i+1), reuse=reuse)) + #if not reuse: + #conv2_16 = tf.compat.v1.get_default_graph().get_tensor_by_name('CNN/conv2/conv2_16/add:0') + #conv4_3 = tf.compat.v1.get_default_graph().get_tensor_by_name('CNN/conv4/conv4_3/add:0') + conv2_16 = self._get_layer_as_input('CNN/conv2/conv2_16/leaky_relu') + conv4_3 = self._get_layer_as_input('CNN/conv4/conv4_3/leaky_relu') + #else: + # conv2_16 = tf.get_default_graph().get_tensor_by_name('CNN_1/conv2/conv2_16/add:0') + # conv4_3 = tf.get_default_graph().get_tensor_by_name('CNN_1/conv4/conv4_3/add:0') + concat = tf.concat([conv2_16, conv4_3] + branches, axis=-1, name='concat') + with tf.compat.v1.variable_scope('fusion',reuse=reuse): + self._add_to_layers('SPP/fusion/conv1', sharedLayers.conv2d(concat, [3, 3, concat.get_shape()[-1].value, 128], name='conv1', \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,reuse=reuse)) + input_layer = self._get_layer_as_input('SPP/fusion/conv1') + self._add_to_layers('SPP/fusion/conv2', sharedLayers.conv2d(input_layer,[1,1,input_layer.get_shape()[-1].value,32], name='conv2', \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,reuse=reuse)) + return self._get_layer_as_input('SPP/fusion/conv2') + + def SPP_branch(self, bottom, pool_size, filters, kernel_size, strides=1, dilation_rate=1, name=None, reuse=None, apply_bn=True, apply_relu=True): + with tf.compat.v1.variable_scope(name,reuse=reuse): + size = tf.shape(bottom)[1:3] + ap = keras.layers.AveragePooling2D(pool_size=(pool_size,pool_size), padding='same', name='avg_pool') + bottom = ap(bottom) + + self._add_to_layers('SPP/'+ name + '/conv', sharedLayers.conv2d(bottom,[kernel_size, kernel_size, bottom.get_shape()[-1].value, filters], \ + strides=strides, name='conv', reuse=reuse, activation=tf.nn.leaky_relu, batch_norm=apply_bn, apply_relu=apply_relu)) + bottom = tf.image.resize(self._get_layer_as_input('SPP/'+ name + '/conv'), size) + return bottom + + def cost_vol(self, left, right, max_disp=192): + with tf.compat.v1.variable_scope('cost_vol'): + disparity_costs = [] + #shape = tf.shape(right) #(N,H,W,F) + for i in range(max_disp // 4): + if i > 0: + left_tensor = keras.backend.spatial_2d_padding(left[:, :, i:, :], padding=((0, 0), (i, 0))) + right_tensor = keras.backend.spatial_2d_padding(right[:, :, :-i, :], padding=((0, 0), (i, 0))) + cost = tf.concat([left_tensor, right_tensor], axis=3) + else: + cost = tf.concat([left, right], axis=3) + disparity_costs.append(cost) + + cost_vol = tf.stack(disparity_costs, axis=1) + return cost_vol + + def CNN3D(self, bottom): + with tf.compat.v1.variable_scope('CNN3D'): + for i in range(2): + self._add_to_layers('CNN3D/3Dconv0_%d' % (i+1), sharedLayers.conv3d(bottom, [3, 3, 3, bottom.get_shape()[-1].value,32], \ + activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True, name='3Dconv0_%d' % (i+1))) + + bottom = self._get_layer_as_input('CNN3D/3Dconv0_2') + _3Dconv1 = self.res_block_3d(bottom, 32, 3, layer_prefix='CNN3D/3Dconv1') + + _3Dstack = [self.hourglass('3d', _3Dconv1, [64, 64, 64, 32], [3, 3, 3, 3], [None, None, -2, _3Dconv1],name='3Dstack1')] + for i in range(1, 3): + _3Dstack.append(self.hourglass('3d', _3Dstack[-1][-1], [64, 64, 64, 32], [3, 3, 3, 3], + [_3Dstack[-1][-2], None, _3Dstack[0][0], _3Dconv1], name='3Dstack%d' % (i+1))) + + self._add_to_layers('output_1_1', sharedLayers.conv3d(_3Dstack[0][3], [3,3,3,_3Dstack[0][3].get_shape()[-1].value,32], name='output_1_1')) + input_layer = self._get_layer_as_input('output_1_1') + self._add_to_layers('output_1', sharedLayers.conv3d(input_layer, [3,3,3,input_layer.get_shape()[-1].value,1],name='output_1', batch_norm=False, apply_relu=False)) + + output_1 = self._get_layer_as_input('output_1') + outputs = [output_1] + + for i in range(1, 3): + self._add_to_layers('output_%d_1' % (i+1), sharedLayers.conv3d(_3Dstack[i][3],[3,3,3,_3Dstack[i][3].get_shape()[-1].value,32], \ + name='output_%d_1' % (i+1))) + input_layer = self._get_layer_as_input('output_%d_1' % (i+1)) + self._add_to_layers('output_%d_2' % (i+1), sharedLayers.conv3d(input_layer,[3,3,3,input_layer.get_shape()[-1].value,1], + name='output_%d_2' % (i+1), batch_norm=False, apply_relu=False)) + prev = self._get_layer_as_input('output_%d_2' % (i+1)) + output = tf.add(prev, outputs[-1], name='output_%d' % (i+1)) + outputs.append(output) + + return outputs + + def res_block_3d(self, bottom, filters, kernel_size, strides=1, dilation_rate=1, layer_prefix=None, reuse=None, projection=False): + with tf.compat.v1.variable_scope(layer_prefix,reuse=reuse): + names = [] + names.append('{}/conv1'.format(layer_prefix)) + + short_cut = bottom + input_layer = bottom + self._add_to_layers(names[-1], sharedLayers.conv3d(input_layer, [kernel_size, kernel_size, kernel_size, input_layer.get_shape()[-1].value, filters], \ + strides=strides, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,name='conv1', reuse=reuse)) + + names.append('{}/conv2'.format(layer_prefix)) + input_layer = self._get_layer_as_input(names[-2]) + self._add_to_layers(names[-1], sharedLayers.conv3d(input_layer, [kernel_size, kernel_size, kernel_size, input_layer.get_shape()[-1].value, filters], \ + strides=1, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False,name='conv2', reuse=reuse)) + + if projection: + names.append('{}/projection'.format(layer_prefix)) + input_layer = short_cut + self._add_to_layers(names[-1], sharedLayers.conv3d(input_layer, [1, 1, 1, input_layer.get_shape()[-1].value, filters], \ + strides=strides, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True,name='projection', reuse=reuse)) + + if projection: + input_layer = self._get_layer_as_input(names[-2]) + else: + input_layer = self._get_layer_as_input(names[-1]) + bottom = tf.add(input_layer, short_cut, 'add') + bottom = tf.nn.relu(bottom, name='relu') + return bottom + + def hourglass(self,strs, bottom, filters_list, kernel_size_list, short_cut_list, dilation_rate=1, name=None): + with tf.compat.v1.variable_scope(name): + output = [] + + for i, (filters, kernel_size, short_cut) in enumerate(zip(filters_list, kernel_size_list, short_cut_list)): + if i < len(filters_list) // 2: + if strs == '2d': + self._add_to_layers(name + '/stack_%d_1' % (i+1), sharedLayers.conv2d(bottom,[kernel_size, kernel_size, \ + bottom.get_shape()[-1].value, filters], strides=2, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True, name='/stack_%d_1' % (i+1))) + bottom = self._get_layer_as_input(name+'/stack_%d_1'%(i+1)) + self._add_to_layers(name + '/stack_%d_2' % (i+1), sharedLayers.conv2d(bottom,[kernel_size, kernel_size, \ + bottom.get_shape()[-1].value, filters], strides=1, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False, name='/stack_%d_2' % (i+1))) + elif strs == '3d': + self._add_to_layers(name + '/stack_%d_1' % (i+1), sharedLayers.conv3d(bottom,[kernel_size, kernel_size, kernel_size, \ + bottom.get_shape()[-1].value, filters], strides=2, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=True, name='/stack_%d_1' % (i+1))) + bottom = self._get_layer_as_input(name+'/stack_%d_1'%(i+1)) + self._add_to_layers(name + '/stack_%d_2' % (i+1), sharedLayers.conv3d(bottom,[kernel_size, kernel_size, kernel_size, \ + bottom.get_shape()[-1].value, filters], strides=1, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False, name='/stack_%d_2' % (i+1))) + + if short_cut is not None: + if type(short_cut) is int: + short_cut = output[short_cut] + bottom = tf.add(bottom, short_cut, name='stack_%d' % (i+1)) + bottom = tf.nn.leaky_relu(bottom, name='relu') + else: + #反卷积有问题,必须确定batch-size,height,weight才能正确运行 + #第三层相加之后在执行relu，第四层卷积运算完之后不执行relu直接相加 + if strs == '2d': + self._add_to_layers(name + '/stack_%d_1' % (i+1), sharedLayers.conv2d_transpose(bottom,[kernel_size, kernel_size, \ + filters,bottom.get_shape()[-1].value ], strides=2, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False, name='/stack_%d_1' % (i+1))) + elif strs == '3d': + self._add_to_layers(name + '/stack_%d_1' % (i+1), sharedLayers.conv3d_transpose(bottom, [kernel_size, kernel_size, kernel_size, filters,\ + bottom.get_shape()[-1].value], strides=2, activation=tf.nn.leaky_relu, batch_norm=True, apply_relu=False, name='/stack_%d_1' % (i+1))) + + bottom = self._get_layer_as_input(name + '/stack_%d_1' % (i+1)) + if short_cut is not None: + if type(short_cut) is int: + short_cut = output[short_cut] + bottom = tf.add(bottom, short_cut, name='stack_%d' % (i + 1)) + if i == 2: + bottom = tf.nn.leaky_relu(bottom, name='relu') + output.append(bottom) + return output + + def output(self, outputs): + disps = [] + for i, output in enumerate(outputs): + squeeze = tf.squeeze(output, [4]) + transpose = tf.transpose(squeeze, [0, 2, 3, 1]) + upsample = tf.transpose(tf.image.resize(transpose, self.image_size_tf), [0, 3, 1, 2]) + disps.append(self.soft_arg_min(upsample, 'soft_arg_min_%d' % (i+1))) + return disps + + def soft_arg_min(self, filtered_cost_volume, name): + with tf.compat.v1.variable_scope(name): + #input.shape (batch, depth, H, W) + # softargmin to disp image, outsize of (B, H, W) + + #print('filtered_cost_volume:',filtered_cost_volume.shape) + probability_volume = tf.nn.softmax(tf.scalar_mul(-1, filtered_cost_volume), + axis=1, name='prob_volume') + #print('probability_volume:',probability_volume.shape) + volume_shape = tf.shape(probability_volume) + soft_1d = tf.cast(tf.range(0, volume_shape[1], dtype=tf.int32),tf.float32) + soft_4d = tf.tile(soft_1d, tf.stack([volume_shape[0] * volume_shape[2] * volume_shape[3]])) + soft_4d = tf.reshape(soft_4d, [volume_shape[0], volume_shape[2], volume_shape[3], volume_shape[1]]) + soft_4d = tf.transpose(soft_4d, [0, 3, 1, 2]) + estimated_disp_image = tf.reduce_sum(soft_4d * probability_volume, axis=1) + #print(estimated_disp_image.shape) + estimated_disp_image = tf.expand_dims(estimated_disp_image, axis=3) + return estimated_disp_image diff --git a/Nets/Stereo_net.py b/Nets/Stereo_net.py index f79a83a..d81038a 100644 --- a/Nets/Stereo_net.py +++ b/Nets/Stereo_net.py @@ -62,8 +62,8 @@ def _add_to_layers(self, name, op): # extract variables scope = '/'.join(op.name.split('/')[0:-1]) - variables = tf.get_collection( - tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope) + variables = tf.compat.v1.get_collection( + tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=scope) self._layer_to_var[name] = variables if not self._after_split: diff --git a/Nets/__init__.py b/Nets/__init__.py index 083d4a7..321a615 100644 --- a/Nets/__init__.py +++ b/Nets/__init__.py @@ -1,9 +1,11 @@ import Nets.DispNet import Nets.MadNet +import Nets.PSMNet STEREO_FACTORY = { Nets.DispNet.DispNet._netName: Nets.DispNet.DispNet, - Nets.MadNet.MadNet._netName: Nets.MadNet.MadNet + Nets.MadNet.MadNet._netName: Nets.MadNet.MadNet, + Nets.PSMNet.PSMNet._netName: Nets.PSMNet.PSMNet } def get_stereo_net(name,args): diff --git a/Nets/sharedLayers.py b/Nets/sharedLayers.py index 376abd5..5bcb1a8 100644 --- a/Nets/sharedLayers.py +++ b/Nets/sharedLayers.py @@ -1,4 +1,5 @@ import tensorflow as tf +import tensorflow.keras as keras import os INITIALIZER_CONV = tf.contrib.layers.xavier_initializer() @@ -29,7 +30,7 @@ def correlation(x,y,max_disp, name='corr', mode=MODE,stride=1): return correlation_native(x,y,max_disp,name=name) def correlation_native(x, y, max_disp, name='corr'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): input_shape = x.get_shape().as_list() x = tf.pad(x, [[0, 0], [0, 0], [max_disp, max_disp], [0, 0]], "CONSTANT") y = tf.pad(y, [[0, 0], [0, 0], [max_disp, max_disp], [0, 0]], "CONSTANT") @@ -39,7 +40,7 @@ def correlation_native(x, y, max_disp, name='corr'): return corr def correlation_tf(x, y, max_disp, stride=1, name='corr'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): corr_tensors = [] y_shape = tf.shape(y) y_feature = tf.pad(y,[[0,0],[0,0],[max_disp,max_disp],[0,0]]) @@ -51,20 +52,50 @@ def correlation_tf(x, y, max_disp, stride=1, name='corr'): return result -def conv2d(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, training=False): - with tf.variable_scope(name, reuse=reuse): - W = tf.get_variable(wName, kernel_shape, initializer=INITIALIZER_CONV) - b = tf.get_variable(bName, kernel_shape[3], initializer=INITIALIZER_BIAS) - x = tf.nn.conv2d(x, W, strides=[1, strides, strides, 1], padding=padding) +def conv2d(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, apply_relu=True, dilation_rate=1, training=False): + with tf.compat.v1.variable_scope(name, reuse=reuse): + W = tf.compat.v1.get_variable(wName, kernel_shape, initializer=INITIALIZER_CONV) + b = tf.compat.v1.get_variable(bName, kernel_shape[3], initializer=INITIALIZER_BIAS) + x = tf.nn.conv2d(x, W, strides=[1, strides, strides, 1], padding=padding,dilations=dilation_rate) x = tf.nn.bias_add(x, b) if batch_norm: - x = tf.layers.batch_normalization(x,training=training,momentum=0.99) - x = activation(x) + bn = keras.layers.BatchNormalization() + x = bn(x,training=training) + if apply_relu: + x = activation(x) + return x + +def conv3d(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, apply_relu=True, training=False): + with tf.compat.v1.variable_scope(name,reuse=reuse): + W = tf.compat.v1.get_variable(wName, kernel_shape, initializer=INITIALIZER_CONV) + b = tf.compat.v1.get_variable(bName, kernel_shape[4], initializer=INITIALIZER_BIAS) + x = tf.nn.conv3d(x, W, strides=[1, strides, strides, strides, 1], padding=padding) + x = tf.nn.bias_add(x, b) + if batch_norm: + bn = keras.layers.BatchNormalization() + x = bn(x,training=training) + if apply_relu: + x = activation(x) + return x + +def conv3d_transpose(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, apply_relu=True, training=False): + with tf.compat.v1.variable_scope(name, reuse=reuse): + W = tf.compat.v1.get_variable(wName, kernel_shape, initializer=INITIALIZER_CONV) + b = tf.compat.v1.get_variable(bName, kernel_shape[3], initializer=INITIALIZER_BIAS) + x = tf.compat.v1.nn.conv3d_transpose(x, filter=W, output_shape=[x.get_shape()[0].value,\ + x.get_shape()[1].value * strides,x.get_shape()[2].value * strides, x.get_shape()[3].value * strides,\ + kernel_shape[3]],strides=[1, strides, strides, strides, 1], padding=padding) + x = tf.nn.bias_add(x, b) + if batch_norm: + bn = keras.layers.BatchNormalization() + x = bn(x,training=training) + if apply_relu: + x = activation(x) return x def dilated_conv2d(x, kernel_shape, rate=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='dilated_conv', reuse=False, wName='weights', bName='biases', batch_norm=False, training=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): weights = tf.get_variable( wName, kernel_shape, initializer=INITIALIZER_CONV) biases = tf.get_variable( @@ -78,7 +109,7 @@ def dilated_conv2d(x, kernel_shape, rate=1, activation=lambda x: tf.maximum(0.1 def conv2d_transpose(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, training=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): W = tf.get_variable(wName, kernel_shape,initializer=INITIALIZER_CONV) tf.add_to_collection(tf.GraphKeys.WEIGHTS, W) b = tf.get_variable(bName, kernel_shape[2], initializer=INITIALIZER_BIAS) @@ -88,11 +119,12 @@ def conv2d_transpose(x, kernel_shape, strides=1, activation=lambda x: tf.maximum x = tf.nn.bias_add(x, b) if batch_norm: x = tf.layers.batch_normalization(x,training=training,momentum=0.99) - x = activation(x) + if apply_relu: + x = activation(x) return x def depthwise_conv(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=False, training=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): w = tf.get_variable(wName, kernel_shape, initializer=INITIALIZER_CONV) b = tf.get_variable(bName, kernel_shape[3]*kernel_shape[2], initializer=INITIALIZER_BIAS) x = tf.nn.depthwise_conv2d(x, w, strides=[1, strides, strides, 1], padding=padding) @@ -103,7 +135,7 @@ def depthwise_conv(x, kernel_shape, strides=1, activation=lambda x: tf.maximum(0 return x def separable_conv2d(x, kernel_shape, channel_multiplier=1, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=True, training=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): #detpthwise conv depthwise_conv_kernel = [kernel_shape[0],kernel_shape[1],kernel_shape[2],channel_multiplier] x = depthwise_conv(x,depthwise_conv_kernel,strides=strides,activation=lambda x: tf.maximum(0.1 * x, x),padding=padding,name='depthwise_conv',reuse=reuse,wName=wName,bName=bName,batch_norm=batch_norm, training=training) @@ -115,7 +147,7 @@ def separable_conv2d(x, kernel_shape, channel_multiplier=1, strides=1, activatio return x def grouped_conv2d(x, kernel_shape, num_groups=1, strides=1, activation=lambda x: tf.maximum(0.1 * x, x), padding='SAME', name='conv', reuse=False, wName='weights', bName='bias', batch_norm=True, training=False): - with tf.variable_scope(name,reuse=reuse): + with tf.compat.v1.variable_scope(name,reuse=reuse): w = tf.get_variable(wName,shape=kernel_shape,initializer=INITIALIZER_CONV) b = tf.get_variable(bName, kernel_shape[3], initializer=INITIALIZER_BIAS) @@ -131,7 +163,7 @@ def grouped_conv2d(x, kernel_shape, num_groups=1, strides=1, activation=lambda x return x def channel_shuffle_inside_group(x, num_groups, name='shuffle'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): _, h, w, c = x.shape.as_list() x_reshaped = tf.reshape(x, [-1, h, w, num_groups, c // num_groups]) x_transposed = tf.transpose(x_reshaped, [0, 1, 2, 4, 3])