def train(rpn_data, roi_data, rpn_cls_score_reshape, rpn_bbox_pred, feature_map, cls_score, bbox_pred, input_image_tensor, input_gt_box_tensor, input_im_info_tensor, pretrain_model=None): output_dir = '/home/give/PycharmProjects/MyFasterRCNN/parameters' saver = tf.train.Saver(max_to_keep=5) # RPN # classification loss(fg bg) rpn_cls_score = tf.reshape(rpn_cls_score_reshape, [-1, 2]) rpn_label = tf.reshape(rpn_data[0], [-1]) rpn_cls_score = tf.reshape( tf.gather(rpn_cls_score, tf.where(tf.not_equal(rpn_label, -1))), [-1, 2]) rpn_label = tf.reshape( tf.gather(rpn_label, tf.where(tf.not_equal(rpn_label, -1))), [-1]) rpn_cross_entropy = tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits(logits=rpn_cls_score, labels=rpn_label)) # bounding box regression L1 loss rpn_bbox_pred = rpn_bbox_pred rpn_bbox_targets = tf.transpose(rpn_data[1], [0, 2, 3, 1]) rpn_bbox_inside_weights = tf.transpose(rpn_data[2], [0, 2, 3, 1]) rpn_bbox_outside_weights = tf.transpose(rpn_data[3], [0, 2, 3, 1]) rpn_smooth_l1 = _modified_smooth_l1(3.0, rpn_bbox_pred, rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights) rpn_loss_box = tf.reduce_mean( tf.reduce_sum(rpn_smooth_l1, reduction_indices=[1, 2, 3])) # classification loss label = tf.reshape(roi_data[1], [-1]) cross_entropy = tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits(logits=cls_score, labels=label)) smooth_l1 = _modified_smooth_l1(1.0, bbox_pred, roi_data[2], roi_data[3], roi_data[4]) loss_box = tf.reduce_mean(tf.reduce_sum(smooth_l1, reduction_indices=[1])) loss = cross_entropy + loss_box + rpn_cross_entropy + rpn_loss_box global_step = tf.Variable(0, trainable=False) lr = tf.train.exponential_decay(cfg.TRAIN.LEARNING_RATE, global_step, cfg.TRAIN.STEPSIZE, 0.001, staircase=True) momentum = cfg.TRAIN.MOMENTUM train_op = tf.train.MomentumOptimizer(lr, momentum).minimize( loss, global_step=global_step) with tf.Session() as sess: args = parse_args() from lib.datasets.factory import get_imdb from lib.fast_rcnn.train import get_training_roidb, filter_roidb from lib.fast_rcnn.train import get_data_layer import lib.roi_data_layer.roidb as rdl_roidb import numpy as np from lib.fast_rcnn.config import cfg_from_file, cfg_from_list if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) imdb = get_imdb('voc_2007_trainval') roidb = get_training_roidb(imdb) roidb = filter_roidb(roidb) bbox_means, bbox_stds = rdl_roidb.add_bbox_regression_targets(roidb) data_layer = get_data_layer(roidb, imdb.num_classes) init_op = tf.global_variables_initializer() sess.run(init_op) if pretrain_model is not None: print('Loading pretrained model ' 'weights from {:s}').format(pretrain_model) load(pretrain_model, sess, saver, True) for iter_index in range(args.max_iters): blobs = data_layer.forward() # blobs keys include: data, gt_boxesm im_info _, rpn_bbox_targets_value, rpn_bbox_inside_weights_value, rpn_bbox_outside_weights_value = sess.run( [ train_op, rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights ], feed_dict={ input_image_tensor: blobs['data'], input_gt_box_tensor: blobs['gt_boxes'], input_im_info_tensor: blobs['im_info'] }) rpn_label_value, loss_value, loss_box_value, cross_entropy_value, rpn_cross_entropy_value, rpn_loss_box_value = sess.run( [ rpn_label, loss, loss_box, cross_entropy, rpn_cross_entropy, rpn_loss_box ], feed_dict={ input_image_tensor: blobs['data'], input_gt_box_tensor: blobs['gt_boxes'], input_im_info_tensor: blobs['im_info'] }) if iter_index % 100 == 0: print 'iter: %d / %d' % (iter_index, args.max_iters) print 'total loss: %.4f, rpn cross entropy: %.4f, rpn_loss_box: %.4f, cross entroopy: %.4f, loss box: %.4f' % ( loss_value, rpn_cross_entropy_value, rpn_loss_box_value, cross_entropy_value, loss_box_value) # print np.shape(rpn_label_value), rpn_label_value if iter_index == 0: continue infix = ('_' + cfg.TRAIN.SNAPSHOT_INFIX if cfg.TRAIN.SNAPSHOT_INFIX != '' else '') filename = (cfg.TRAIN.SNAPSHOT_PREFIX + infix + '_iter_{:d}'.format(iter_index + 1) + '.ckpt') filename = os.path.join(output_dir, filename) saver.save(sess, filename)
sys.exit(1) args = parser.parse_args() return args if __name__ == '__main__': args = parse_args() print('Called with args:') print(args) if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) cfg.GPU_ID = args.gpu_id print('Using config:') pprint.pprint(cfg) if not args.randomize: # fix the random seeds (numpy and caffe) for reproducibility np.random.seed(cfg.RNG_SEED) #imdb = get_imdb(args.imdb_name) #print 'Loaded dataset `{:s}` for training'.format(imdb.name) #roidb = get_training_roidb(imdb) ''' output_dir = get_output_dir(imdb, None) log_dir = get_log_dir(imdb) print 'Output will be saved to `{:s}`'.format(output_dir)
parser.print_help() # sys.exit(1) args = parser.parse_args() return args if __name__ == '__main__': args = parse_args() print('Called with args:') print(args) if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) print('Using config:') pprint.pprint(cfg) if not args.randomize: # fix the random seeds (numpy and caffe) for reproducibility np.random.seed(cfg.RNG_SEED) imdb = get_imdb(args.imdb_name) print 'Loaded dataset `{:s}` for training'.format(imdb.name) roidb = get_training_roidb(imdb) output_dir = get_output_dir(imdb, None) log_dir = get_log_dir(imdb) print 'Output will be saved to `{:s}`'.format(output_dir) print 'Logs will be saved to `{:s}`'.format(log_dir)