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