def im_detect(sess, net, im, timers, im_idx=None, nbr_gts=None):
# Setup image blob
blobs = {}
blobs['data'], im_scales, blobs['im_shape_orig'] = get_image_blob(im)
im_blob = blobs['data']
blobs['im_info'] = np.array([im_blob.shape[1], im_blob.shape[2], im_scales[0]])
# Run drl-RPN
scores, pred_bboxes, timers, stats \
= run_drl_rpn(sess, net, blobs, timers, 'test', cfg.DRL_RPN_TEST.BETA,
im_idx, nbr_gts)
return scores, pred_bboxes, timers, stats
def train_model(self, sess, max_iters):
# Build data layers for both training and validation set
self.data_layer = RoIDataLayer(self.roidb, cfg.NBR_CLASSES)
self.data_layer_val = RoIDataLayer(self.valroidb, cfg.NBR_CLASSES,
True)
# Construct the computation graph corresponding to the original Faster R-CNN
# architecture first (and potentially the post-hist module of drl-RPN)
lr_det_op, train_op, lr_post_op, train_op_post \
= self.construct_graph(sess)
# Initialize the variables or restore them from the last snapshot
rate, last_snapshot_iter, stepsizes, np_paths, ss_paths \
= self.initialize(sess)
# We will handle the snapshots ourselves
self.saver = tf.train.Saver(max_to_keep=100000)
# Initialize
self.net.init_rl_train(sess)
# Setup initial learning rates
lr_rl = cfg.DRL_RPN_TRAIN.LEARNING_RATE
lr_det = cfg.TRAIN.LEARNING_RATE
sess.run(tf.assign(lr_det_op, lr_det))
if cfg.DRL_RPN.USE_POST:
lr_post = cfg.DRL_RPN_TRAIN.POST_LR
sess.run(tf.assign(lr_post_op, lr_post))
# Sample first beta
if cfg.DRL_RPN_TRAIN.USE_POST:
betas = cfg.DRL_RPN_TRAIN.POST_BETAS
else:
betas = cfg.DRL_RPN_TRAIN.BETAS
beta_idx = 0
beta = betas[beta_idx]
# Setup drl-RPN timers
timers = {
'init': Timer(),
'fulltraj': Timer(),
'upd-obs-vol': Timer(),
'upd-seq': Timer(),
'upd-rl': Timer(),
'action-rl': Timer(),
'coll-traj': Timer(),
'run-drl-rpn': Timer(),
'train-drl-rpn': Timer()
}
# Create StatCollector (tracks various RL training statistics)
stat_strings = [
'reward', 'rew-done', 'traj-len', 'frac-area', 'gt >= 0.5 frac',
'gt-IoU-frac'
]
sc = StatCollector(max_iters, stat_strings)
timer = Timer()
iter = 0
snapshot_add = 0
while iter < max_iters:
# Get training data, one batch at a time (assumes batch size 1)
blobs = self.data_layer.forward()
# Allows the possibility to start at arbitrary image, rather
# than always starting from first image in dataset. Useful if
# want to load parameters and keep going from there, rather
# than having those and encountering visited images again.
iter, max_iters, snapshot_add, do_continue \
= self._check_if_continue(iter, max_iters, snapshot_add)
if do_continue:
continue
if not cfg.DRL_RPN_TRAIN.USE_POST:
# Potentially update drl-RPN learning rate
if (iter + 1) % cfg.DRL_RPN_TRAIN.STEPSIZE == 0:
lr_rl *= cfg.DRL_RPN_TRAIN.GAMMA
# Run drl-RPN in training mode
timers['run-drl-rpn'].tic()
stats = run_drl_rpn(sess,
self.net,
blobs,
timers,
mode='train',
beta=beta,
im_idx=None,
extra_args=lr_rl)
timers['run-drl-rpn'].toc()
if (iter + 1) % cfg.DRL_RPN_TRAIN.BATCH_SIZE == 0:
print(
"\n##### DRL-RPN BATCH GRADIENT UPDATE - START ##### \n"
)
print('iter: %d / %d' % (iter + 1, max_iters))
print('lr-rl: %f' % lr_rl)
timers['train-drl-rpn'].tic()
self.net.train_drl_rpn(sess, lr_rl, sc, stats)
timers['train-drl-rpn'].toc()
sc.print_stats()
print('TIMINGS:')
print('runnn-drl-rpn: %.4f' %
timers['run-drl-rpn'].get_avg())
print('train-drl-rpn: %.4f' %
timers['train-drl-rpn'].get_avg())
print(
"\n##### DRL-RPN BATCH GRADIENT UPDATE - DONE ###### \n"
)
# Also sample new beta for next batch
beta_idx += 1
beta_idx %= len(betas)
beta = betas[beta_idx]
else:
sc.update(0, stats)
# At this point we assume that an RL-trajectory has been performed.
# We next train detector with drl-RPN running in deterministic mode.
# Potentially train detector component of network
if cfg.DRL_RPN_TRAIN.DET_START >= 0 and \
iter >= cfg.DRL_RPN_TRAIN.DET_START:
# Run drl-RPN in deterministic mode
net_conv, rois_drl_rpn, gt_boxes, im_info, timers, _ \
= run_drl_rpn(sess, self.net, blobs, timers, mode='train_det',
beta=beta, im_idx=None)
# Learning rate
if (iter + 1) % cfg.TRAIN.STEPSIZE[0] == 0:
lr_det *= cfg.TRAIN.GAMMA
sess.run(tf.assign(lr_det_op, lr_det))
timer.tic()
# Train detector part
loss_cls, loss_box, tot_loss \
= self.net.train_step_det(sess, train_op, net_conv, rois_drl_rpn,
gt_boxes, im_info)
timer.toc()
# Display training information
self._print_det_loss(iter, max_iters, tot_loss, loss_cls,
loss_box, lr_det, timer)
# Train post-hist module AFTER we have trained rest of drl-RPN! Specifically
# once rest of drl-RPN has been trained already, copy those weights into
# the folder of pretrained weights and rerun training with those as initial
# weights, which will then train only the posterior-history module
else:
# The very first time we need to assign the ordinary detector weights
# as starting point
if iter == 0:
self.net.assign_post_hist_weights(sess)
# Sample beta
beta = betas[beta_idx]
beta_idx += 1
beta_idx %= len(betas)
# Run drl-RPN in deterministic mode
net_conv, rois_drl_rpn, gt_boxes, im_info, timers, cls_hist \
= run_drl_rpn(sess, self.net, blobs, timers, mode='train_det',
beta=beta, im_idx=None)
# Learning rate (assume only one learning rate iter for now!)
if (iter + 1) % cfg.DRL_RPN_TRAIN.POST_SS[0] == 0:
lr_post *= cfg.TRAIN.GAMMA
sess.run(tf.assign(lr_post_op, lr_post))
# Train post-hist detector part
tot_loss = self.net.train_step_post(sess, train_op_post,
net_conv, rois_drl_rpn,
gt_boxes, im_info,
cls_hist)
# Display training information
self._print_det_loss(iter, max_iters, tot_loss, None, None,
lr_post, timer, 'post-hist')
# Snapshotting
if (iter + 1) % cfg.TRAIN.SNAPSHOT_ITERS == 0:
last_snapshot_iter = iter + 1
ss_path, np_path = self.snapshot(sess, iter + 1 + snapshot_add)
np_paths.append(np_path)
ss_paths.append(ss_path)
# Remove the old snapshots if there are too many
if len(np_paths) > cfg.TRAIN.SNAPSHOT_KEPT:
self.remove_snapshot(np_paths, ss_paths)
# Increase iteration counter
iter += 1
# Potentially save one last time
if last_snapshot_iter != iter:
self.snapshot(sess, iter + snapshot_add)