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')
