def pred(sess,classes,input_images,image_list):#,input_image,input_im_info,bbox_pred, cls_pred, cls_prob):
logger.info("开始探测图片")
start = time.time()
image_list = data_util.prepare4vgg(image_list)
_classes = sess.run(classes,feed_dict={input_images: image_list})
logger.info("探测图片完成，耗时: %f", (time.time() - start))
return _classes
def validate(sess,cls_pred,ph_input_image,ph_label):
#### 加载验证数据,随机加载FLAGS.validate_batch张
accuracy = 0
precision = 0
recall = 0
f1 = 0
for step in range(FLAGS.validate_times):
image_list, image_label = data_provider.load_validate_data(FLAGS.validate_label,FLAGS.validate_batch)
logger.debug("加载了验证集%d张",len(image_list))
classes = sess.run(cls_pred,feed_dict={
ph_input_image: data_util.prepare4vgg(image_list),
ph_label: image_label
}) # data[3]是图像的路径，传入sess是为了调试画图用
logger.debug("预测结果为：%r",classes)
logger.debug("Label为：%r",image_label)
# pred和label格式如:[2,1,0,1,1,3]，0-3是对应的方向，0朝上，1朝右倒，2倒立，3朝左倒
# accuracy: (tp + tn) / (p + n)
accuracy = accuracy + accuracy_score(image_label, classes)
# precision tp / (tp + fp)
precision = precision + precision_score(image_label, classes,labels=[0,1,2,3],average='micro')
# recall: tp / (tp + fn)
recall = recall + recall_score(image_label, classes,labels=[0,1,2,3],average='micro')
# f1: 2 tp / (2 tp + fp + fn)
f1 = f1 + f1_score(image_label, classes,labels=[0,1,2,3],average='micro')
accuracy = accuracy/FLAGS.validate_times
precision = precision/FLAGS.validate_times
recall = recall/FLAGS.validate_times
f1 = f1/FLAGS.validate_times
return accuracy,precision,recall,f1
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))