def tower_loss(images,
score_maps,
geo_maps,
training_masks,
reuse_variables=None):
# Build inference graph
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
# 模型定义！！！，f_score是和原图大小一样的是否是前景的概率图， f_geometry是5张图，4张是上下左右值，1张是旋转角度值
f_score, f_geometry = model.model(images, is_training=True)
# def loss(y_true_cls, y_pred_cls, y_true_geo, y_pred_geo,training_mask):
model_loss = model.loss(score_maps, f_score, geo_maps, f_geometry,
training_masks)
total_loss = tf.add_n(
[model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
tf.summary.image('input', images)
tf.summary.image('score_map', score_maps)
tf.summary.image('score_map_pred', f_score * 255)
tf.summary.image('geo_map_0', geo_maps[:, :, :, 0:1])
tf.summary.image('geo_map_#0_pred', f_geometry[:, :, :, 0:1])
tf.summary.image('geo_map_#1_pred', f_geometry[:, :, :, 0:1])
tf.summary.image('training_masks', training_masks)
tf.summary.scalar('model_loss', model_loss)
tf.summary.scalar('total_loss', total_loss)
return total_loss, model_loss, f_score, f_geometry
def tower_loss(images, annotation,class_labels,reuse_variables=None):
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
logits = model.model(images, is_training=True)
pred = tf.argmax(logits, dimension=3)
model_loss = model.loss(annotation, logits,class_labels)
total_loss = tf.add_n([model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
# add summary
if reuse_variables is None:
tf.summary.scalar('model_loss', model_loss)
tf.summary.scalar('total_loss', total_loss)
return total_loss, model_loss,pred
def tower_loss(images,
score_maps,
geo_maps,
training_masks,
reuse_variables=None):
# Build inference graph
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
f_score, f_geometry = model.model(images, is_training=True)
model_loss = model.loss(score_maps, f_score, geo_maps, f_geometry,
training_masks)
total_loss = tf.add_n(
[model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
# cls_score = f_score
cls_score = tf.nn.softmax(f_score)
geo_score_1 = tf.nn.softmax(f_geometry[:, :, :, 0:2])
# geo_score_2 = tf.nn.softmax(f_geometry[:,:,:,2:4])
# geo_score_3 = tf.nn.softmax(f_geometry[:,:,:,4:6])
# geo_score_4 = tf.nn.softmax(f_geometry[:,:,:,6:8])
# geo_score_5 = tf.nn.softmax(f_geometry[:,:,:,8:10])
# geo_score_6 = tf.nn.softmax(f_geometry[:,:,:,10:12])
# geo_score_7 = tf.nn.softmax(f_geometry[:,:,:,12:14])
# geo_score_8 = tf.nn.softmax(f_geometry[:,:,:,14:16])
# add summary
if reuse_variables is None:
tf.summary.image('input', images, max_outputs=1)
tf.summary.image('score_map', score_maps, max_outputs=1)
# tf.summary.image('score_map_pred', cls_score * 255, max_outputs=1)
tf.summary.image('score_map_pred',
cls_score[:, :, :, 1:2] * 255,
max_outputs=1)
tf.summary.image('geo_map_0', geo_maps[:, :, :, 0:1], max_outputs=1)
tf.summary.image('geo_map_1_pred',
geo_score_1[:, :, :, 1:2] * 255,
max_outputs=1)
# tf.summary.image('geo_map_2_pred', geo_score_2[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_3_pred', geo_score_3[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_4_pred', geo_score_4[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_5_pred', geo_score_5[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_6_pred', geo_score_6[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_7_pred', geo_score_7[:, :, :, 1:2] * 255, max_outputs=1)
# tf.summary.image('geo_map_8_pred', geo_score_8[:, :, :, 1:2] * 255, max_outputs=1)
tf.summary.scalar('model_loss', model_loss)
tf.summary.scalar('total_loss', total_loss)
return total_loss, model_loss
def tower_loss(images, seg_maps_gt, training_masks, reuse_variables=None):
# Build inference graph
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
seg_maps_pred = model.model(images, is_training=True)
model_loss = model.loss(seg_maps_gt, seg_maps_pred, training_masks)
total_loss = tf.add_n([model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
# add summary
if reuse_variables is None:
tf.summary.image('input', images)
tf.summary.image('seg_map_0_gt', seg_maps_gt[:, :, :, 0:1] * 255)
tf.summary.image('seg_map_0_pred', seg_maps_pred[:, :, :, 0:1] * 255)
tf.summary.image('training_masks', training_masks)
tf.summary.scalar('model_loss', model_loss)
tf.summary.scalar('total_loss', total_loss)
return total_loss, model_loss
def main(argv=None):
# 选择GPU
if FLAGS.gpu!="1" and FLAGS.gpu!="0":
logger.error("无法确定使用哪一个GPU，退出")
exit()
logger.info("使用GPU%s显卡进行训练",FLAGS.gpu)
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
logger.info(
"本次使用的参数：\nlearning_rate:%f\ndecay_steps:%f\nmax_steps:%d\nevaluate_steps:%d\nmodel:%s\nlambda1:%d\nlogs_path:%s\nrestore:%r\ndebug:%r\nsave_checkpoint_steps:%d", \
FLAGS.learning_rate,
FLAGS.decay_steps,
FLAGS.max_steps,
FLAGS.evaluate_steps,
FLAGS.model,
FLAGS.lambda1,
FLAGS.logs_path,
FLAGS.restore,
FLAGS.debug,
FLAGS.save_checkpoint_steps)
now = datetime.datetime.now()
StyleTime = now.strftime("%Y-%m-%d-%H-%M-%S")
os.makedirs(os.path.join(FLAGS.logs_path, StyleTime))
if not os.path.exists(FLAGS.model):
os.makedirs(FLAGS.model)
# 输入图像数据的维度[批次, 高度, 宽度, 3通道]
ph_input_image = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='ph_input_image')
ph_label = tf.placeholder(tf.int64, shape=[None], name='ph_label')
global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)
learning_rate = tf.Variable(FLAGS.learning_rate, trainable=False)
tf.summary.scalar('learning_rate', learning_rate)
adam_opt = tf.train.AdamOptimizer(learning_rate) # 默认是learning_rate是0.001，而且后期会不断的根据梯度调整，一般不用设这个数，所以我索性去掉了
# gpu_id = int(FLAGS.gpu)
# with tf.device('/gpu:%d' % gpu_id):
# with tf.name_scope('model_%d' % gpu_id) as scope:
cls_prob,cls_preb = model.model(ph_input_image)
cross_entropy = model.loss(cls_prob,ph_label)
batch_norm_updates_op = tf.group(*tf.get_collection(tf.GraphKeys.UPDATE_OPS))
#计算梯度
grads = adam_opt.compute_gradients(cross_entropy)
# logger.info("计算图定义完毕，定义在gpu:%d上", gpu_id)
# 使用计算得到的梯度来更新对应的variable
apply_gradient_op = adam_opt.apply_gradients(grads, global_step=global_step)
# 这个是定义召回率、精确度和F1
v_recall = tf.Variable(0.001, trainable=False)
v_precision = tf.Variable(0.001, trainable=False)
v_accuracy = tf.Variable(0.001, trainable=False)
v_f1 = tf.Variable(0.001, trainable=False)
tf.summary.scalar("Recall",v_recall)
tf.summary.scalar("Precision",v_precision)
tf.summary.scalar("F1",v_f1)
summary_op = tf.summary.merge_all()
logger.info("summary定义完毕")
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# 某些操作执行的依赖关系，这时我们可以使用tf.control_dependencies()来实现
# 我依赖于
with tf.control_dependencies([variables_averages_op, apply_gradient_op, batch_norm_updates_op]):
train_op = tf.no_op(name='train_op') # no_op啥也不干，但是它依赖的操作都会被干一遍
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
summary_writer = tf.summary.FileWriter(os.path.join(FLAGS.logs_path,StyleTime), tf.get_default_graph())
if FLAGS.pretrained_model_path is not None:
logger.info('加载vgg模型：%s',FLAGS.pretrained_model_path)
variable_restore_op = slim.assign_from_checkpoint_fn(FLAGS.pretrained_model_path,
slim.get_trainable_variables(),
ignore_missing_vars=True)
# 早停用的变量
best_f1 = 0
early_stop_counter = 0
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.95
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
if FLAGS.restore:
ckpt = tf.train.latest_checkpoint(FLAGS.model)
logger.debug("最新的模型文件:%s",ckpt) #有点担心learning rate也被恢复
saver.restore(sess, ckpt)
else:
logger.info("从头开始训练模型")
sess.run(tf.global_variables_initializer())
if FLAGS.pretrained_model_path is not None:
variable_restore_op(sess)
logger.debug("开始加载训练数据")
# 是的，get_batch返回的是一个generator
data_generator = data_provider.get_batch(num_workers=FLAGS.num_readers,label_file=FLAGS.train_label,batch_num=FLAGS.train_batch)
start = time.time()
train_start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(start))
logger.debug("开始训练")
for step in range(FLAGS.max_steps):
image_list,label_list = next(data_generator) # next(<迭代器>）来返回下一个结果
logger.debug("成功加载图片%d张，标签%d个：",len(image_list),len(label_list))
image_list = data_util.prepare4vgg(image_list)
logger.debug("开始第%d步训练，运行sess.run,数据shape：%r",step,image_list.shape)
_, summary_str,classes = sess.run([train_op, summary_op, cls_prob],
feed_dict = {ph_input_image: image_list , ph_label: label_list}) # data[3]是图像的路径，传入sess是为了调试画图用 np.array(image_list)
logger.info("结束第%d步训练，结束sess.run",step)
summary_writer.add_summary(summary_str, global_step=step)
if step!=0 and step % FLAGS.evaluate_steps == 0:
logger.info("在第%d步，开始进行模型评估",step)
# data[4]是大框的坐标，是个数组，8个值
accuracy_value,precision_value,recall_value,f1_value = validate(sess,cls_preb,ph_input_image,ph_label)
if f1_value>best_f1:
logger.info("新F1值[%f]大于过去最好的F1值[%f]，早停计数器重置",f1_value,best_f1)
best_f1 = f1_value
early_stop_counter = 0
# 每次效果好的话，就保存一个模型
filename = ('ctpn-{:s}-{:d}'.format(train_start_time,step + 1) + '.ckpt')
filename = os.path.join(FLAGS.model, filename)
saver.save(sess, filename)
logger.info("在第%d步，保存了最好的模型文件：%s，F1：%f",step,filename,best_f1)
else:
logger.info("新F1值[%f]小于过去最好的F1值[%f]，早停计数器+1", f1_value, best_f1)
early_stop_counter+= 1
# 更新F1,Recall和Precision
sess.run([tf.assign(v_f1, f1_value),
tf.assign(v_recall, recall_value),
tf.assign(v_precision,precision_value),
tf.assign(v_accuracy, accuracy_value)])
logger.info("在第%d步，模型评估结束", step)
if early_stop_counter> FLAGS.early_stop:
logger.warning("达到了早停计数次数：%d次，训练提前结束",early_stop_counter)
break
if step != 0 and step % FLAGS.decay_steps == 0:
logger.info("学习率(learning rate)衰减：%f=>%f",learning_rate.eval(),learning_rate.eval() * FLAGS.decay_rate)
sess.run(tf.assign(learning_rate, learning_rate.eval() * FLAGS.decay_rate))
