text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  # File: train.py
  
  import os
  import argparse
  import cv2
  import shutil
  import itertools
  import tqdm
  import numpy as np
  import json
  import six
  import tensorflow as tf
  try:
      import horovod.tensorflow as hvd
  except ImportError:
      pass
  
  assert six.PY3, "FasterRCNN requires Python 3!"
  
  from tensorpack import *
  from tensorpack.tfutils.summary import add_moving_summary
  from tensorpack.tfutils import optimizer
  from tensorpack.tfutils.common import get_tf_version_tuple
  import tensorpack.utils.viz as tpviz
  
  from coco import COCODetection
  from basemodel import (
      image_preprocess, resnet_c4_backbone, resnet_conv5,
      resnet_fpn_backbone)
  
  import model_frcnn
  import model_mrcnn
  from model_frcnn import (
      sample_fast_rcnn_targets,
      fastrcnn_outputs, fastrcnn_losses, fastrcnn_predictions)
  from model_mrcnn import maskrcnn_upXconv_head, maskrcnn_loss
  from model_rpn import rpn_head, rpn_losses, generate_rpn_proposals
  from model_fpn import (
      fpn_model, multilevel_roi_align,
      multilevel_rpn_losses, generate_fpn_proposals)
  from model_box import (
      clip_boxes, decode_bbox_target, encode_bbox_target,
      crop_and_resize, roi_align, RPNAnchors)
  
  from data import (
      get_train_dataflow, get_eval_dataflow,
      get_all_anchors, get_all_anchors_fpn)
  from viz import (
      draw_annotation, draw_proposal_recall,
      draw_predictions, draw_final_outputs)
  from eval import (
      eval_coco, detect_one_image, print_evaluation_scores, DetectionResult)
  from config import finalize_configs, config as cfg
  
  
  class DetectionModel(ModelDesc):
      def preprocess(self, image):
          image = tf.expand_dims(image, 0)
          image = image_preprocess(image, bgr=True)
          return tf.transpose(image, [0, 3, 1, 2])
  
      def optimizer(self):
          lr = tf.get_variable('learning_rate', initializer=0.003, trainable=False)
          tf.summary.scalar('learning_rate-summary', lr)
  
          factor = cfg.TRAIN.NUM_GPUS / 8.
          if factor != 1:
              lr = lr * factor
          opt = tf.train.MomentumOptimizer(lr, 0.9)
          if cfg.TRAIN.NUM_GPUS < 8:
              opt = optimizer.AccumGradOptimizer(opt, 8 // cfg.TRAIN.NUM_GPUS)
          return opt
  
      def fastrcnn_training(self, image,
                            rcnn_labels, fg_rcnn_boxes, gt_boxes_per_fg,
                            rcnn_label_logits, fg_rcnn_box_logits):
          """
          Args:
              image (NCHW):
              rcnn_labels (n): labels for each sampled targets
              fg_rcnn_boxes (fg x 4): proposal boxes for each sampled foreground targets
              gt_boxes_per_fg (fg x 4): matching gt boxes for each sampled foreground targets
              rcnn_label_logits (n): label logits for each sampled targets
              fg_rcnn_box_logits (fg x #class x 4): box logits for each sampled foreground targets
          """
  
          with tf.name_scope('fg_sample_patch_viz'):
              fg_sampled_patches = crop_and_resize(
                  image, fg_rcnn_boxes,
                  tf.zeros([tf.shape(fg_rcnn_boxes)[0]], dtype=tf.int32), 300)
              fg_sampled_patches = tf.transpose(fg_sampled_patches, [0, 2, 3, 1])
              fg_sampled_patches = tf.reverse(fg_sampled_patches, axis=[-1])  # BGR->RGB
              tf.summary.image('viz', fg_sampled_patches, max_outputs=30)
  
          encoded_boxes = encode_bbox_target(
              gt_boxes_per_fg, fg_rcnn_boxes) * tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32)
          fastrcnn_label_loss, fastrcnn_box_loss = fastrcnn_losses(
              rcnn_labels, rcnn_label_logits,
              encoded_boxes,
              fg_rcnn_box_logits)
          return fastrcnn_label_loss, fastrcnn_box_loss
  
      def fastrcnn_inference(self, image_shape2d,
                             rcnn_boxes, rcnn_label_logits, rcnn_box_logits):
          """
          Args:
              image_shape2d: h, w
              rcnn_boxes (nx4): the proposal boxes
              rcnn_label_logits (n):
              rcnn_box_logits (nx #class x 4):
  
          Returns:
              boxes (mx4):
              labels (m): each >= 1
          """
          rcnn_box_logits = rcnn_box_logits[:, 1:, :]
          rcnn_box_logits.set_shape([None, cfg.DATA.NUM_CATEGORY, None])
          label_probs = tf.nn.softmax(rcnn_label_logits, name='fastrcnn_all_probs')  # #proposal x #Class
          anchors = tf.tile(tf.expand_dims(rcnn_boxes, 1), [1, cfg.DATA.NUM_CATEGORY, 1])   # #proposal x #Cat x 4
          decoded_boxes = decode_bbox_target(
              rcnn_box_logits /
              tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32), anchors)
          decoded_boxes = clip_boxes(decoded_boxes, image_shape2d, name='fastrcnn_all_boxes')
  
          # indices: Nx2. Each index into (#proposal, #category)
          pred_indices, final_probs = fastrcnn_predictions(decoded_boxes, label_probs)
          final_probs = tf.identity(final_probs, 'final_probs')
          final_boxes = tf.gather_nd(decoded_boxes, pred_indices, name='final_boxes')
          final_labels = tf.add(pred_indices[:, 1], 1, name='final_labels')
          return final_boxes, final_labels
  
      def get_inference_tensor_names(self):
          """
          Returns two lists of tensor names to be used to create an inference callable.
  
          Returns:
              [str]: input names
              [str]: output names
          """
          out = ['final_boxes', 'final_probs', 'final_labels']
          if cfg.MODE_MASK:
              out.append('final_masks')
          return ['image'], out
  
  
  class ResNetC4Model(DetectionModel):
      def inputs(self):
          ret = [
              tf.placeholder(tf.float32, (None, None, 3), 'image'),
              tf.placeholder(tf.int32, (None, None, cfg.RPN.NUM_ANCHOR), 'anchor_labels'),
              tf.placeholder(tf.float32, (None, None, cfg.RPN.NUM_ANCHOR, 4), 'anchor_boxes'),
              tf.placeholder(tf.float32, (None, 4), 'gt_boxes'),
              tf.placeholder(tf.int64, (None,), 'gt_labels')]  # all > 0
          if cfg.MODE_MASK:
              ret.append(
                  tf.placeholder(tf.uint8, (None, None, None), 'gt_masks')
              )   # NR_GT x height x width
          return ret
  
      def build_graph(self, *inputs):
          is_training = get_current_tower_context().is_training
          if cfg.MODE_MASK:
              image, anchor_labels, anchor_boxes, gt_boxes, gt_labels, gt_masks = inputs
          else:
              image, anchor_labels, anchor_boxes, gt_boxes, gt_labels = inputs
          image = self.preprocess(image)     # 1CHW
  
          featuremap = resnet_c4_backbone(image, cfg.BACKBONE.RESNET_NUM_BLOCK[:3])
          rpn_label_logits, rpn_box_logits = rpn_head('rpn', featuremap, cfg.RPN.HEAD_DIM, cfg.RPN.NUM_ANCHOR)
  
          anchors = RPNAnchors(get_all_anchors(), anchor_labels, anchor_boxes)
          anchors = anchors.narrow_to(featuremap)
  
          image_shape2d = tf.shape(image)[2:]     # h,w
          pred_boxes_decoded = anchors.decode_logits(rpn_box_logits)  # fHxfWxNAx4, floatbox
          proposal_boxes, proposal_scores = generate_rpn_proposals(
              tf.reshape(pred_boxes_decoded, [-1, 4]),
              tf.reshape(rpn_label_logits, [-1]),
              image_shape2d,
              cfg.RPN.TRAIN_PRE_NMS_TOPK if is_training else cfg.RPN.TEST_PRE_NMS_TOPK,
              cfg.RPN.TRAIN_POST_NMS_TOPK if is_training else cfg.RPN.TEST_POST_NMS_TOPK)
  
          if is_training:
              # sample proposal boxes in training
              rcnn_boxes, rcnn_labels, fg_inds_wrt_gt = sample_fast_rcnn_targets(
                  proposal_boxes, gt_boxes, gt_labels)
          else:
              # The boxes to be used to crop RoIs.
              # Use all proposal boxes in inference
              rcnn_boxes = proposal_boxes
  
          boxes_on_featuremap = rcnn_boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE)
          roi_resized = roi_align(featuremap, boxes_on_featuremap, 14)
  
          feature_fastrcnn = resnet_conv5(roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCK[-1])    # nxcx7x7
          # Keep C5 feature to be shared with mask branch
          feature_gap = GlobalAvgPooling('gap', feature_fastrcnn, data_format='channels_first')
          fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs('fastrcnn', feature_gap, cfg.DATA.NUM_CLASS)
  
          if is_training:
              # rpn loss
              rpn_label_loss, rpn_box_loss = rpn_losses(
                  anchors.gt_labels, anchors.encoded_gt_boxes(), rpn_label_logits, rpn_box_logits)
  
              # fastrcnn loss
              matched_gt_boxes = tf.gather(gt_boxes, fg_inds_wrt_gt)
  
              fg_inds_wrt_sample = tf.reshape(tf.where(rcnn_labels > 0), [-1])   # fg inds w.r.t all samples
              fg_sampled_boxes = tf.gather(rcnn_boxes, fg_inds_wrt_sample)
              fg_fastrcnn_box_logits = tf.gather(fastrcnn_box_logits, fg_inds_wrt_sample)
  
              fastrcnn_label_loss, fastrcnn_box_loss = self.fastrcnn_training(
                  image, rcnn_labels, fg_sampled_boxes,
                  matched_gt_boxes, fastrcnn_label_logits, fg_fastrcnn_box_logits)
  
              if cfg.MODE_MASK:
                  # maskrcnn loss
                  fg_labels = tf.gather(rcnn_labels, fg_inds_wrt_sample)
                  # In training, mask branch shares the same C5 feature.
                  fg_feature = tf.gather(feature_fastrcnn, fg_inds_wrt_sample)
                  mask_logits = maskrcnn_upXconv_head(
                      'maskrcnn', fg_feature, cfg.DATA.NUM_CATEGORY, num_convs=0)   # #fg x #cat x 14x14
  
                  target_masks_for_fg = crop_and_resize(
                      tf.expand_dims(gt_masks, 1),
                      fg_sampled_boxes,
                      fg_inds_wrt_gt, 14,
                      pad_border=False)  # nfg x 1x14x14
                  target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1, 'sampled_fg_mask_targets')
                  mrcnn_loss = maskrcnn_loss(mask_logits, fg_labels, target_masks_for_fg)
              else:
                  mrcnn_loss = 0.0
  
              wd_cost = regularize_cost(
                  '(?:group1|group2|group3|rpn|fastrcnn|maskrcnn)/.*W',
                  l2_regularizer(cfg.TRAIN.WEIGHT_DECAY), name='wd_cost')
  
              total_cost = tf.add_n([
                  rpn_label_loss, rpn_box_loss,
                  fastrcnn_label_loss, fastrcnn_box_loss,
                  mrcnn_loss, wd_cost], 'total_cost')
  
              add_moving_summary(total_cost, wd_cost)
              return total_cost * (1. / cfg.TRAIN.NUM_GPUS)
          else:
              final_boxes, final_labels = self.fastrcnn_inference(
                  image_shape2d, rcnn_boxes, fastrcnn_label_logits, fastrcnn_box_logits)
  
              if cfg.MODE_MASK:
                  roi_resized = roi_align(featuremap, final_boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE), 14)
                  feature_maskrcnn = resnet_conv5(roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCK[-1])
                  mask_logits = maskrcnn_upXconv_head(
                      'maskrcnn', feature_maskrcnn, cfg.DATA.NUM_CATEGORY, 0)   # #result x #cat x 14x14
                  indices = tf.stack([tf.range(tf.size(final_labels)), tf.to_int32(final_labels) - 1], axis=1)
                  final_mask_logits = tf.gather_nd(mask_logits, indices)   # #resultx14x14
                  tf.sigmoid(final_mask_logits, name='final_masks')
  
  
  class ResNetFPNModel(DetectionModel):
  
      def inputs(self):
          ret = [
              tf.placeholder(tf.float32, (None, None, 3), 'image')]
          num_anchors = len(cfg.RPN.ANCHOR_RATIOS)
          for k in range(len(cfg.FPN.ANCHOR_STRIDES)):
              ret.extend([
                  tf.placeholder(tf.int32, (None, None, num_anchors),
                                 'anchor_labels_lvl{}'.format(k + 2)),
                  tf.placeholder(tf.float32, (None, None, num_anchors, 4),
                                 'anchor_boxes_lvl{}'.format(k + 2))])
          ret.extend([
              tf.placeholder(tf.float32, (None, 4), 'gt_boxes'),
              tf.placeholder(tf.int64, (None,), 'gt_labels')])  # all > 0
          if cfg.MODE_MASK:
              ret.append(
                  tf.placeholder(tf.uint8, (None, None, None), 'gt_masks')
              )   # NR_GT x height x width
          return ret
  
      def slice_feature_and_anchors(self, image_shape2d, p23456, anchors):
          for i, stride in enumerate(cfg.FPN.ANCHOR_STRIDES):
              with tf.name_scope('FPN_slice_lvl{}'.format(i)):
                  if i < 3:
                      # Images are padded for p5, which are too large for p2-p4.
                      # This seems to have no effect on mAP.
                      pi = p23456[i]
                      target_shape = tf.to_int32(tf.ceil(tf.to_float(image_shape2d) * (1.0 / stride)))
                      p23456[i] = tf.slice(pi, [0, 0, 0, 0],
                                           tf.concat([[-1, -1], target_shape], axis=0))
                      p23456[i].set_shape([1, pi.shape[1], None, None])
  
                  anchors[i] = anchors[i].narrow_to(p23456[i])
  
      def build_graph(self, *inputs):
          num_fpn_level = len(cfg.FPN.ANCHOR_STRIDES)
          assert len(cfg.RPN.ANCHOR_SIZES) == num_fpn_level
          is_training = get_current_tower_context().is_training
          image = inputs[0]
          input_anchors = inputs[1: 1 + 2 * num_fpn_level]
          multilevel_anchors = [RPNAnchors(*args) for args in
                                zip(get_all_anchors_fpn(), input_anchors[0::2], input_anchors[1::2])]
          gt_boxes, gt_labels = inputs[11], inputs[12]
          if cfg.MODE_MASK:
              gt_masks = inputs[-1]
  
          image = self.preprocess(image)     # 1CHW
          image_shape2d = tf.shape(image)[2:]     # h,w
  
          c2345 = resnet_fpn_backbone(image, cfg.BACKBONE.RESNET_NUM_BLOCK)
          p23456 = fpn_model('fpn', c2345)
          self.slice_feature_and_anchors(image_shape2d, p23456, multilevel_anchors)
  
          # Multi-Level RPN Proposals
          rpn_outputs = [rpn_head('rpn', pi, cfg.FPN.NUM_CHANNEL, len(cfg.RPN.ANCHOR_RATIOS))
                         for pi in p23456]
          multilevel_label_logits = [k[0] for k in rpn_outputs]
          multilevel_box_logits = [k[1] for k in rpn_outputs]
  
          proposal_boxes, proposal_scores = generate_fpn_proposals(
              multilevel_anchors, multilevel_label_logits,
              multilevel_box_logits, image_shape2d)
  
          if is_training:
              rcnn_boxes, rcnn_labels, fg_inds_wrt_gt = sample_fast_rcnn_targets(
                  proposal_boxes, gt_boxes, gt_labels)
          else:
              # The boxes to be used to crop RoIs.
              rcnn_boxes = proposal_boxes
  
          roi_feature_fastrcnn = multilevel_roi_align(p23456[:4], rcnn_boxes, 7)
  
          fastrcnn_head_func = getattr(model_frcnn, cfg.FPN.FRCNN_HEAD_FUNC)
          fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_head_func(
              'fastrcnn', roi_feature_fastrcnn, cfg.DATA.NUM_CLASS)
  
          if is_training:
              # rpn loss:
              rpn_label_loss, rpn_box_loss = multilevel_rpn_losses(
                  multilevel_anchors, multilevel_label_logits, multilevel_box_logits)
  
              # fastrcnn loss:
              matched_gt_boxes = tf.gather(gt_boxes, fg_inds_wrt_gt)
  
              fg_inds_wrt_sample = tf.reshape(tf.where(rcnn_labels > 0), [-1])   # fg inds w.r.t all samples
              fg_sampled_boxes = tf.gather(rcnn_boxes, fg_inds_wrt_sample)
              fg_fastrcnn_box_logits = tf.gather(fastrcnn_box_logits, fg_inds_wrt_sample)
  
              fastrcnn_label_loss, fastrcnn_box_loss = self.fastrcnn_training(
                  image, rcnn_labels, fg_sampled_boxes,
                  matched_gt_boxes, fastrcnn_label_logits, fg_fastrcnn_box_logits)
  
              if cfg.MODE_MASK:
                  # maskrcnn loss
                  fg_labels = tf.gather(rcnn_labels, fg_inds_wrt_sample)
                  roi_feature_maskrcnn = multilevel_roi_align(
                      p23456[:4], fg_sampled_boxes, 14,
                      name_scope='multilevel_roi_align_mask')
                  maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                  mask_logits = maskrcnn_head_func(
                      'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)   # #fg x #cat x 28 x 28
  
                  target_masks_for_fg = crop_and_resize(
                      tf.expand_dims(gt_masks, 1),
                      fg_sampled_boxes,
                      fg_inds_wrt_gt, 28,
                      pad_border=False)  # fg x 1x28x28
                  target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1, 'sampled_fg_mask_targets')
                  mrcnn_loss = maskrcnn_loss(mask_logits, fg_labels, target_masks_for_fg)
              else:
                  mrcnn_loss = 0.0
  
              wd_cost = regularize_cost(
                  '(?:group1|group2|group3|rpn|fpn|fastrcnn|maskrcnn)/.*W',
                  l2_regularizer(cfg.TRAIN.WEIGHT_DECAY), name='wd_cost')
  
              total_cost = tf.add_n([rpn_label_loss, rpn_box_loss,
                                     fastrcnn_label_loss, fastrcnn_box_loss,
                                     mrcnn_loss, wd_cost], 'total_cost')
  
              add_moving_summary(total_cost, wd_cost)
              return total_cost * (1. / cfg.TRAIN.NUM_GPUS)
          else:
              final_boxes, final_labels = self.fastrcnn_inference(
                  image_shape2d, rcnn_boxes, fastrcnn_label_logits, fastrcnn_box_logits)
              if cfg.MODE_MASK:
                  # Cascade inference needs roi transform with refined boxes.
                  roi_feature_maskrcnn = multilevel_roi_align(p23456[:4], final_boxes, 14)
                  maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                  mask_logits = maskrcnn_head_func(
                      'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)   # #fg x #cat x 28 x 28
                  indices = tf.stack([tf.range(tf.size(final_labels)), tf.to_int32(final_labels) - 1], axis=1)
                  final_mask_logits = tf.gather_nd(mask_logits, indices)   # #resultx28x28
                  tf.sigmoid(final_mask_logits, name='final_masks')
  
  
  def visualize(model, model_path, nr_visualize=100, output_dir='output'):
      """
      Visualize some intermediate results (proposals, raw predictions) inside the pipeline.
      """
      df = get_train_dataflow()   # we don't visualize mask stuff
      df.reset_state()
  
      pred = OfflinePredictor(PredictConfig(
          model=model,
          session_init=get_model_loader(model_path),
          input_names=['image', 'gt_boxes', 'gt_labels'],
          output_names=[
              'generate_{}_proposals/boxes'.format('fpn' if cfg.MODE_FPN else 'rpn'),
              'generate_{}_proposals/probs'.format('fpn' if cfg.MODE_FPN else 'rpn'),
              'fastrcnn_all_probs',
              'final_boxes',
              'final_probs',
              'final_labels',
          ]))
  
      if os.path.isdir(output_dir):
          shutil.rmtree(output_dir)
      utils.fs.mkdir_p(output_dir)
      with tqdm.tqdm(total=nr_visualize) as pbar:
          for idx, dp in itertools.islice(enumerate(df.get_data()), nr_visualize):
              img = dp[0]
              if cfg.MODE_MASK:
                  gt_boxes, gt_labels, gt_masks = dp[-3:]
              else:
                  gt_boxes, gt_labels = dp[-2:]
  
              rpn_boxes, rpn_scores, all_probs, \
                  final_boxes, final_probs, final_labels = pred(img, gt_boxes, gt_labels)
  
              # draw groundtruth boxes
              gt_viz = draw_annotation(img, gt_boxes, gt_labels)
              # draw best proposals for each groundtruth, to show recall
              proposal_viz, good_proposals_ind = draw_proposal_recall(img, rpn_boxes, rpn_scores, gt_boxes)
              # draw the scores for the above proposals
              score_viz = draw_predictions(img, rpn_boxes[good_proposals_ind], all_probs[good_proposals_ind])
  
              results = [DetectionResult(*args) for args in
                         zip(final_boxes, final_probs, final_labels,
                             [None] * len(final_labels))]
              final_viz = draw_final_outputs(img, results)
  
              viz = tpviz.stack_patches([
                  gt_viz, proposal_viz,
                  score_viz, final_viz], 2, 2)
  
              if os.environ.get('DISPLAY', None):
                  tpviz.interactive_imshow(viz)
              cv2.imwrite("{}/{:03d}.png".format(output_dir, idx), viz)
              pbar.update()
  
  
  def offline_evaluate(pred_func, output_file):
      df = get_eval_dataflow()
      all_results = eval_coco(
          df, lambda img: detect_one_image(img, pred_func))
      with open(output_file, 'w') as f:
          json.dump(all_results, f)
      print_evaluation_scores(output_file)
  
  
  def predict(pred_func, input_file):
      img = cv2.imread(input_file, cv2.IMREAD_COLOR)
      results = detect_one_image(img, pred_func)
      final = draw_final_outputs(img, results)
      viz = np.concatenate((img, final), axis=1)
      tpviz.interactive_imshow(viz)
  
  
  class EvalCallback(Callback):
      def __init__(self, in_names, out_names):
          self._in_names, self._out_names = in_names, out_names
  
      def _setup_graph(self):
          self.pred = self.trainer.get_predictor(self._in_names, self._out_names)
          self.df = get_eval_dataflow()
  
      def _before_train(self):
          EVAL_TIMES = 5  # eval 5 times during training
          interval = self.trainer.max_epoch // (EVAL_TIMES + 1)
          self.epochs_to_eval = set([interval * k for k in range(1, EVAL_TIMES + 1)])
          self.epochs_to_eval.add(self.trainer.max_epoch)
          logger.info("[EvalCallback] Will evaluate at epoch " + str(sorted(self.epochs_to_eval)))
  
      def _eval(self):
          all_results = eval_coco(self.df, lambda img: detect_one_image(img, self.pred))
          output_file = os.path.join(
              logger.get_logger_dir(), 'outputs{}.json'.format(self.global_step))
          with open(output_file, 'w') as f:
              json.dump(all_results, f)
          try:
              scores = print_evaluation_scores(output_file)
          except Exception:
              logger.exception("Exception in COCO evaluation.")
              scores = {}
          for k, v in scores.items():
              self.trainer.monitors.put_scalar(k, v)
  
      def _trigger_epoch(self):
          if self.epoch_num in self.epochs_to_eval:
              self._eval()
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--load', help='load a model for evaluation or training. Can overwrite BACKBONE.WEIGHTS')
      parser.add_argument('--logdir', help='log directory', default='train_log/maskrcnn')
      parser.add_argument('--visualize', action='store_true', help='visualize intermediate results')
      parser.add_argument('--evaluate', help="Run evaluation on COCO. "
                                             "This argument is the path to the output json evaluation file")
      parser.add_argument('--predict', help="Run prediction on a given image. "
                                            "This argument is the path to the input image file")
      parser.add_argument('--config', help="A list of KEY=VALUE to overwrite those defined in config.py",
                          nargs='+')
  
      if get_tf_version_tuple() < (1, 6):
          # https://github.com/tensorflow/tensorflow/issues/14657
          logger.warn("TF<1.6 has a bug which may lead to crash in FasterRCNN training if you're unlucky.")
  
      args = parser.parse_args()
      if args.config:
          cfg.update_args(args.config)
  
      MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model()
  
      if args.visualize or args.evaluate or args.predict:
          assert args.load
          finalize_configs(is_training=False)
  
          if args.predict or args.visualize:
              cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS
  
          if args.visualize:
              visualize(MODEL, args.load)
          else:
              pred = OfflinePredictor(PredictConfig(
                  model=MODEL,
                  session_init=get_model_loader(args.load),
                  input_names=MODEL.get_inference_tensor_names()[0],
                  output_names=MODEL.get_inference_tensor_names()[1]))
              if args.evaluate:
                  assert args.evaluate.endswith('.json'), args.evaluate
                  offline_evaluate(pred, args.evaluate)
              elif args.predict:
                  COCODetection(cfg.DATA.BASEDIR, 'val2014')   # Only to load the class names into caches
                  predict(pred, args.predict)
      else:
          is_horovod = cfg.TRAINER == 'horovod'
          if is_horovod:
              hvd.init()
              logger.info("Horovod Rank={}, Size={}".format(hvd.rank(), hvd.size()))
  
          if not is_horovod or hvd.rank() == 0:
              logger.set_logger_dir(args.logdir, 'd')
  
          finalize_configs(is_training=True)
          factor = 8. / cfg.TRAIN.NUM_GPUS
          stepnum = cfg.TRAIN.STEPS_PER_EPOCH
  
          # warmup is step based, lr is epoch based
          warmup_schedule = [(0, cfg.TRAIN.BASE_LR / 3), (cfg.TRAIN.WARMUP * factor, cfg.TRAIN.BASE_LR)]
          warmup_end_epoch = cfg.TRAIN.WARMUP * factor * 1. / stepnum
          lr_schedule = [(int(np.ceil(warmup_end_epoch)), warmup_schedule[-1][1])]
          for idx, steps in enumerate(cfg.TRAIN.LR_SCHEDULE[:-1]):
              mult = 0.1 ** (idx + 1)
              lr_schedule.append(
                  (steps * factor // stepnum, cfg.TRAIN.BASE_LR * mult))
          logger.info("Warm Up Schedule (steps, value): " + str(warmup_schedule))
          logger.info("LR Schedule (epochs, value): " + str(lr_schedule))
  
          callbacks = [
              PeriodicCallback(
                  ModelSaver(max_to_keep=10, keep_checkpoint_every_n_hours=1),
                  every_k_epochs=20),
              # linear warmup
              ScheduledHyperParamSetter(
                  'learning_rate', warmup_schedule, interp='linear', step_based=True),
              ScheduledHyperParamSetter('learning_rate', lr_schedule),
              EvalCallback(*MODEL.get_inference_tensor_names()),
              PeakMemoryTracker(),
              EstimatedTimeLeft(median=True),
              SessionRunTimeout(60000).set_chief_only(True),   # 1 minute timeout
          ]
          if not is_horovod:
              callbacks.append(GPUUtilizationTracker())
  
          if args.load:
              session_init = get_model_loader(args.load)
          else:
              session_init = get_model_loader(cfg.BACKBONE.WEIGHTS) if cfg.BACKBONE.WEIGHTS else None
  
          traincfg = TrainConfig(
              model=MODEL,
              data=QueueInput(get_train_dataflow()),
              callbacks=callbacks,
              steps_per_epoch=stepnum,
              max_epoch=cfg.TRAIN.LR_SCHEDULE[-1] * factor // stepnum,
              session_init=session_init,
          )
          if is_horovod:
              # horovod mode has the best speed for this model
              trainer = HorovodTrainer(average=False)
          else:
              # nccl mode has better speed than cpu mode
              trainer = SyncMultiGPUTrainerReplicated(cfg.TRAIN.NUM_GPUS, average=False, mode='nccl')
          launch_train_with_config(traincfg, trainer)
  