def resnet_fpn_model_fn(features, labels, mode, params): """parameter definition part""" is_train = (mode == tf.estimator.ModeKeys.TRAIN) resnet_num_blocks = params["RESNET_NUM_BLOCKS"] num_anchors = params["num_anchors"] head_dim = params["head_dim"] resolution = params["resolution"] num_classes = params["num_classes"] bbox_reg_weights = params["bbox_reg_weights"] weight_decay = params["weight_decay"] learning_rate = params["learning_rate"] lr_schedule = params["lr_schedule"] """model definition part""" image = image_preprocess(features["image"]) c2345 = resnet_fpn_backbone(image, resnet_num_blocks, is_train) p23456 = fpn_model(c2345) image_shape2d = tf.shape(image)[1:3] all_anchors_fpn = tf_get_all_anchors_fpn() model_rpn_head = RPNHead(_C.FPN.NUM_CHANNEL, len(_C.RPN.ANCHOR_RATIOS)) rpn_outputs = [model_rpn_head(pi) for pi in p23456] multilevel_label_logits = [k[0] for k in rpn_outputs] multilevel_box_logits = [k[1] for k in rpn_outputs] # debug_op = tf.print({"debug_inf": tf.convert_to_tensor("now in here")}) # with tf.control_dependencies([debug_op]): # image_shape2d = tf.identity(image_shape2d) if mode != tf.estimator.ModeKeys.PREDICT: multilevel_anchors = [ RPNAnchors(all_anchors_fpn[i], features['anchor_labels_lvl{}'.format(i + 2)], features['anchor_boxes_lvl{}'.format(i + 2)]) for i in range(len(all_anchors_fpn)) ] else: multilevel_anchors = [ RPNAnchors(all_anchors_fpn[i], None, None) for i in range(len(all_anchors_fpn)) ] slice_feature_and_anchors(p23456, multilevel_anchors) # Multi-Level RPN Proposals multilevel_pred_boxes = [ anchor.decode_logits(logits) for anchor, logits in zip(multilevel_anchors, multilevel_box_logits) ] proposal_boxes, proposal_scores = generate_fpn_proposals( multilevel_pred_boxes, multilevel_label_logits, image_shape2d, is_train) proposals = BoxProposals(proposal_boxes) gt_boxes = None gt_labels = None if mode != tf.estimator.ModeKeys.PREDICT: losses = multilevel_rpn_losses(multilevel_anchors, multilevel_label_logits, multilevel_box_logits) gt_boxes = features['boxes'] gt_labels = features['gt_labels'] proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels) fastrcnn_head_func = getattr(model_frcnn, _C.FPN.FRCNN_HEAD_FUNC) if not _C.FPN.CASCADE: roi_feature_fastrcnn = multilevel_roi_align(p23456[:4], proposals.boxes, 7) head_feature = fastrcnn_head_func(roi_feature_fastrcnn) fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs( head_feature, num_classes) fastrcnn_head = FastRCNNHead( proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes, tf.convert_to_tensor(bbox_reg_weights, tf.float32)) else: def roi_func(boxes): return multilevel_roi_align(p23456[:4], boxes, 7) fastrcnn_head = CascadeRCNNHead(proposals, roi_func, fastrcnn_head_func, (gt_boxes, gt_labels), image_shape2d, num_classes, mode != tf.estimator.ModeKeys.PREDICT) decoded_boxes = fastrcnn_head.decoded_output_boxes() decoded_boxes = clip_boxes(decoded_boxes, image_shape2d) label_scores = fastrcnn_head.output_scores() # final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores) final_boxes, final_scores, final_labels, valid_detections = fastrcnn_predictions_v2( decoded_boxes, label_scores) global_step = tf.train.get_or_create_global_step() if mode != tf.estimator.ModeKeys.PREDICT: all_losses = fastrcnn_head.losses() trainable_weights = tf.trainable_variables() weight_loss = 0.0 for i, ele in enumerate(trainable_weights): if re.search('.*/kernel', ele.name): weight_loss += tf.reduce_sum(tf.square(ele) * weight_decay) # print_op = tf.print({'rpn_loss': tf.add_n(losses), # 'frcnn_loss': tf.add_n(all_losses)}) # with tf.control_dependencies([print_op]): total_cost = tf.add_n(losses + all_losses, "total_cost") total_cost = tf.add(total_cost, weight_loss, 'all_total_cost') if is_train: learning_rate = tf.train.piecewise_constant( global_step, lr_schedule, values=[tf.convert_to_tensor(0.01 * 0.33 * 0.375, tf.float32)] + [learning_rate * (0.1**i) for i in range(len(lr_schedule))]) tf.summary.scalar("learning_rate", learning_rate) opt = tf.train.MomentumOptimizer(learning_rate, 0.9) # opt = tf.train.AdamOptimizer(learning_rate) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) print(update_ops) with tf.control_dependencies(update_ops): train_op = opt.minimize(total_cost, global_step) return tf.estimator.EstimatorSpec(mode, loss=total_cost, train_op=train_op) else: return tf.estimator.EstimatorSpec(mode, loss=total_cost) else: predictions = { 'boxes': final_boxes[0, :valid_detections[0]], 'labels': final_labels[0, :valid_detections[0]], 'scores': final_scores[0, :valid_detections[0]], # 'image': features['image'], # 'original_image': features['original_image'], 'h_pre': features['h_pre'], 'w_pre': features['w_pre'], 'h_now': features['h_now'], 'w_now': features['w_now'], 'scale': features['scale'], 'image_id': features.get('image_id', tf.convert_to_tensor(0)), 'valid_detection': valid_detections } # predictions = {'boxes': final_boxes, # 'labels': final_labels, # 'scores': final_scores, # 'image': features['image'], # 'original_image': features['original_image'], # 'h_pre': features['h_pre'], # 'w_pre': features['w_pre'], # 'h_now': features['h_now'], # 'w_now': features['w_now'], # 'scale': features['scale'], # 'image_id': features.get('image_id', tf.convert_to_tensor(0)) # } return tf.estimator.EstimatorSpec(mode, predictions)
def resnet_c4_model_fn(features, labels, mode, params): """parameter defination part""" is_train = (mode == tf.estimator.ModeKeys.TRAIN) resnet_num_blocks = params["RESNET_NUM_BLOCKS"] num_anchors = params["num_anchors"] head_dim = params["head_dim"] resolution = params["resolution"] num_classes = params["num_classes"] bbox_reg_weights = params["bbox_reg_weights"] weight_decay = params["weight_decay"] learning_rate = params["learning_rate"] lr_schedule = params["lr_schedule"] """model definition part""" image = image_preprocess(features['image']) featuremap = resnet_c4_backbone(image, resnet_num_blocks[:3], is_train) image_shape2d = tf.shape(image)[1:3] rpn_label_logits, rpn_box_logits = rpn_head(featuremap, head_dim, num_anchors) if mode != tf.estimator.ModeKeys.PREDICT: anchors = RPNAnchors(tf_get_all_anchors(), features['anchor_labels'], features['anchor_boxes']) else: anchors = RPNAnchors(tf_get_all_anchors(), None, None) anchors = anchors.narrow_to(featuremap) pred_boxes_decoded = anchors.decode_logits(rpn_box_logits) # x1y1x2y2 proposals, proposal_scores = generate_rpn_proposals( tf.reshape(pred_boxes_decoded, [-1, 4]), tf.reshape(rpn_label_logits, [-1]), image_shape2d, _C.RPN.TRAIN_PRE_NMS_TOPK if mode == tf.estimator.ModeKeys.TRAIN else _C.RPN.TEST_PRE_NMS_TOPK, _C.RPN.TRAIN_POST_NMS_TOPK if mode == tf.estimator.ModeKeys.TRAIN else _C.RPN.TEST_POST_NMS_TOPK) # rpn_size = tf.shape(proposals)[0] # rpn_boxes = tf.gather(proposals, tf.where(tf.greater(proposals, 0.5))) proposals = BoxProposals(proposals) if mode != tf.estimator.ModeKeys.PREDICT: rpn_loss = rpn_losses(anchors.gt_labels, anchors.encoded_gt_boxes(), rpn_label_logits, rpn_box_logits) # targets = [features[k] for k in ['boxes', 'gt_labels', 'gt_masks'] if k in features.keys()] # gt_boxes, gt_labels, *_ = targets gt_boxes = features['boxes'] gt_labels = features['gt_labels'] proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels) boxes_on_featuremap = proposals.boxes * (1.0 / _C.RPN.ANCHOR_STRIDE) roi_resized = roi_align(featuremap, boxes_on_featuremap, resolution) feature_fastrcnn = resnet_conv5(roi_resized, resnet_num_blocks[-1], is_train) # Keep C5 feature to be shared with mask branch feature_gap = keras.layers.GlobalAveragePooling2D()(feature_fastrcnn) fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs( feature_gap, num_classes) bbox_reg_weights_tensor = tf.convert_to_tensor(bbox_reg_weights, tf.float32) if mode != tf.estimator.ModeKeys.PREDICT: fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes, bbox_reg_weights_tensor) all_loss = fastrcnn_head.losses() label_scores = tf.nn.softmax(fastrcnn_label_logits) decoded_boxes = decoded_output_boxes(proposals, num_classes, fastrcnn_box_logits, bbox_reg_weights_tensor) decoded_boxes = clip_boxes(decoded_boxes, image_shape2d) final_boxes, final_scores, final_labels, valid_detections = fastrcnn_predictions_v2( decoded_boxes, label_scores) # final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores) global_step = tf.train.get_or_create_global_step() if mode != tf.estimator.ModeKeys.PREDICT: trainable_weights = tf.trainable_variables() weight_loss = 0.0 for i, ele in enumerate(trainable_weights): if re.search('.*/kernel', ele.name): weight_loss += tf.reduce_sum(tf.square(ele) * weight_decay) total_cost = tf.add_n(rpn_loss + all_loss, 'total_cost') tf.summary.scalar('total_cost', total_cost) if is_train: learning_rate = tf.train.piecewise_constant( global_step, lr_schedule, values=[tf.convert_to_tensor(0.01 * 0.33, tf.float32)] + [learning_rate * (0.1**i) for i in range(len(lr_schedule))]) opt = tf.train.MomentumOptimizer(learning_rate, 0.9) # opt = tf.train.AdamOptimizer(learning_rate) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = opt.minimize(total_cost, global_step) return tf.estimator.EstimatorSpec(mode, loss=total_cost, train_op=train_op) else: return tf.estimator.EstimatorSpec(mode, loss=total_cost) else: predictions = { 'boxes': final_boxes[0, :valid_detections[0]], 'labels': final_labels[0, :valid_detections[0]], 'scores': final_scores[0, :valid_detections[0]], 'image': features['image'], # 'rpn_boxes': rpn_boxes, # 'rpn_size': rpn_size, 'valid_detection': valid_detections } return tf.estimator.EstimatorSpec(mode, predictions)
tf.reshape(pred_boxes_decoded, [-1, 4]), tf.reshape(rpn_label_logits, [-1]), image_shape2d, _C.RPN.TRAIN_PRE_NMS_TOPK if is_train else _C.RPN.TEST_PRE_NMS_TOPK, _C.RPN.TRAIN_POST_NMS_TOPK if is_train else _C.RPN.TEST_POST_NMS_TOPK) proposals = BoxProposals(proposals) if is_train: rpn_loss = rpn_losses(anchors.gt_labels, anchors.encoded_gt_boxes(), rpn_label_logits, rpn_box_logits) # targets = [features[k] for k in ['boxes', 'gt_labels', 'gt_masks'] if k in features.keys()] # gt_boxes, gt_labels, *_ = targets gt_boxes = features['boxes'] gt_labels = features['gt_labels'] proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels) boxes_on_featuremap = proposals.boxes * (1.0 / _C.RPN.ANCHOR_STRIDE) roi_resized = roi_align(featuremap, boxes_on_featuremap, resolution) feature_fastrcnn = resnet_conv5(roi_resized, resnet_num_blocks[-1], is_train) # Keep C5 feature to be shared with mask branch feature_gap = keras.layers.GlobalAveragePooling2D()(feature_fastrcnn) fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(feature_gap, num_classes) bbox_reg_weights_tensor = tf.convert_to_tensor(bbox_reg_weights, tf.float32) if is_train: fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes, bbox_reg_weights_tensor) all_loss = fastrcnn_head.losses() label_scores = tf.nn.softmax(fastrcnn_label_logits) decoded_boxes = decoded_output_boxes(proposals, num_classes, fastrcnn_box_logits, bbox_reg_weights_tensor) decoded_boxes = clip_boxes(decoded_boxes, image_shape2d) final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores) print("final_boxes", final_boxes) print("final_scores", final_scores) print("final_labels", final_labels)