3 from __future__ import absolute_import
4 from __future__ import print_function
5 from __future__ import division
7 import tensorflow as tf
8 import tensorflow.contrib.slim as slim
12 sys.path.append("../")
14 from libs.configs import cfgs
15 # from libs.networks import build_whole_network2
16 from libs.networks import build_whole_network
17 from data.io.read_tfrecord import next_batch
18 from libs.box_utils import show_box_in_tensor
19 from help_utils import tools
21 os.environ["CUDA_VISIBLE_DEVICES"] = "2"
26 faster_rcnn = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
29 with tf.name_scope('get_batch'):
30 img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \
31 next_batch(dataset_name=cfgs.DATASET_NAME, # 'pascal', 'coco'
32 batch_size=cfgs.BATCH_SIZE,
33 shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
35 gtboxes_and_label = tf.reshape(gtboxes_and_label_batch, [-1, 5])
37 biases_regularizer = tf.no_regularizer
38 weights_regularizer = tf.contrib.layers.l2_regularizer(cfgs.WEIGHT_DECAY)
40 # list as many types of layers as possible, even if they are not used now
41 with slim.arg_scope([slim.conv2d, slim.conv2d_in_plane, \
42 slim.conv2d_transpose, slim.separable_conv2d, slim.fully_connected],
43 weights_regularizer=weights_regularizer,
44 biases_regularizer=biases_regularizer,
45 biases_initializer=tf.constant_initializer(0.0)):
46 final_bbox, final_scores, final_category, loss_dict = faster_rcnn.build_whole_detection_network(
47 input_img_batch=img_batch,
48 gtboxes_batch=gtboxes_and_label)
50 # ----------------------------------------------------------------------------------------------------build loss
51 weight_decay_loss = tf.add_n(slim.losses.get_regularization_losses())
52 rpn_location_loss = loss_dict['rpn_loc_loss']
53 rpn_cls_loss = loss_dict['rpn_cls_loss']
54 rpn_total_loss = rpn_location_loss + rpn_cls_loss
56 fastrcnn_cls_loss = loss_dict['fastrcnn_cls_loss']
57 fastrcnn_loc_loss = loss_dict['fastrcnn_loc_loss']
58 fastrcnn_total_loss = fastrcnn_cls_loss + fastrcnn_loc_loss
60 total_loss = rpn_total_loss + fastrcnn_total_loss + weight_decay_loss
61 # ____________________________________________________________________________________________________build loss
64 # ---------------------------------------------------------------------------------------------------add summary
66 tf.summary.scalar('RPN_LOSS/cls_loss', rpn_cls_loss)
67 tf.summary.scalar('RPN_LOSS/location_loss', rpn_location_loss)
68 tf.summary.scalar('RPN_LOSS/rpn_total_loss', rpn_total_loss)
70 tf.summary.scalar('FAST_LOSS/fastrcnn_cls_loss', fastrcnn_cls_loss)
71 tf.summary.scalar('FAST_LOSS/fastrcnn_location_loss', fastrcnn_loc_loss)
72 tf.summary.scalar('FAST_LOSS/fastrcnn_total_loss', fastrcnn_total_loss)
74 tf.summary.scalar('LOSS/total_loss', total_loss)
75 tf.summary.scalar('LOSS/regular_weights', weight_decay_loss)
77 gtboxes_in_img = show_box_in_tensor.draw_boxes_with_categories(img_batch=img_batch,
78 boxes=gtboxes_and_label[:, :-1],
79 labels=gtboxes_and_label[:, -1])
80 if cfgs.ADD_BOX_IN_TENSORBOARD:
81 detections_in_img = show_box_in_tensor.draw_boxes_with_categories_and_scores(img_batch=img_batch,
83 labels=final_category,
85 tf.summary.image('Compare/final_detection', detections_in_img)
86 tf.summary.image('Compare/gtboxes', gtboxes_in_img)
88 # ___________________________________________________________________________________________________add summary
90 global_step = slim.get_or_create_global_step()
91 lr = tf.train.piecewise_constant(global_step,
92 boundaries=[np.int64(cfgs.DECAY_STEP[0]), np.int64(cfgs.DECAY_STEP[1])],
93 values=[cfgs.LR, cfgs.LR / 10., cfgs.LR / 100.])
94 tf.summary.scalar('lr', lr)
95 optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM)
96 # optimizer = tf.train.AdamOptimizer(lr)
98 # ---------------------------------------------------------------------------------------------compute gradients
99 gradients = faster_rcnn.get_gradients(optimizer, total_loss)
101 # enlarge_gradients for bias
102 if cfgs.MUTILPY_BIAS_GRADIENT:
103 gradients = faster_rcnn.enlarge_gradients_for_bias(gradients)
105 if cfgs.GRADIENT_CLIPPING_BY_NORM:
106 with tf.name_scope('clip_gradients_YJR'):
107 gradients = slim.learning.clip_gradient_norms(gradients,
108 cfgs.GRADIENT_CLIPPING_BY_NORM)
109 # _____________________________________________________________________________________________compute gradients
114 train_op = optimizer.apply_gradients(grads_and_vars=gradients,
115 global_step=global_step)
116 summary_op = tf.summary.merge_all()
118 tf.global_variables_initializer(),
119 tf.local_variables_initializer()
122 restorer, restore_ckpt = faster_rcnn.get_restorer()
123 saver = tf.train.Saver(max_to_keep=30)
125 config = tf.ConfigProto()
126 config.gpu_options.allow_growth = True
130 with tf.Session(config=config) as sess:
132 if not restorer is None:
133 restorer.restore(sess, restore_ckpt)
134 print('restore model')
135 coord = tf.train.Coordinator()
136 threads = tf.train.start_queue_runners(sess, coord)
138 summary_path = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION)
139 tools.mkdir(summary_path)
140 summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph)
141 training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
143 for step in range(cfgs.MAX_ITERATION):
144 training_time1 = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
145 if step % cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % cfgs.SMRY_ITER != 0:
146 _, global_stepnp = sess.run([train_op, global_step])
149 if step % cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % cfgs.SMRY_ITER != 0:
152 _, global_stepnp, img_name, rpnLocLoss, rpnClsLoss, rpnTotalLoss, \
153 fastrcnnLocLoss, fastrcnnClsLoss, fastrcnnTotalLoss, totalLoss = \
155 [train_op, global_step, img_name_batch, rpn_location_loss, rpn_cls_loss, rpn_total_loss,
156 fastrcnn_loc_loss, fastrcnn_cls_loss, fastrcnn_total_loss, total_loss])
159 compute_time = compute_time + (end - start)
160 compute_imgnum = compute_imgnum + 1
161 print(""" {}: step{} image_name:{} |\t
162 rpn_loc_loss:{} |\t rpn_cla_loss:{} |\t rpn_total_loss:{} |
163 fast_rcnn_loc_loss:{} |\t fast_rcnn_cla_loss:{} |\t fast_rcnn_total_loss:{} |
164 total_loss:{} |\t per_cost_time:{}s""" \
165 .format(training_time1, global_stepnp, str(img_name[0]), rpnLocLoss, rpnClsLoss,
166 rpnTotalLoss, fastrcnnLocLoss, fastrcnnClsLoss, fastrcnnTotalLoss, totalLoss,
169 if step % cfgs.SMRY_ITER == 0:
170 _, global_stepnp, summary_str = sess.run([train_op, global_step, summary_op])
171 summary_writer.add_summary(summary_str, global_stepnp)
172 summary_writer.flush()
174 if (step > 0 and step % cfgs.SAVE_WEIGHTS_INTE == 0) or (step == cfgs.MAX_ITERATION - 1):
176 save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION)
177 if not os.path.exists(save_dir):
180 #save_ckpt = os.path.join(save_dir, 'voc_' + str(global_stepnp) + 'model.ckpt')
181 save_ckpt = os.path.join(save_dir, 'pcb_' + str(global_stepnp) + 'model.ckpt')
182 saver.save(sess, save_ckpt)
183 print(' weights had been saved')
184 print('average_training_time_per_image is' + str(compute_time / compute_imgnum))
185 print('traning start time is ' + training_time)
186 end_training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
187 print('traning end time is ' + end_training_time)
192 if __name__ == '__main__':