def build_inputs(self): """Input fetching and batching Outputs: self.images: image batch of shape [batch, hz, wz, 3] labels: image batch of shape [batch, hx, wx, num_classes] """ if self.mode in ['train', 'validation', 'test']: # Put data loading and preprocessing in CPU is substantially faster # DataSet prepare with tf.device("/cpu:0"): dataset = DataLoader(self.data_config, self.train_config['DataSet'], self.train_config['class_dict']) self.images, labels = dataset.get_one_batch() self.labels = tf.one_hot(labels, self.num_classes) else: self.images_feed = tf.placeholder(shape=[None, None, None, 3], dtype=tf.uint8, name='images_input') self.images = tf.to_float(self.images_feed)/255
class BiseNet(object): def __init__(self, model_config, train_config, num_classes, mode): self.model_config = model_config self.train_config = train_config self.num_classes = num_classes self.mode = mode assert mode in ['train', 'validation', 'inference', 'test'] if self.mode == 'train': self.data_config = self.train_config['train_data_config'] elif self.mode == 'validation': self.data_config = self.train_config['validation_data_config'] elif self.mode == 'test': self.data_config = self.train_config['test_data_config'] self.images = None self.images_feed = None self.labels = None self.net = None self.sup1 = None self.sup2 = None self.init_fn = None self.loss = None self.total_loss = [] self.response = None with tf.device("/cpu:0"): self.dataset = DataLoader(self.data_config, self.train_config['DataSet'], self.train_config['dataset_dir'], self.train_config['class_dict']) def build_inputs(self): """Input fetching and batching Outputs: self.images: image batch of shape [batch, hz, wz, 3] labels: image batch of shape [batch, hx, wx, num_classes] """ if self.mode in ['train', 'validation', 'test']: # Put data loading and preprocessing in CPU is substantially faster # DataSet prepare self.images, labels = self.dataset.get_one_batch() # labels = tf.Print(labels, [tf.unique(tf.reshape(labels,[-1,]))[0]], message="labels:", summarize=10) self.labels = tf.one_hot(labels, self.num_classes) else: self.images_feed = tf.placeholder(shape=[None, None, None, 3], dtype=tf.uint8, name='images_input') self.images = tf.to_float(self.images_feed) / 255 def is_training(self): """Returns true if the model is built for training mode""" return self.mode == 'train' def setup_global_step(self): global_step = tf.Variable(initial_value=0, name='global_step', trainable=False, collections=[ tf.GraphKeys.GLOBAL_STEP, tf.GraphKeys.GLOBAL_VARIABLES ]) self.global_step = global_step def build_bisenet_custom(self, reuse=False): """ Builds the BiSeNet model. Arguments: reuse: Reuse variable or not Returns: BiSeNet model """ ### The spatial path ### The number of feature maps for each convolution is not specified in the paper ### It was chosen here to be equal to the number of feature maps of a classification ### model at each corresponding stage batch_norm_params = self.model_config['batch_norm_params'] init_method = self.model_config['conv_config']['init_method'] down_16x_end_points = self.model_config['net_node']['16xdown:50'] down_32x_end_points = self.model_config['net_node']['32xdown:25'] if init_method == 'kaiming_normal': initializer = slim.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False) else: initializer = slim.xavier_initializer() with tf.variable_scope('spatial_net', reuse=reuse): with slim.arg_scope([slim.conv2d], biases_initializer=None, weights_initializer=initializer): with slim.arg_scope([slim.batch_norm], is_training=self.is_training(), **batch_norm_params): # inference/spatial_net/Conv/Conv2D run 1 average cost 250.552994 ms, 25.405 %, FlopsRate: 9.064 % # conv2d spatial_net = slim.conv2d(self.images, 16, [3, 3], stride=[2, 2], activation_fn=None) spatial_net = hard_swish( slim.batch_norm(spatial_net, fused=True)) # bneck1 exp_size = _make_divisible(16) spatial_net = slim.conv2d(spatial_net, exp_size, [1, 1], stride=[1, 1], activation_fn=None) spatial_net = slim.batch_norm(spatial_net, fused=True) spatial_net = DepthSepConv(spatial_net, 16, kernel=[3, 3], stride=2) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) # bneck2 exp_size = _make_divisible(72) spatial_net = slim.conv2d(spatial_net, exp_size, [1, 1], stride=[1, 1], activation_fn=None) spatial_net = slim.batch_norm(spatial_net, fused=True) spatial_net = DepthSepConv(spatial_net, 24, kernel=[3, 3], stride=2) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) # bneck3 exp_size = _make_divisible(88) spatial_net = slim.conv2d(spatial_net, exp_size, [1, 1], stride=[1, 1], activation_fn=None) spatial_net = slim.batch_norm(spatial_net, fused=True) spatial_net = DepthSepConv(spatial_net, 24, kernel=[3, 3], stride=1) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) # bneck4 exp_size = _make_divisible(96) spatial_net = slim.conv2d(spatial_net, exp_size, [1, 1], stride=[1, 1], activation_fn=None) spatial_net = slim.batch_norm(spatial_net, fused=True) spatial_net = DepthSepConv(spatial_net, 40, kernel=[3, 3], stride=1) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) # bneck5 spatial_net = DepthSepConv(spatial_net, 80, kernel=[3, 3], stride=1) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) # bneck6 spatial_net = DepthSepConv(spatial_net, 128, kernel=[3, 3], stride=1) spatial_net = tf.nn.relu( slim.batch_norm(spatial_net, fused=True)) frontend_config = self.model_config['frontend_config'] ### Context path logits, end_points, frontend_scope, init_fn = frontend_builder.build_frontend( self.images, frontend_config, self.is_training(), reuse) ### Combining the paths with tf.variable_scope('combine_path', reuse=reuse): with slim.arg_scope([slim.conv2d], biases_initializer=None, weights_initializer=initializer): with slim.arg_scope([slim.batch_norm], is_training=self.is_training(), **batch_norm_params): # tail part global_context = tf.reduce_mean( end_points[down_32x_end_points], [1, 2], keep_dims=True) global_context = slim.conv2d(global_context, 128, 1, [1, 1], activation_fn=None) global_context = tf.nn.relu( slim.batch_norm(global_context, fused=True)) ARM_out1 = AttentionRefinementModule_Custom( end_points[down_32x_end_points], n_filters=128) ARM_out2 = AttentionRefinementModule_Custom( end_points[down_16x_end_points], n_filters=128) ARM_out1 = tf.add(ARM_out1, global_context) ARM_out1 = Upsampling(ARM_out1, scale=2) # inference/combine_path/Conv_6/Conv2D run 1 average cost 23.034000 ms, 2.336 %, FlopsRate: 8.879 % exp_size = _make_divisible(256) ARM_out1 = slim.conv2d(ARM_out1, exp_size, [1, 1], stride=[1, 1], activation_fn=None) ARM_out1 = slim.batch_norm(ARM_out1, fused=True) ARM_out1 = DepthSepConv(ARM_out1, 128, kernel=[3, 3], stride=1) ARM_out1 = tf.nn.relu(slim.batch_norm(ARM_out1, fused=True)) ARM_out2 = tf.add(ARM_out2, ARM_out1) ARM_out2 = Upsampling(ARM_out2, scale=2) # inference/combine_path/Conv_13/Conv2D run 1 average cost 23.034000 ms, 2.336 %, FlopsRate: 8.879 % exp_size = _make_divisible(256) ARM_out2 = slim.conv2d(ARM_out2, exp_size, [1, 1], stride=[1, 1], activation_fn=None) ARM_out2 = slim.batch_norm(ARM_out2, fused=True) ARM_out2 = DepthSepConv(ARM_out2, 128, kernel=[3, 3], stride=1) ARM_out2 = tf.nn.relu(slim.batch_norm(ARM_out2, fused=True)) context_net = ARM_out2 FFM_out = FeatureFusionModule_Custom(input_1=spatial_net, input_2=context_net, n_filters=256) ARM_out1 = ConvBlock(ARM_out1, n_filters=128, kernel_size=[3, 3]) ARM_out2 = ConvBlock(ARM_out2, n_filters=128, kernel_size=[3, 3]) exp_size = _make_divisible(128) FFM_out = slim.conv2d(FFM_out, exp_size, [1, 1], stride=[1, 1], activation_fn=None) FFM_out = slim.batch_norm(FFM_out, fused=True) FFM_out = DepthSepConv(FFM_out, 64, kernel=[3, 3], stride=1) FFM_out = tf.nn.relu(slim.batch_norm(FFM_out, fused=True)) # Upsampling + dilation or only Upsampling FFM_out = Upsampling(FFM_out, scale=2) # inference/combine_path/Conv_12/Conv2D run 1 average cost 32.151001 ms, 3.260 %, FlopsRate: 8.879 % exp_size = _make_divisible(128) FFM_out = slim.conv2d(FFM_out, exp_size, [1, 1], stride=[1, 1], activation_fn=None) FFM_out = DepthSepConv(FFM_out, 64, kernel=[3, 3], stride=1, rate=2) FFM_out = tf.nn.relu(slim.batch_norm(FFM_out, fused=True)) FFM_out = slim.conv2d(FFM_out, self.num_classes, [1, 1], activation_fn=None, scope='logits') self.net = Upsampling(FFM_out, 4) if self.mode in ['train', 'validation', 'test']: sup1 = slim.conv2d(ARM_out1, self.num_classes, [1, 1], activation_fn=None, scope='supl1') sup2 = slim.conv2d(ARM_out2, self.num_classes, [1, 1], activation_fn=None, scope='supl2') self.sup1 = Upsampling(sup1, scale=16) self.sup2 = Upsampling(sup2, scale=8) self.init_fn = init_fn def build_bisenet(self, reuse=False): """ Builds the BiSeNet model. Arguments: reuse: Reuse variable or not Returns: BiSeNet model """ ### The spatial path ### The number of feature maps for each convolution is not specified in the paper ### It was chosen here to be equal to the number of feature maps of a classification ### model at each corresponding stage batch_norm_params = self.model_config['batch_norm_params'] init_method = self.model_config['conv_config']['init_method'] down_16x_end_points = self.model_config['net_node']['16xdown:50'] down_32x_end_points = self.model_config['net_node']['32xdown:25'] if init_method == 'kaiming_normal': initializer = slim.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False) else: initializer = slim.xavier_initializer() with tf.variable_scope('spatial_net', reuse=reuse): with slim.arg_scope([slim.conv2d], biases_initializer=None, weights_initializer=initializer): with slim.arg_scope([slim.batch_norm], is_training=self.is_training(), **batch_norm_params): #print("*"*20) print(self.images) #print("*" * 20) spatial_net = ConvBlock(self.images, n_filters=64, kernel_size=[7, 7], strides=2) spatial_net = ConvBlock(spatial_net, n_filters=64, kernel_size=[3, 3], strides=2) spatial_net = ConvBlock(spatial_net, n_filters=64, kernel_size=[3, 3], strides=2) spatial_net = ConvBlock(spatial_net, n_filters=128, kernel_size=[1, 1]) frontend_config = self.model_config['frontend_config'] ### Context path logits, end_points, frontend_scope, init_fn = frontend_builder.build_frontend( self.images, frontend_config, self.is_training(), reuse) ### Combining the paths with tf.variable_scope('combine_path', reuse=reuse): with slim.arg_scope([slim.conv2d], biases_initializer=None, weights_initializer=initializer): with slim.arg_scope([slim.batch_norm], is_training=self.is_training(), **batch_norm_params): # tail part size = tf.shape(end_points[down_32x_end_points])[1:3] global_context = tf.reduce_mean( end_points[down_32x_end_points], [1, 2], keep_dims=True) global_context = slim.conv2d(global_context, 128, 1, [1, 1], activation_fn=None) global_context = tf.nn.relu( slim.batch_norm(global_context, fused=True)) net_5 = AttentionRefinementModule( end_points[down_32x_end_points], n_filters=128) net_4 = AttentionRefinementModule( end_points[down_16x_end_points], n_filters=128) net_5 = tf.add(net_5, global_context) net_5 = Upsampling(net_5, scale=2) net_5 = ConvBlock(net_5, n_filters=128, kernel_size=[3, 3]) #net_4=net_5 net_4 = tf.add(net_4, net_5) net_4 = Upsampling(net_4, scale=2) net_4 = ConvBlock(net_4, n_filters=128, kernel_size=[3, 3]) context_net = net_4 net = FeatureFusionModule(input_1=spatial_net, input_2=context_net, n_filters=256) net_5 = ConvBlock(net_5, n_filters=128, kernel_size=[3, 3]) net_4 = ConvBlock(net_4, n_filters=128, kernel_size=[3, 3]) net = ConvBlock(net, n_filters=64, kernel_size=[3, 3]) # Upsampling + dilation or only Upsampling net = Upsampling(net, scale=2) net = slim.conv2d(net, 64, [3, 3], rate=2, activation_fn=tf.nn.relu, biases_initializer=None, normalizer_fn=slim.batch_norm) net = slim.conv2d(net, self.num_classes, [1, 1], activation_fn=None, scope='logits') self.net = Upsampling(net, 4) # net = slim.conv2d(net, self.num_classes, [1, 1], activation_fn=None, scope='logits') # self.net = Upsampling(net, scale=8) if self.mode in ['train', 'validation', 'test']: sup1 = slim.conv2d(net_5, self.num_classes, [1, 1], activation_fn=None, scope='supl1') sup2 = slim.conv2d(net_4, self.num_classes, [1, 1], activation_fn=None, scope='supl2') self.sup1 = Upsampling(sup1, scale=16) self.sup2 = Upsampling(sup2, scale=8) self.init_fn = init_fn def build_loss(self): # self.labels = tf.Print(self.labels, [tf.unique(tf.reshape(self.labels, (-1,)))[0]], message="label:", summarize=10) loss1 = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(logits=self.net, labels=self.labels)) loss2 = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(logits=self.sup1, labels=self.labels)) loss3 = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(logits=self.sup2, labels=self.labels)) loss = tf.add_n([loss1, loss2, loss3]) # self.loss = loss1 # self.total_loss = tf.losses.get_total_loss() return loss def summarize(self): shape = tf.shape(self.labels) # Tensorboard inspection tf.summary.image('image', self.images, family=self.mode, max_outputs=1) # tf.Print(self.labels, [tf.shape(self.labels)], message="label size:", summarize=10) color_map = colors_dict[self.train_config['DataSet']] tf.summary.image('GT', tf.reshape( tf.matmul( tf.reshape(self.labels, [-1, self.num_classes]), color_map), [-1, shape[1], shape[2], 3]), family=self.mode, max_outputs=1) tf.summary.image('response', tf.reshape( tf.matmul( tf.reshape( tf.one_hot(tf.argmax(self.net, -1), self.num_classes), [-1, self.num_classes]), color_map), [-1, shape[1], shape[2], 3]), family=self.mode, max_outputs=1) tf.summary.scalar('total_loss', tf.reduce_mean(self.total_loss), family=self.mode) # tf.summary.scalar('loss', self.loss, family=self.mode) accuracy, accuracy_update = tf.contrib.metrics.streaming_accuracy( tf.argmax(self.net, -1), tf.argmax(self.labels, -1)) mean_IOU, mean_IOU_update = tf.contrib.metrics.streaming_mean_iou( predictions=tf.argmax(self.net, -1), labels=tf.argmax(self.labels, -1), num_classes=self.num_classes) with tf.control_dependencies([accuracy_update, mean_IOU_update]): tf.summary.scalar('accuracy', accuracy, family=self.mode) tf.summary.scalar('mean_IOU', mean_IOU, family=self.mode) def predict(self): self.response = self.net def build(self, num_gpus=1, reuse=False): """Creates all ops for training and evaluation""" with tf.name_scope(self.mode): if self.mode in ['train', 'validation', 'test']: tower_losses = [] for i in range(num_gpus): self.build_inputs() with tf.device('/gpu:%d' % i): # First tower has default name scope. name_scope = ('clone_%d' % i) if i else '' with tf.name_scope(name_scope) as scope: with tf.variable_scope( tf.get_variable_scope(), reuse=True if i != 0 else None): if self.model_config['use_custom']: self.build_bisenet_custom(reuse=reuse) else: self.build_bisenet(reuse=reuse) loss = self.build_loss() self.total_loss.append(loss) with tf.device('/cpu:0'): self.summarize() else: self.build_inputs() self.build_bisenet(reuse=reuse) self.predict() if self.is_training(): self.setup_global_step()