· 8 years ago · Feb 28, 2018, 09:42 PM
1from __future__ import absolute_import
2from __future__ import division
3from __future__ import print_function
4
5import sys
6import os
7import glob
8
9import numpy as np
10import tensorflow as tf
11import time
12
13from tensorflow.contrib import slim
14from tensorflow.python.ops import control_flow_ops
15
16sys.path.append('./datasets')
17import ressep_bg_seg
18from mask_rcnn_tfrecords import get_dataset, batch_segmentation_masks_fast,\
19 visualize_masks
20from mask_rcnn_stream import MaskRCNNMultiStream
21
22sys.path.append('./tf_models/research/slim/deployment')
23import model_deploy as model_deploy
24
25tf.app.flags.DEFINE_string(
26 'master', '', 'The address of the TensorFlow master to use.')
27
28tf.app.flags.DEFINE_string(
29 'train_dir', '/tmp/tfmodel/',
30 'Directory where checkpoints and event logs are written to.')
31
32tf.app.flags.DEFINE_integer('num_clones', 1,
33 'Number of model clones to deploy.')
34
35tf.app.flags.DEFINE_boolean('clone_on_cpu', False,
36 'Use CPUs to deploy clones.')
37
38tf.app.flags.DEFINE_integer('worker_replicas', 1, 'Number of worker replicas.')
39
40tf.app.flags.DEFINE_integer(
41 'num_ps_tasks', 0,
42 'The number of parameter servers. If the value is 0, then the parameters '
43 'are handled locally by the worker.')
44
45tf.app.flags.DEFINE_integer(
46 'log_every_n_steps', 10,
47 'The frequency with which logs are print.')
48
49tf.app.flags.DEFINE_integer(
50 'save_summaries_secs', 120,
51 'The frequency with which summaries are saved, in seconds.')
52
53tf.app.flags.DEFINE_integer(
54 'save_interval_secs', 300,
55 'The frequency with which the model is saved, in seconds.')
56
57tf.app.flags.DEFINE_integer('startup_delay_steps', 15,
58 'Number of training steps between replicas startup.')
59
60tf.app.flags.DEFINE_integer(
61 'task', 0, 'Task id of the replica running the training.')
62
63######################
64# Optimization Flags #
65######################
66
67tf.app.flags.DEFINE_float(
68 'weight_decay', 0.00004, 'The weight decay on the model weights.')
69
70tf.app.flags.DEFINE_string(
71 'optimizer', 'momentum',
72 'The name of the optimizer, one of "adadelta", "adagrad", "adam",'
73 '"ftrl", "momentum", "sgd" or "rmsprop".')
74
75tf.app.flags.DEFINE_float(
76 'adadelta_rho', 0.95,
77 'The decay rate for adadelta.')
78
79tf.app.flags.DEFINE_float(
80 'adagrad_initial_accumulator_value', 0.1,
81 'Starting value for the AdaGrad accumulators.')
82
83tf.app.flags.DEFINE_float(
84 'adam_beta1', 0.9,
85 'The exponential decay rate for the 1st moment estimates.')
86
87tf.app.flags.DEFINE_float(
88 'adam_beta2', 0.999,
89 'The exponential decay rate for the 2nd moment estimates.')
90
91tf.app.flags.DEFINE_float('opt_epsilon', 1.0, 'Epsilon term for the optimizer.')
92
93tf.app.flags.DEFINE_float('ftrl_learning_rate_power', -0.5,
94 'The learning rate power.')
95
96tf.app.flags.DEFINE_float(
97 'ftrl_initial_accumulator_value', 0.1,
98 'Starting value for the FTRL accumulators.')
99
100tf.app.flags.DEFINE_float(
101 'ftrl_l1', 0.0, 'The FTRL l1 regularization strength.')
102
103tf.app.flags.DEFINE_float(
104 'ftrl_l2', 0.0, 'The FTRL l2 regularization strength.')
105
106tf.app.flags.DEFINE_float(
107 'momentum', 0.9,
108 'The momentum for the MomentumOptimizer and RMSPropOptimizer.')
109
110tf.app.flags.DEFINE_float('rmsprop_decay', 0.9, 'Decay term for RMSProp.')
111
112#######################
113# Learning Rate Flags #
114#######################
115
116tf.app.flags.DEFINE_string(
117 'learning_rate_decay_type',
118 'exponential',
119 'Specifies how the learning rate is decayed. One of "fixed", "exponential",'
120 ' or "polynomial"')
121
122tf.app.flags.DEFINE_float('learning_rate', 0.005, 'Initial learning rate.')
123
124tf.app.flags.DEFINE_float(
125 'end_learning_rate', 0.00001,
126 'The minimal end learning rate used by a polynomial decay learning rate.')
127
128tf.app.flags.DEFINE_float(
129 'label_smoothing', 0.1, 'The amount of label smoothing.')
130
131tf.app.flags.DEFINE_float(
132 'learning_rate_decay_factor', 0.1, 'Learning rate decay factor.')
133
134tf.app.flags.DEFINE_float(
135 'num_epochs_per_decay', 10.0,
136 'Number of epochs after which learning rate decays.')
137
138tf.app.flags.DEFINE_bool(
139 'sync_replicas', False,
140 'Whether or not to synchronize the replicas during training.')
141
142tf.app.flags.DEFINE_integer(
143 'replicas_to_aggregate', 1,
144 'The Number of gradients to collect before updating params.')
145
146tf.app.flags.DEFINE_float(
147 'moving_average_decay', 0.9999,
148 'The decay to use for the moving average.'
149 'If left as None, then moving averages are not used.')
150
151#######################
152# Dataset Flags #
153#######################
154
155tf.app.flags.DEFINE_string(
156 'train_data_path', None, 'Directory containing training records.')
157
158tf.app.flags.DEFINE_string(
159 'train_regex', '*.tfrecords', 'Pattern to match training records.')
160
161tf.app.flags.DEFINE_integer(
162 'height', 720, 'The number of samples in each batch.')
163
164tf.app.flags.DEFINE_integer(
165 'width', 1280, 'The number of samples in each batch.')
166
167tf.app.flags.DEFINE_integer(
168 'batch_size', 4, 'The number of samples in each batch.')
169
170tf.app.flags.DEFINE_integer('max_frames', 100000,
171 'The maximum number of training steps.')
172
173tf.app.flags.DEFINE_integer('num_samples_per_epoch', 4000,
174 'Number of samples per epoch.')
175
176tf.app.flags.DEFINE_string(
177 'detections_prefix', None, 'Path to video file.')
178
179#####################
180# Fine-Tuning Flags #
181#####################
182
183tf.app.flags.DEFINE_string(
184 'checkpoint_path', None,
185 'The path to a checkpoint from which to fine-tune.')
186
187tf.app.flags.DEFINE_string(
188 'checkpoint_exclude_scopes', None,
189 'Comma-separated list of scopes of variables to exclude when restoring '
190 'from a checkpoint.')
191
192tf.app.flags.DEFINE_string(
193 'trainable_scopes', None,
194 'Comma-separated list of scopes to filter the set of variables to train.'
195 'By default, None would train all the variables.')
196
197tf.app.flags.DEFINE_boolean(
198 'ignore_missing_vars', False,
199 'When restoring a checkpoint would ignore missing variables.')
200
201tf.app.flags.DEFINE_boolean(
202 'use_seperable_convolution', False,
203 'Use a seperable convolution block.')
204
205tf.app.flags.DEFINE_float(
206 'filter_depth_multiplier', 1.0,
207 'Filter depth multipler for encoder.')
208
209tf.app.flags.DEFINE_integer(
210 'num_units', 1,
211 'Number of units in each ressep block.')
212
213tf.app.flags.DEFINE_float(
214 'scale', 1.0, 'Input scale factor')
215
216tf.app.flags.DEFINE_integer(
217 'foreground_weight', 10,
218 'Weights for foreground objects.')
219
220tf.app.flags.DEFINE_integer(
221 'background_weight', 10,
222 'Weights for background objects.')
223
224tf.app.flags.DEFINE_boolean(
225 'use_batch_norm', True,
226 'Use batch normalization.')
227
228tf.app.flags.DEFINE_boolean(
229 'only_train_bnorm', False,
230 'Only train the batch norm layers.')
231
232tf.app.flags.DEFINE_boolean(
233 'fine_classes', False,
234 'Only train the batch norm layers.')
235
236tf.app.flags.DEFINE_integer(
237 'training_stride', 25,
238 'Stride at which training is done.')
239
240tf.app.flags.DEFINE_integer(
241 'inference_stride', 10,
242 'Stride at which inference is done.')
243
244tf.app.flags.DEFINE_boolean(
245 'no_summary', True,
246 'Do not compute or write summaries.')
247
248tf.app.flags.DEFINE_string(
249 'stats_path', '',
250 'If set, will output stats to stats_path')
251
252FLAGS = tf.app.flags.FLAGS
253
254def _configure_optimizer(learning_rate):
255 """Configures the optimizer used for training.
256
257 Args:
258 learning_rate: A scalar or `Tensor` learning rate.
259
260 Returns:
261 An instance of an optimizer.
262
263 Raises:
264 ValueError: if FLAGS.optimizer is not recognized.
265 """
266 if FLAGS.optimizer == 'adadelta':
267 optimizer = tf.train.AdadeltaOptimizer(
268 learning_rate,
269 rho=FLAGS.adadelta_rho,
270 epsilon=FLAGS.opt_epsilon)
271 elif FLAGS.optimizer == 'adagrad':
272 optimizer = tf.train.AdagradOptimizer(
273 learning_rate,
274 initial_accumulator_value=FLAGS.adagrad_initial_accumulator_value)
275 elif FLAGS.optimizer == 'adam':
276 optimizer = tf.train.AdamOptimizer(
277 learning_rate,
278 beta1=FLAGS.adam_beta1,
279 beta2=FLAGS.adam_beta2,
280 epsilon=FLAGS.opt_epsilon)
281 elif FLAGS.optimizer == 'ftrl':
282 optimizer = tf.train.FtrlOptimizer(
283 learning_rate,
284 learning_rate_power=FLAGS.ftrl_learning_rate_power,
285 initial_accumulator_value=FLAGS.ftrl_initial_accumulator_value,
286 l1_regularization_strength=FLAGS.ftrl_l1,
287 l2_regularization_strength=FLAGS.ftrl_l2)
288 elif FLAGS.optimizer == 'momentum':
289 optimizer = tf.train.MomentumOptimizer(
290 learning_rate,
291 momentum=FLAGS.momentum,
292 name='Momentum')
293 elif FLAGS.optimizer == 'rmsprop':
294 optimizer = tf.train.RMSPropOptimizer(
295 learning_rate,
296 decay=FLAGS.rmsprop_decay,
297 momentum=FLAGS.momentum,
298 epsilon=FLAGS.opt_epsilon)
299 elif FLAGS.optimizer == 'sgd':
300 optimizer = tf.train.GradientDescentOptimizer(learning_rate)
301 else:
302 raise ValueError('Optimizer [%s] was not recognized', FLAGS.optimizer)
303 return optimizer
304
305def _get_init_fn():
306 """Returns a function run by the chief worker to warm-start the training.
307
308 Note that the init_fn is only run when initializing the model during the very
309 first global step.
310
311 Returns:
312 An init function run by the supervisor.
313 """
314 if FLAGS.checkpoint_path is None:
315 return None
316
317 # Warn the user if a checkpoint exists in the train_dir. Then we'll be
318 # ignoring the checkpoint anyway.
319 if tf.train.latest_checkpoint(FLAGS.train_dir):
320 tf.logging.info(
321 'Ignoring --checkpoint_path because a checkpoint already exists in %s'
322 % FLAGS.train_dir)
323 return None
324
325 exclusions = []
326 if FLAGS.checkpoint_exclude_scopes:
327 exclusions = [scope.strip()
328 for scope in FLAGS.checkpoint_exclude_scopes.split(',')]
329
330 # TODO(sguada) variables.filter_variables()
331 variables_to_restore = []
332 for var in slim.get_model_variables():
333 excluded = False
334 for exclusion in exclusions:
335 if var.op.name.startswith(exclusion):
336 excluded = True
337 break
338 if not excluded:
339 variables_to_restore.append(var)
340
341 if tf.gfile.IsDirectory(FLAGS.checkpoint_path):
342 checkpoint_path = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
343 else:
344 checkpoint_path = FLAGS.checkpoint_path
345
346 tf.logging.info('Fine-tuning from %s' % checkpoint_path)
347
348 return slim.assign_from_checkpoint_fn(
349 checkpoint_path,
350 variables_to_restore,
351 ignore_missing_vars=FLAGS.ignore_missing_vars)
352
353def _get_variables_to_train():
354 """Returns a list of variables to train.
355
356 Returns:
357 A list of variables to train by the optimizer.
358 """
359 if FLAGS.trainable_scopes is None:
360 return tf.trainable_variables()
361 else:
362 scopes = [scope.strip() for scope in FLAGS.trainable_scopes.split(',')]
363
364 variables_to_train = []
365 for scope in scopes:
366 variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)
367 variables_to_train.extend(variables)
368 return variables_to_train
369
370def group_labels(class_groups, batch_segmentation_masks,
371 batch_boxes, batch_classes, batch_masks,
372 batch_num_objects):
373
374 labels = tf.zeros((FLAGS.batch_size, FLAGS.height, FLAGS.width), tf.int32)
375 label_weights = tf.zeros((FLAGS.batch_size, FLAGS.height, FLAGS.width), tf.int32)
376 for g in range(len(class_groups)):
377 seg_group, box_group = tf.py_func(batch_segmentation_masks,
378 [FLAGS.batch_size,
379 (FLAGS.height, FLAGS.width, 3),
380 batch_boxes, batch_classes, batch_masks,
381 batch_num_objects, class_groups[g], True],
382 (tf.bool, tf.bool))
383
384 seg_group.set_shape([FLAGS.batch_size, FLAGS.height, FLAGS.width])
385 box_group.set_shape([FLAGS.batch_size, FLAGS.height, FLAGS.width])
386
387 group_labels = tf.cast(seg_group, tf.int32) * (g + 1)
388 box_labels = tf.cast(box_group, tf.int32) * (g + 1)
389
390 labels = tf.where(tf.cast(seg_group, tf.int32) > 0, group_labels, labels)
391 label_weights = tf.where(tf.cast(box_group, tf.int32) > 0, box_labels, label_weights)
392
393 return labels, label_weights
394
395def ressep_model(input, height, width, scale, weight_decay,
396 use_seperable_convolution, num_classes,
397 filter_depth_multiplier=1, num_units=1,
398 is_training=True, use_batch_norm=True):
399 input = tf.cast(input, tf.float32)
400 if (scale < 1.0):
401 original_dims = [height, width]
402 rescale_dims = [int(scale * height), int(scale * width)]
403 input = tf.image.resize_images(input, rescale_dims)
404 with slim.arg_scope(ressep_bg_seg.ressep_arg_scope(weight_decay=weight_decay,
405 use_batch_norm=use_batch_norm)):
406 net, end_points = ressep_bg_seg.ressep_factory(
407 input,
408 use_seperable_convolution = use_seperable_convolution,
409 filter_depth_multiplier = filter_depth_multiplier,
410 is_training = is_training,
411 use_batch_norm = use_batch_norm,
412 num_units = num_units)
413
414 if (scale < 1.0):
415 net = tf.image.resize_images(net, original_dims)
416
417 logits = slim.conv2d(net, num_classes, [1, 1], normalizer_fn=None,
418 activation_fn=None, scope='logits')
419
420 return logits, end_points
421
422def get_class_groups():
423 people_cls = [1]
424 twowheeler_cls = [2, 4]
425 vehicle_cls = [3, 6, 7, 8]
426
427 #(40, 'bottle')
428 #(41, 'wine glass')
429 #(42, 'cup')
430 #(43, 'fork')
431 #(44, 'knife')
432 #(45, 'spoon')
433 #(46, 'bowl')
434
435 utensils_cls = [40, 41, 42, 43, 44, 45, 46]
436
437 #(14, 'bench')
438 #(57, 'chair')
439 #(58, 'couch')
440 #(61, 'dining table')
441 furniture_cls = [14, 57, 58, 61]
442
443 if FLAGS.fine_classes:
444 cls = [1, 2, 4, 10, 40, 42, 46, 57, 61]
445 class_groups = [[x] for x in cls]
446 class_groups.append([3, 6, 8])
447 else:
448 class_groups = [people_cls, twowheeler_cls, vehicle_cls, utensils_cls, furniture_cls]
449
450 return class_groups
451
452def update_stats(labels, pred_vals, class_tp, class_fp, class_fn,
453 class_total, class_correct, weight_mask, frame_stats,
454 frame_id):
455 eps = 1e-06
456 num_classes = len(class_total)
457 '''
458 curr_tp = np.zeros(num_classes, np.float32)
459 curr_fp = np.zeros(num_classes, np.float32)
460 curr_fn = np.zeros(num_classes, np.float32)
461 curr_iou = np.zeros(num_classes, np.float32)
462 curr_correct = np.zeros(num_classes, np.float32)
463 curr_total = np.zeros(num_classes, np.float32)
464 correct_mask = (pred_vals == labels)
465
466 for g in range(num_classes):
467 cls_mask = np.logical_and((labels == g), weight_mask)
468 cls_tp_mask = np.logical_and(cls_mask, correct_mask)
469 cls_tp = np.sum(cls_tp_mask)
470 curr_tp[g] = cls_tp
471 class_tp[g] = class_tp[g] + cls_tp
472
473 cls_total = np.sum(cls_mask)
474 curr_total[g] = cls_total
475 curr_correct[g] = cls_tp
476 class_total[g] = class_total[g] + cls_total
477 class_correct[g] = class_correct[g] + cls_tp
478
479 pred_mask = np.logical_and((pred_vals == g), weight_mask)
480 cls_fp_mask = np.logical_and(np.logical_not(cls_mask), pred_mask)
481 cls_fn_mask = np.logical_and(cls_mask, np.logical_not(pred_mask))
482
483 cls_fp = np.sum(cls_fp_mask)
484 cls_fn = np.sum(cls_fn_mask)
485 curr_fp[g] = cls_fp
486 curr_fn[g] = cls_fn
487 class_fp[g] = class_fp[g] + cls_fp
488 class_fn[g] = class_fn[g] + cls_fn
489
490 cls_iou = (cls_tp + eps) / (cls_tp + cls_fp + cls_fn + eps)
491 curr_iou[g] = cls_iou
492
493 frame_stats[frame_id] = { 'tp': curr_tp,
494 'fp': curr_fp,
495 'fn': curr_fn,
496 'iou': curr_iou,
497 'correct': curr_correct,
498 'total': curr_total }
499 print(frame_stats[frame_id]["iou"])
500 '''
501 curr_tp = np.zeros(num_classes, np.float32)
502 curr_fp = np.zeros(num_classes, np.float32)
503 curr_fn = np.zeros(num_classes, np.float32)
504 curr_iou = np.zeros(num_classes, np.float32)
505 curr_correct = np.zeros(num_classes, np.float32)
506 curr_total = np.zeros(num_classes, np.float32)
507 # assuming (1, height, width)
508 height = pred_vals.shape[1]
509 width = pred_vals.shape[2]
510 correct_mask = (pred_vals == labels)
511 # vectorize correct_mask
512 correct_mask = np.tile(correct_mask, (num_classes, 1)).reshape((1, num_classes, height, width))
513 cls_mask = np.zeros((1, num_classes, height, width), dtype=bool)
514 pred_mask = np.zeros((1, num_classes, height, width), dtype=bool)
515 for g in range(num_classes):
516 cls_mask[0, g] = np.logical_and((labels == g), weight_mask)
517 pred_mask[0, g] = np.logical_and((pred_vals == g), weight_mask)
518
519 cls_tp_mask = np.logical_and(cls_mask, correct_mask)
520 cls_tp = np.sum(cls_tp_mask, axis=(2, 3))
521 curr_tp = cls_tp
522 class_tp = class_tp + cls_tp
523
524 cls_total = np.sum(cls_mask, axis=(2, 3))
525 curr_total = cls_total
526 curr_correct = cls_tp
527 class_total = class_total + cls_total
528 class_correct = class_correct + cls_tp
529
530 cls_fp_mask = np.logical_and(np.logical_not(cls_mask), pred_mask)
531 cls_fn_mask = np.logical_and(cls_mask, np.logical_not(pred_mask))
532
533 cls_fp = np.sum(cls_fp_mask, axis=(2, 3))
534 cls_fn = np.sum(cls_fn_mask, axis=(2, 3))
535 curr_fp = cls_fp
536 curr_fn = cls_fn
537 class_fp = class_fp + cls_fp
538 class_fn = class_fn + cls_fn
539
540 for g in range(num_classes):
541 cls_iou = (cls_tp[0][g] + eps) / (cls_tp[0][g] + cls_fp[0][g] + cls_fn[0][g] + eps)
542 curr_iou[g] = cls_iou
543
544 frame_stats[frame_id] = { 'tp': curr_tp,
545 'fp': curr_fp,
546 'fn': curr_fn,
547 'iou': curr_iou,
548 'correct': curr_correct,
549 'total': curr_total }
550
551 #print(frame_stats[frame_id]["iou"])
552
553
554def main(_):
555 tf.logging.set_verbosity(tf.logging.INFO)
556 with tf.Graph().as_default():
557 #######################
558 # Config model_deploy #
559 #######################
560 deploy_config = model_deploy.DeploymentConfig(
561 num_clones=FLAGS.num_clones,
562 clone_on_cpu=FLAGS.clone_on_cpu,
563 replica_id=FLAGS.task,
564 num_replicas=FLAGS.worker_replicas,
565 num_ps_tasks=FLAGS.num_ps_tasks)
566
567 # Create global_step
568 with tf.device(deploy_config.variables_device()):
569 global_step = slim.create_global_step()
570
571 class_groups = get_class_groups()
572
573 with tf.device(deploy_config.inputs_device()):
574 video_files = []
575 detections_paths = []
576 video_paths = FLAGS.train_data_path.split(',')
577 for path in video_paths:
578 video_files = video_files + glob.glob(os.path.join(path, FLAGS.train_regex))
579 segment_names = [ s.split('/')[-1].split('.')[0] for s in video_files ]
580 for s in segment_names:
581 detections_paths.append(os.path.join(path, FLAGS.detections_prefix + '_' + s + '.npy'))
582
583 batch_image = tf.placeholder(tf.float32, (FLAGS.batch_size,
584 FLAGS.height,
585 FLAGS.width, 3))
586 labels = tf.placeholder(tf.int32, (FLAGS.batch_size,
587 FLAGS.height,
588 FLAGS.width))
589
590 label_weights_in = tf.placeholder(tf.int32, (FLAGS.batch_size,
591 FLAGS.height,
592 FLAGS.width))
593
594 low_weights = tf.constant(FLAGS.background_weight,
595 dtype=tf.int32, shape=labels.shape)
596 high_weights = tf.constant(FLAGS.foreground_weight,
597 dtype=tf.int32, shape=labels.shape)
598 label_weights = tf.where(label_weights_in > 0,
599 high_weights, low_weights)
600
601 #labels, label_weights = group_labels(class_groups, batch_segmentation_masks,
602 # batch_boxes, batch_classes, batch_masks,
603 # batch_num_objects)
604
605 def clone_fn(batch_image, labels, label_weights):
606
607 labels = tf.cast(labels, tf.int32)
608 labels.set_shape([FLAGS.batch_size, FLAGS.height, FLAGS.width])
609
610 batch_image.set_shape([1, FLAGS.height, FLAGS.width, 3])
611
612 num_classes = len(class_groups) + 1
613
614 logits, end_points = ressep_model(batch_image, FLAGS.height,
615 FLAGS.width, FLAGS.scale, FLAGS.weight_decay,
616 FLAGS.use_seperable_convolution,
617 num_classes,
618 is_training=True,
619 use_batch_norm=FLAGS.use_batch_norm,
620 num_units=FLAGS.num_units,
621 filter_depth_multiplier=FLAGS.filter_depth_multiplier)
622
623 # Specify loss
624 cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels,
625 logits,
626 weights=label_weights,
627 scope='xentropy')
628 tf.losses.add_loss(cross_entropy)
629
630 return end_points, labels, label_weights, logits
631
632 # Gather initial summaries.
633 summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
634
635 clones = model_deploy.create_clones(deploy_config, clone_fn, [batch_image, labels, label_weights])
636 first_clone_scope = deploy_config.clone_scope(0)
637 # Gather update_ops from the first clone. These contain, for example,
638 # the updates for the batch_norm variables created by network_fn.
639 update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, first_clone_scope)
640
641 #################################
642 # Configure the moving averages #
643 #################################
644 if FLAGS.moving_average_decay:
645 moving_average_variables = slim.get_model_variables()
646 variable_averages = tf.train.ExponentialMovingAverage(
647 FLAGS.moving_average_decay, global_step)
648 else:
649 moving_average_variables, variable_averages = None, None
650
651 #########################################
652 # Configure the optimization procedure. #
653 #########################################
654 with tf.device(deploy_config.optimizer_device()):
655 num_training_records = FLAGS.num_samples_per_epoch
656 learning_rate = FLAGS.learning_rate
657 optimizer = _configure_optimizer(learning_rate)
658 summaries.add(tf.summary.scalar('learning_rate', learning_rate))
659
660 if FLAGS.sync_replicas:
661 # If sync_replicas is enabled, the averaging will be done in the chief
662 # queue runner.
663 optimizer = tf.train.SyncReplicasOptimizer(
664 opt=optimizer,
665 replicas_to_aggregate=FLAGS.replicas_to_aggregate,
666 variable_averages=variable_averages,
667 variables_to_average=moving_average_variables,
668 replica_id=tf.constant(FLAGS.task, tf.int32, shape=()),
669 total_num_replicas=FLAGS.worker_replicas)
670 elif FLAGS.moving_average_decay:
671 # Update ops executed locally by trainer.
672 update_ops.append(variable_averages.apply(moving_average_variables))
673
674 end_points, labels_tensor, label_weights_tensor, logits_tensor = clones[0].outputs
675
676 predictions = tf.argmax(logits_tensor, axis=3)
677
678 # Add accuracy summaries
679 correct_preds = tf.equal(tf.cast(predictions, tf.int32), labels_tensor)
680 total_preds = tf.reduce_sum(tf.cast(labels_tensor >=0, tf.float32))
681 batch_accuracy = tf.reduce_sum(tf.cast(correct_preds, tf.float32)) / total_preds
682
683 acc_collection = tf.get_collection('Accuracy')
684 iou_collection = tf.get_collection('IOU')
685
686 summaries.add(tf.summary.scalar('Accuracy', batch_accuracy, collections = acc_collection))
687
688 eps = 1e-06
689 for cls in range(len(class_groups) + 1):
690 cls_mask = tf.equal(labels_tensor, cls)
691 cls_correct = tf.reduce_sum(
692 tf.cast(tf.logical_and(cls_mask, correct_preds), tf.float32))
693 cls_total = tf.reduce_sum(tf.cast(cls_mask, tf.float32))
694 cls_accuracy = cls_correct / tf.add(cls_total, eps)
695
696 pred_mask = tf.equal(predictions, cls)
697 pred_not_mask = tf.not_equal(predictions, cls)
698 label_mask = tf.equal(labels, cls)
699 label_not_mask = tf.not_equal(labels, cls)
700 cls_tp = cls_correct
701 cls_fp = tf.reduce_sum(tf.cast(tf.logical_and(pred_mask, label_not_mask), tf.float32))
702 cls_fn = tf.reduce_sum(tf.cast(tf.logical_and(pred_not_mask, label_mask), tf.float32))
703 cls_iou = cls_tp / tf.add(cls_tp + cls_fp + cls_fn, eps)
704 summaries.add(tf.summary.scalar('Accuracy_cls_%d'%(cls), cls_accuracy, collections=acc_collection))
705 summaries.add(tf.summary.scalar('IOU_cls_%d'%(cls), cls_iou, collections=iou_collection))
706
707 num_classes = len(class_groups) + 1
708 labels_vis = tf.py_func(visualize_masks, [labels_tensor, FLAGS.batch_size,
709 (FLAGS.height, FLAGS.width, 3), num_classes],
710 tf.uint8)
711 pred_vis = tf.py_func(visualize_masks, [predictions, FLAGS.batch_size,
712 (FLAGS.height, FLAGS.width, 3), num_classes],
713 tf.uint8)
714
715 label_weights_tensor = tf.cast(label_weights_tensor, tf.float32)
716 min_weight = tf.reduce_min(label_weights_tensor)
717 max_weight = tf.reduce_max(label_weights_tensor)
718 scale_weights = (tf.add(label_weights_tensor, eps) - min_weight)/(tf.add(max_weight, eps) - min_weight)
719
720 labels_vis = 0.5 * tf.cast(batch_image, tf.float32) + 0.5 * tf.cast(labels_vis, tf.float32)
721 pred_vis = 0.5 * tf.cast(batch_image, tf.float32) + 0.5 * tf.cast(pred_vis, tf.float32)
722
723 label_weights_vis = tf.cast(batch_image, tf.float32) * tf.expand_dims(scale_weights, axis=3)
724
725 def rearrage_channels(img):
726 img.set_shape([FLAGS.batch_size, FLAGS.height, FLAGS.width, 3])
727 channels = tf.unstack(img, axis=-1)
728 img = tf.stack([channels[2], channels[1], channels[0]], axis=-1)
729 return img
730
731 labels_vis = tf.cast(labels_vis, tf.uint8)
732 pred_vis = tf.cast(pred_vis, tf.uint8)
733 label_weights_vis = tf.cast(label_weights_vis, tf.uint8)
734
735 labels_vis = rearrage_channels(labels_vis)
736 pred_vis = rearrage_channels(pred_vis)
737 label_weights_vis = rearrage_channels(label_weights_vis)
738
739 summary_img = tf.concat([labels_vis, pred_vis, label_weights_vis], axis=1)
740
741 tf.summary.image('summary_image', summary_img, max_outputs=4)
742
743 # Variables to train.
744 variables_to_train = _get_variables_to_train()
745
746 if FLAGS.only_train_bnorm:
747 variables_to_train = tf.contrib.framework.filter_variables(variables_to_train,
748 include_patterns=['BatchNorm', 'logits'])
749
750 # and returns a train_tensor and summary_op
751 total_loss, clones_gradients = model_deploy.optimize_clones(clones,
752 optimizer,
753 var_list=variables_to_train)
754 # Add total_loss to summary.
755 summaries.add(tf.summary.scalar('total_loss', total_loss))
756
757 # Create gradient updates.
758 grad_updates = optimizer.apply_gradients(clones_gradients,
759 global_step=global_step)
760 update_ops.append(grad_updates)
761
762 update_op = tf.group(*update_ops)
763 train_tensor = control_flow_ops.with_dependencies([update_op], total_loss,
764 name='train_op')
765
766 # Add the summaries from the first clone. These contain the summaries
767 # created by model_fn and either optimize_clones() or _gather_clone_loss().
768 summaries |= set(tf.get_collection(tf.GraphKeys.SUMMARIES,
769 first_clone_scope))
770 # Merge all summaries together.
771 summary_op = tf.summary.merge(list(summaries), name='summary_op')
772
773 # for stats training, reuse memory
774 per_frame_stats = {}
775 class_correct = np.zeros(num_classes, np.float32)
776 class_total = np.zeros(num_classes, np.float32)
777
778 class_tp = np.zeros(num_classes, np.float32)
779 class_fp = np.zeros(num_classes, np.float32)
780 class_fn = np.zeros(num_classes, np.float32)
781 class_iou = np.zeros(num_classes, np.float32)
782
783 stats_path = FLAGS.stats_path
784
785 if FLAGS.sync_replicas:
786 sync_optimizer = opt
787 startup_delay_steps = 0
788 else:
789 sync_optimizer = None
790 startup_delay_steps = FLAGS.task * FLAGS.startup_delay_steps
791
792 ###########################
793 # Kicks off the training. #
794 ###########################
795 #with tf.contrib.tfprof.ProfileContext(FLAGS.train_dir) as pctx:
796
797 print(video_files, detections_paths)
798 input_streams = MaskRCNNMultiStream(video_files, detections_paths,
799 start_frame=0, stride=FLAGS.inference_stride)
800
801 init_fn = _get_init_fn()
802
803 with tf.Session() as sess:
804 sess.run(tf.global_variables_initializer())
805 sess.run(tf.local_variables_initializer())
806
807 init_fn(sess)
808
809 curr_frame = 0
810
811 summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
812
813 zero_input = np.zeros((FLAGS.batch_size, FLAGS.height, FLAGS.width),
814 np.int32)
815 inference_str = ""
816 training_str = ""
817 stats_str = ""
818 for frame, boxes, classes, scores, masks, num_objects, frame_id in input_streams:
819 if curr_frame > FLAGS.max_frames:
820 break
821 frame = np.expand_dims(frame, axis=0)
822 boxes = np.expand_dims(boxes, axis=0)
823 classes = np.expand_dims(classes, axis=0)
824 scores = np.expand_dims(scores, axis=0)
825 masks = np.expand_dims(masks, axis=0)
826 num_objects = np.expand_dims(num_objects, axis=0)
827
828 preds = None
829 summary_str = ""
830 if curr_frame % FLAGS.training_stride == 0:
831 start = time.time()
832 labels_val, label_weights_val = batch_segmentation_masks_fast(1,
833 (FLAGS.height, FLAGS.width),
834 boxes, classes, masks,
835 num_objects, True,
836 class_groups)
837
838 labels_val = labels_val.astype(np.int32)
839 label_weights_val = label_weights_val.astype(np.int32)
840
841 if FLAGS.no_summary or curr_frame%FLAGS.summary_stride != 0:
842 step, _, loss, preds = sess.run([global_step, train_tensor,
843 total_loss, predictions],
844 feed_dict={batch_image: frame,
845 labels: labels_val,
846 label_weights_in: label_weights_val})
847 else:
848 step, summary, _, loss, preds = sess.run([global_step, summary_op,
849 train_tensor, total_loss, predictions],
850 feed_dict={batch_image: frame,
851 labels: labels_val,
852 label_weights_in: label_weights_val})
853 summary_str = "summary: {0:5d}".format(step)
854 summary_writer.add_summary(summary, step)
855 end = time.time()
856 training_str = "training: {0:.5f}s loss: {1:.5f}".format(end - start, loss)
857 elif curr_frame % FLAGS.inference_stride == 0:
858 start = time.time()
859 step, preds = sess.run([global_step, predictions],
860 feed_dict={ batch_image: frame,
861 labels: zero_input,
862 label_weights_in: zero_input})
863 end = time.time()
864 inference_str = "inference: {0:.5f}s".format(end - start)
865 # compute stats
866 if stats_path:
867 start = time.time()
868 labels_vals, _ = batch_segmentation_masks_fast(1,
869 (FLAGS.height, FLAGS.width),
870 boxes, classes, masks,
871 num_objects, True,
872 class_groups)
873
874 update_stats(labels_vals, preds, class_tp, class_fp, class_fn,
875 class_total, class_correct, np.ones(labels_vals.shape, dtype=np.bool),
876 per_frame_stats, curr_frame)
877 end = time.time()
878 stats_str = "stats: {0:.5f}s".format(end - start)
879
880 frame_str = "frame: {0:05d}".format(curr_frame)
881 print(" ".join([frame_str, training_str, inference_str, stats_str,
882 summary_str]), end="\r", file=sys.stderr)
883 curr_frame = curr_frame + 1
884 print(file=sys.stderr)
885 summary_writer.close()
886 if stats_path:
887 np.save(stats_path, [per_frame_stats])
888
889if __name__ == '__main__':
890 tf.app.run()