def build_model(): """ """ graph = tf.Graph() with graph.as_default(): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 512, 1024, 3], name='input_tensor') label_tensor = tf.placeholder(dtype=tf.int32, shape=[1, 512, 1024], name='input_label') net = bisenet_v2.BiseNetV2(phase='train') loss = net.compute_loss(input_tensor=input_tensor, label_tensor=label_tensor, name='BiseNetV2', reuse=False) final_output_tensor = graph.get_tensor_by_name( 'BiseNetV2/logits/segmentation_head_logits:0') stage_1_output_tensor = graph.get_tensor_by_name( 'BiseNetV2/detail_branch/stage_1/conv_block_2_repeat_1/3x3_conv/relu:0' ) stage_2_output_tensor = graph.get_tensor_by_name( 'BiseNetV2/detail_branch/stage_2/conv_block_2_repeat_2/3x3_conv/relu:0' ) stage_3_output_tensor = graph.get_tensor_by_name( 'BiseNetV2/detail_branch/stage_3/conv_block_2_repeat_1/3x3_conv/conv/conv:0' ) y_c = tf.reduce_sum(tf.multiply(final_output_tensor, label_tensor), axis=1) return graph
def load_graph_from_ckpt_file(weights_path): """ :param weights_path: return: """ # construct compute graph input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 256, 3], name='input_tensor') net = bisenet_v2.BiseNetV2(phase='test', cfg=CFG) prediction = net.inference(input_tensor=input_tensor, name='BiseNetV2', reuse=False) prediction = tf.squeeze(prediction, axis=0, name='final_output') prediction = tf.identity(prediction, name='final_output') sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = 0.9 sess_config.gpu_options.allow_growth = True sess_config.gpu_options.allocator_type = 'BFC' # define moving average version of the learned variables for eval with tf.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage(0.9995) variables_to_restore = variable_averages.variables_to_restore() saver = tf.train.Saver(variables_to_restore) # create a session sess = tf.Session(config=sess_config) # import best model saver.restore(sess, weights_path) # variables # get graph definition gd = graph_util.remove_training_nodes(sess.graph_def) return gd, sess, prediction
def test_bisenet_celebamaskhq(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ # define bisenet input_tensor_size = CFG.AUG.EVAL_CROP_SIZE input_tensor_size = [int(tmp) for tmp in input_tensor_size] input_tensor = tf.placeholder( dtype=tf.float32, shape=[1, input_tensor_size[1], input_tensor_size[0], 3], name='input_tensor') bisenet_model = bisenet_v2.BiseNetV2(phase='test', cfg=CFG) prediction = bisenet_model.inference(input_tensor=input_tensor, name='BiseNetV2', reuse=False) # define session and gpu config sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) # define moving average version of the learned variables for eval with tf.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage( CFG.SOLVER.MOVING_AVE_DECAY) variables_to_restore = variable_averages.variables_to_restore() # define saver saver = tf.train.Saver(variables_to_restore) # prepare input image src_image = cv2.imread(image_path, cv2.IMREAD_COLOR) preprocessed_image = preprocess_image(src_image, input_tensor_size) # run net and decode output prediction with sess.as_default(): saver.restore(sess, weights_path) t_start = time.time() loop_times = 2000 for i in range(loop_times): prediction_value = sess.run( fetches=prediction, feed_dict={input_tensor: [preprocessed_image]}) t_cost = (time.time() - t_start) / loop_times print('Mean cost time: {:.5f}s'.format(t_cost)) print('Mean fps: {:.5f}fps'.format(1.0 / t_cost)) prediction_value = np.squeeze(prediction_value, axis=0) prediction_value = cv2.resize(prediction_value, dsize=(input_tensor_size[0] * 2, input_tensor_size[1] * 2), interpolation=cv2.INTER_NEAREST) print('Prediction mask unique label ids: {}'.format( np.unique(prediction_value))) prediction_mask_color = decode_prediction_mask(prediction_value) plt.figure('src_image') plt.imshow(src_image[:, :, (2, 1, 0)]) plt.figure('prediction_mask_color') plt.imshow(prediction_mask_color[:, :, (2, 1, 0)]) plt.show()
def __init__(self): """ initialize bisenetv2 trainner """ # define solver params and dataset self._mixed_8K_io = mixed_8K_tf_io.mixed_8K_TfIO() self._train_dataset = self._mixed_8K_io.train_dataset_reader self._steps_per_epoch = len(self._train_dataset) self._model_name = CFG.MODEL.MODEL_NAME self._train_epoch_nums = CFG.TRAIN.EPOCH_NUMS self._batch_size = CFG.TRAIN.BATCH_SIZE self._snapshot_epoch = CFG.TRAIN.SNAPSHOT_EPOCH self._model_save_dir = ops.join(CFG.TRAIN.MODEL_SAVE_DIR, self._model_name) self._tboard_save_dir = ops.join(CFG.TRAIN.TBOARD_SAVE_DIR, self._model_name) self._enable_miou = CFG.TRAIN.COMPUTE_MIOU.ENABLE if self._enable_miou: self._record_miou_epoch = CFG.TRAIN.COMPUTE_MIOU.EPOCH self._input_tensor_size = [int(tmp) for tmp in CFG.AUG.TRAIN_CROP_SIZE] self._init_learning_rate = CFG.SOLVER.LR self._moving_ave_decay = CFG.SOLVER.MOVING_AVE_DECAY self._momentum = CFG.SOLVER.MOMENTUM self._lr_polynimal_decay_power = CFG.SOLVER.LR_POLYNOMIAL_POWER self._optimizer_mode = CFG.SOLVER.OPTIMIZER.lower() if CFG.TRAIN.RESTORE_FROM_SNAPSHOT.ENABLE: self._initial_weight = CFG.TRAIN.RESTORE_FROM_SNAPSHOT.SNAPSHOT_PATH else: self._initial_weight = None if CFG.TRAIN.WARM_UP.ENABLE: self._warmup_epoches = CFG.TRAIN.WARM_UP.EPOCH_NUMS self._warmup_init_learning_rate = self._init_learning_rate / 1000.0 else: self._warmup_epoches = 0 # define tensorflow session sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self._sess = tf.Session(config=sess_config) # define graph input tensor with tf.variable_scope(name_or_scope='graph_input_node'): self._input_src_image, self._input_label_image = self._train_dataset.next_batch( batch_size=self._batch_size ) # define model loss self._model = bisenet_v2.BiseNetV2(phase='train', cfg=CFG) loss_set = self._model.compute_loss( input_tensor=self._input_src_image, label_tensor=self._input_label_image, name='BiseNetV2', reuse=False ) self._prediciton = self._model.inference( input_tensor=self._input_src_image, name='BiseNetV2', reuse=True ) self._loss = loss_set['total_loss'] self._l2_loss = loss_set['l2_loss'] # define miou if self._enable_miou: with tf.variable_scope('miou'): pred = tf.reshape(self._prediciton, [-1, ]) gt = tf.reshape(self._input_label_image, [-1, ]) indices = tf.squeeze(tf.where(tf.less_equal(gt, CFG.DATASET.NUM_CLASSES - 1)), 1) gt = tf.gather(gt, indices) pred = tf.gather(pred, indices) self._miou, self._miou_update_op = tf.metrics.mean_iou( labels=gt, predictions=pred, num_classes=CFG.DATASET.NUM_CLASSES ) # define learning rate with tf.variable_scope('learning_rate'): self._global_step = tf.Variable(1.0, dtype=tf.float32, trainable=False, name='global_step') warmup_steps = tf.constant( self._warmup_epoches * self._steps_per_epoch, dtype=tf.float32, name='warmup_steps' ) train_steps = tf.constant( self._train_epoch_nums * self._steps_per_epoch, dtype=tf.float32, name='train_steps' ) self._learn_rate = tf.cond( pred=self._global_step < warmup_steps, true_fn=lambda: self._compute_warmup_lr(warmup_steps=warmup_steps, name='warmup_lr'), false_fn=lambda: tf.train.polynomial_decay( learning_rate=self._init_learning_rate, global_step=self._global_step, decay_steps=train_steps, end_learning_rate=0.000001, power=self._lr_polynimal_decay_power) ) self._learn_rate = tf.identity(self._learn_rate, 'lr') global_step_update = tf.assign_add(self._global_step, 1.0) # define moving average op with tf.variable_scope(name_or_scope='moving_avg'): if CFG.TRAIN.FREEZE_BN.ENABLE: train_var_list = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] else: train_var_list = tf.trainable_variables() moving_ave_op = tf.train.ExponentialMovingAverage( self._moving_ave_decay).apply(train_var_list + tf.moving_average_variables()) # define training op with tf.variable_scope(name_or_scope='train_step'): if CFG.TRAIN.FREEZE_BN.ENABLE: train_var_list = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] else: train_var_list = tf.trainable_variables() if self._optimizer_mode == 'sgd': optimizer = tf.train.MomentumOptimizer( learning_rate=self._learn_rate, momentum=self._momentum ) elif self._optimizer_mode == 'adam': optimizer = tf.train.AdamOptimizer( learning_rate=self._learn_rate, ) else: raise ValueError('Not support optimizer: {:s}'.format(self._optimizer_mode)) optimize_op = optimizer.minimize(self._loss, var_list=train_var_list) with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): with tf.control_dependencies([optimize_op, global_step_update]): with tf.control_dependencies([moving_ave_op]): self._train_op = tf.no_op() # define saver and loader with tf.variable_scope('loader_and_saver'): self._net_var = [vv for vv in tf.global_variables() if 'lr' not in vv.name] self._loader = tf.train.Saver(self._net_var) self._saver = tf.train.Saver(tf.global_variables(), max_to_keep=5) # define summary with tf.variable_scope('summary'): summary_merge_list = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total", self._loss), tf.summary.scalar('l2_loss', self._l2_loss) ] if self._enable_miou: with tf.control_dependencies([self._miou_update_op]): summary_merge_list_with_miou = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total", self._loss), tf.summary.scalar('l2_loss', self._l2_loss), tf.summary.scalar('miou', self._miou) ] self._write_summary_op_with_miou = tf.summary.merge(summary_merge_list_with_miou) if ops.exists(self._tboard_save_dir): shutil.rmtree(self._tboard_save_dir) os.makedirs(self._tboard_save_dir, exist_ok=True) model_params_file_save_path = ops.join(self._tboard_save_dir, CFG.TRAIN.MODEL_PARAMS_CONFIG_FILE_NAME) with open(model_params_file_save_path, 'w', encoding='utf-8') as f_obj: CFG.dump_to_json_file(f_obj) self._write_summary_op = tf.summary.merge(summary_merge_list) self._summary_writer = tf.summary.FileWriter(self._tboard_save_dir, graph=self._sess.graph) LOG.info('Initialize mixed_8K bisenetv2 trainner complete')
def __init__(self): """ initialize bisenetv2 multi gpu trainner """ # define solver params and dataset self._cityscapes_io = cityscapes_tf_io.CityScapesTfIO() self._train_dataset = self._cityscapes_io.train_dataset_reader self._steps_per_epoch = len(self._train_dataset) self._model_name = CFG.MODEL.MODEL_NAME self._train_epoch_nums = CFG.TRAIN.EPOCH_NUMS self._batch_size = CFG.TRAIN.BATCH_SIZE self._snapshot_epoch = CFG.TRAIN.SNAPSHOT_EPOCH self._model_save_dir = ops.join(CFG.TRAIN.MODEL_SAVE_DIR, self._model_name) self._tboard_save_dir = ops.join(CFG.TRAIN.TBOARD_SAVE_DIR, self._model_name) self._enable_miou = CFG.TRAIN.COMPUTE_MIOU.ENABLE if self._enable_miou: self._record_miou_epoch = CFG.TRAIN.COMPUTE_MIOU.EPOCH self._input_tensor_size = [ int(tmp / 2) for tmp in CFG.AUG.TRAIN_CROP_SIZE ] self._gpu_devices = CFG.TRAIN.MULTI_GPU.GPU_DEVICES self._gpu_nums = len(self._gpu_devices) self._chief_gpu_index = CFG.TRAIN.MULTI_GPU.CHIEF_DEVICE_INDEX self._batch_size_per_gpu = int(self._batch_size / self._gpu_nums) self._init_learning_rate = CFG.SOLVER.LR self._moving_ave_decay = CFG.SOLVER.MOVING_AVE_DECAY self._momentum = CFG.SOLVER.MOMENTUM self._lr_polynimal_decay_power = CFG.SOLVER.LR_POLYNOMIAL_POWER self._optimizer_mode = CFG.SOLVER.OPTIMIZER.lower() if CFG.TRAIN.RESTORE_FROM_SNAPSHOT.ENABLE: self._initial_weight = CFG.TRAIN.RESTORE_FROM_SNAPSHOT.SNAPSHOT_PATH else: self._initial_weight = None if CFG.TRAIN.WARM_UP.ENABLE: self._warmup_epoches = CFG.TRAIN.WARM_UP.EPOCH_NUMS self._warmup_init_learning_rate = self._init_learning_rate / 1000.0 else: self._warmup_epoches = 0 # define tensorflow session sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self._sess = tf.Session(config=sess_config) # define graph input tensor with tf.variable_scope(name_or_scope='graph_input_node'): self._input_src_image_list = [] self._input_label_image_list = [] for i in range(self._gpu_nums): src_imgs, label_imgs = self._train_dataset.next_batch( batch_size=self._batch_size_per_gpu) self._input_src_image_list.append(src_imgs) self._input_label_image_list.append(label_imgs) # define model self._model = bisenet_v2.BiseNetV2(phase='train') # define average container tower_grads = [] tower_total_loss = [] tower_l2_loss = [] batchnorm_updates = None # define learning rate with tf.variable_scope('learning_rate'): self._global_step = tf.Variable(1.0, dtype=tf.float32, trainable=False, name='global_step') warmup_steps = tf.constant(self._warmup_epoches * self._steps_per_epoch, dtype=tf.float32, name='warmup_steps') train_steps = tf.constant(self._train_epoch_nums * self._steps_per_epoch, dtype=tf.float32, name='train_steps') self._learn_rate = tf.cond( pred=self._global_step < warmup_steps, true_fn=lambda: self._compute_warmup_lr( warmup_steps=warmup_steps, name='warmup_lr'), false_fn=lambda: tf.train.polynomial_decay( learning_rate=self._init_learning_rate, global_step=self._global_step, decay_steps=train_steps, end_learning_rate=0.000000001, power=self._lr_polynimal_decay_power)) self._learn_rate = tf.identity(self._learn_rate, 'lr') # define optimizer if self._optimizer_mode == 'sgd': optimizer = tf.train.MomentumOptimizer( learning_rate=self._learn_rate, momentum=self._momentum) elif self._optimizer_mode == 'adam': optimizer = tf.train.AdamOptimizer( learning_rate=self._learn_rate, ) else: raise NotImplementedError( 'Not support optimizer: {:s} for now'.format( self._optimizer_mode)) # define distributed train op with tf.variable_scope(tf.get_variable_scope()): is_network_initialized = False for i in range(self._gpu_nums): with tf.device('/gpu:{:d}'.format(i)): with tf.name_scope('tower_{:d}'.format(i)) as _: input_images = self._input_src_image_list[i] input_labels = self._input_label_image_list[i] tmp_loss, tmp_grads = self._compute_net_gradients( input_images, input_labels, optimizer, is_net_first_initialized=is_network_initialized) is_network_initialized = True # Only use the mean and var in the chief gpu tower to update the parameter if i == self._chief_gpu_index: batchnorm_updates = tf.get_collection( tf.GraphKeys.UPDATE_OPS) tower_grads.append(tmp_grads) tower_total_loss.append(tmp_loss['total_loss']) tower_l2_loss.append(tmp_loss['l2_loss']) grads = self._average_gradients(tower_grads) self._loss = tf.reduce_mean(tower_total_loss, name='reduce_mean_tower_total_loss') self._l2_loss = tf.reduce_mean(tower_l2_loss, name='reduce_mean_tower_l2_loss') # define moving average op with tf.variable_scope(name_or_scope='moving_avg'): if CFG.TRAIN.FREEZE_BN.ENABLE: train_var_list = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] else: train_var_list = tf.trainable_variables() moving_ave_op = tf.train.ExponentialMovingAverage( self._moving_ave_decay).apply(train_var_list + tf.moving_average_variables()) # group all the op needed for training batchnorm_updates_op = tf.group(*batchnorm_updates) apply_gradient_op = optimizer.apply_gradients( grads, global_step=self._global_step) self._train_op = tf.group(apply_gradient_op, moving_ave_op, batchnorm_updates_op) # define prediction self._prediciton = self._model.inference( input_tensor=self._input_src_image_list[self._chief_gpu_index], name='BiseNetV2', reuse=True) # define miou if self._enable_miou: with tf.variable_scope('miou'): pred = tf.reshape(self._prediciton, [ -1, ]) gt = tf.reshape( self._input_label_image_list[self._chief_gpu_index], [ -1, ]) indices = tf.squeeze( tf.where(tf.less_equal(gt, CFG.DATASET.NUM_CLASSES - 1)), 1) gt = tf.gather(gt, indices) pred = tf.gather(pred, indices) self._miou, self._miou_update_op = tf.metrics.mean_iou( labels=gt, predictions=pred, num_classes=CFG.DATASET.NUM_CLASSES) # define saver and loader with tf.variable_scope('loader_and_saver'): self._net_var = [ vv for vv in tf.global_variables() if 'lr' not in vv.name ] self._loader = tf.train.Saver(self._net_var) self._saver = tf.train.Saver(max_to_keep=5) # define summary with tf.variable_scope('summary'): summary_merge_list = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total", self._loss), tf.summary.scalar('l2_loss', self._l2_loss) ] if self._enable_miou: with tf.control_dependencies([self._miou_update_op]): summary_merge_list_with_miou = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total", self._loss), tf.summary.scalar('l2_loss', self._l2_loss), tf.summary.scalar('miou', self._miou) ] self._write_summary_op_with_miou = tf.summary.merge( summary_merge_list_with_miou) if ops.exists(self._tboard_save_dir): shutil.rmtree(self._tboard_save_dir) os.makedirs(self._tboard_save_dir, exist_ok=True) model_params_file_save_path = ops.join( self._tboard_save_dir, CFG.TRAIN.MODEL_PARAMS_CONFIG_FILE_NAME) with open(model_params_file_save_path, 'w', encoding='utf-8') as f_obj: CFG.dump_to_json_file(f_obj) self._write_summary_op = tf.summary.merge(summary_merge_list) self._summary_writer = tf.summary.FileWriter( self._tboard_save_dir, graph=self._sess.graph) LOG.info('Initialize cityscapes bisenetv2 multi gpu trainner complete')