def evaluate():
"""Evaluate."""
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
if not os.path.exists(FLAGS.eval_dir):
os.makedirs(FLAGS.eval_dir)
with tf.Graph().as_default() as g:
#Select model to evaluate
if FLAGS.net == 'squeezeDet+PruneFilter':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.IS_TRAINING = False
mc.IS_PRUNING = False
mc.LOAD_PRETRAINED_MODEL = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilter(mc)
elif FLAGS.net == 'squeezeDet+PruneFilterShape':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.IS_TRAINING = False
mc.IS_PRUNING = False
mc.LOAD_PRETRAINED_MODEL = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilterShape(mc)
elif FLAGS.net == 'squeezeDet+PruneLayer':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.IS_TRAINING = False
mc.IS_PRUNING = False
mc.LOAD_PRETRAINED_MODEL = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneLayer(mc)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
no_evaluation = True
# Logic to load checkpoints if available
for file in os.listdir(FLAGS.checkpoint_path):
if file.endswith(".meta"):
saver = tf.train.Saver(model.model_params)
ckpt_path = FLAGS.checkpoint_path + "/" + file[:-5]
step = file[11:-5]
eval_once(saver, ckpt_path, imdb, model, step, True)
no_evaluation = False
if no_evaluation:
saver = None
eval_once(saver, '-', imdb, model, '0', False)
def evaluate():
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default() as g:
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, state='val')
output_node_names = ["out_det_probs", "out_det_boxes", "out_det_class"]
det_probs = tf.identity(model.det_probs, name=output_node_names[0])
det_boxes = tf.identity(model.det_boxes, name=output_node_names[1])
det_class = tf.identity(model.det_class, name=output_node_names[2])
imdb = kitti('val', FLAGS.data_path, mc)
saver = tf.train.Saver(model.model_params)
eval_once(saver, FLAGS.checkpoint_path, imdb, model, output_node_names)
def evaluate():
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default() as g:
mc = kitti_squeezeDetPlus_config()
#mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, state='val')
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# add summary ops and placeholders
saver = tf.train.Saver(model.model_params)
summary_writer = tf.summary.FileWriter(FLAGS.eval_dir, g)
ckpts = set()
while True:
if FLAGS.run_once:
# When run_once is true, checkpoint_path should point to the exact
# checkpoint file.
eval_once(saver, FLAGS.checkpoint_path, summary_writer, imdb,
model)
return
else:
# When run_once is false, checkpoint_path should point to the directory
# that stores checkpoint files.
from os.path import basename
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
#base_name, max_idx = ckpt.model_checkpoint_path.split('-')
#for idx in range(0, int(max_idx), )
if ckpt and ckpt.model_checkpoint_path:
if ckpt.model_checkpoint_path in ckpts:
# Do not evaluate on the same checkpoint
print(
'Wait {:d}s for new checkpoints to be saved ... '.
format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
else:
ckpts.add(ckpt.model_checkpoint_path)
print('Evaluating {}...'.format(
ckpt.model_checkpoint_path))
eval_once(saver, ckpt.model_checkpoint_path,
summary_writer, imdb, model)
else:
print('No checkpoint file found')
if not FLAGS.run_once:
print(
'Wait {:d}s for new checkpoints to be saved ... '.
format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
def evaluate():
"""Evaluate."""
assert FLAGS.net in ['vgg16', 'vgg16_v2', 'vgg16_v3', 'yolo_v2'], \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
assert FLAGS.dataset in ['KITTI', 'PASCAL_VOC', 'VID'], \
'Either KITTI / PASCAL_VOC / VID'
if FLAGS.dataset == 'KITTI':
if FLAGS.net == 'yolo_v2':
raise NotImplementedError
else:
mc = kitti_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
imdb = kitti(FLAGS.train_set, FLAGS.data_path, mc)
elif FLAGS.dataset == 'PASCAL_VOC':
if FLAGS.net == 'yolo_v2':
mc = pascal_voc_yolo_config()
else:
mc = pascal_voc_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
imdb = pascal_voc(FLAGS.train_set, FLAGS.year, FLAGS.data_path, mc)
elif FLAGS.dataset == 'VID':
if FLAGS.net == 'yolo_v2':
mc = vid_yolo_config()
else:
mc = vid_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
imdb = vid(FLAGS.train_set, FLAGS.data_path, mc)
with tf.Graph().as_default() as g:
if FLAGS.dataset == 'KITTI':
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
elif FLAGS.dataset == 'PASCAL_VOC':
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = True
imdb = pascal_voc(FLAGS.image_set, FLAGS.year, FLAGS.data_path, mc)
elif FLAGS.dataset == 'VID':
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
imdb = vid(FLAGS.image_set, FLAGS.data_path, mc)
if FLAGS.net == 'vgg16':
model = VGG16ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'vgg16_v2':
model = VGG16ConvDetV2(mc, FLAGS.gpu)
elif FLAGS.net == 'vgg16_v3':
model = VGG16ConvDetV3(mc, FLAGS.gpu)
elif FLAGS.net == 'yolo_v2':
model = YOLO_V2(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
summary_writer = tf.train.SummaryWriter(FLAGS.eval_dir, g)
ckpts = set()
while True:
if FLAGS.run_once:
# When run_once is true, checkpoint_path should point to the exact
# checkpoint file.
eval_once(saver, FLAGS.checkpoint_path, summary_writer, imdb, model)
return
else:
# When run_once is false, checkpoint_path should point to the directory
# that stores checkpoint files.
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
if ckpt and ckpt.model_checkpoint_path:
if ckpt.model_checkpoint_path in ckpts:
# Do not evaluate on the same checkpoint
print ('Wait {:d}s for new checkpoints to be saved ... '
.format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
else:
ckpts.add(ckpt.model_checkpoint_path)
print ('Evaluating {}...'.format(ckpt.model_checkpoint_path))
eval_once(saver, ckpt.model_checkpoint_path,
summary_writer, imdb, model)
else:
print('No checkpoint file found')
if not FLAGS.run_once:
print ('Wait {:d}s for new checkpoints to be saved ... '
.format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
def evaluate():
"""Evaluate."""
assert FLAGS.dataset == 'KITTI', \
'Currently only supports KITTI dataset'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default() as g:
assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
mc = kitti_vgg16_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = VGG16ConvDet(mc)
elif FLAGS.net == 'resnet50':
mc = kitti_res50_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = ResNet50ConvDet(mc)
elif FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc)
elif FLAGS.net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# add summary ops and placeholders
ap_names = []
for cls in imdb.classes:
ap_names.append(cls+'_easy')
ap_names.append(cls+'_medium')
ap_names.append(cls+'_hard')
eval_summary_ops = []
eval_summary_phs = {}
for ap_name in ap_names:
ph = tf.placeholder(tf.float32)
eval_summary_phs['APs/'+ap_name] = ph
eval_summary_ops.append(tf.summary.scalar('APs/'+ap_name, ph))
ph = tf.placeholder(tf.float32)
eval_summary_phs['APs/mAP'] = ph
eval_summary_ops.append(tf.summary.scalar('APs/mAP', ph))
ph = tf.placeholder(tf.float32)
eval_summary_phs['timing/im_detect'] = ph
eval_summary_ops.append(tf.summary.scalar('timing/im_detect', ph))
ph = tf.placeholder(tf.float32)
eval_summary_phs['timing/im_read'] = ph
eval_summary_ops.append(tf.summary.scalar('timing/im_read', ph))
ph = tf.placeholder(tf.float32)
eval_summary_phs['timing/post_proc'] = ph
eval_summary_ops.append(tf.summary.scalar('timing/post_proc', ph))
ph = tf.placeholder(tf.float32)
eval_summary_phs['num_det_per_image'] = ph
eval_summary_ops.append(tf.summary.scalar('num_det_per_image', ph))
saver = tf.train.Saver(model.model_params)
summary_writer = tf.summary.FileWriter(FLAGS.eval_dir, g)
ckpts = set()
while True:
if FLAGS.run_once:
# When run_once is true, checkpoint_path should point to the exact
# checkpoint file.
eval_once(
saver, FLAGS.checkpoint_path, summary_writer, eval_summary_ops,
eval_summary_phs, imdb, model)
return
else:
# When run_once is false, checkpoint_path should point to the directory
# that stores checkpoint files.
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
if ckpt and ckpt.model_checkpoint_path:
if ckpt.model_checkpoint_path in ckpts:
# Do not evaluate on the same checkpoint
print ('Wait {:d}s for new checkpoints to be saved ... '
.format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
else:
ckpts.add(ckpt.model_checkpoint_path)
print ('Evaluating {}...'.format(ckpt.model_checkpoint_path))
eval_once(
saver, ckpt.model_checkpoint_path, summary_writer,
eval_summary_ops, eval_summary_phs, imdb, model)
else:
print('No checkpoint file found')
if not FLAGS.run_once:
print ('Wait {:d}s for new checkpoints to be saved ... '
.format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
def train():
"""Train SqueezeDet model"""
assert FLAGS.dataset == 'KITTI', \
'Currently only support KITTI dataset'
with tf.Graph().as_default():
assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
mc = kitti_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = VGG16ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'resnet50':
mc = kitti_res50_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = ResNet50ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlus(mc, FLAGS.gpu)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'),
'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
saver = tf.train.Saver(tf.all_variables())
summary_op = tf.merge_all_summaries()
init = tf.initialize_all_variables()
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(init)
tf.train.start_queue_runners(sess=sess)
summary_writer = tf.train.SummaryWriter(FLAGS.train_dir, sess.graph)
for step in xrange(FLAGS.max_steps):
start_time = time.time()
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(
label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend([[i, aidx_per_batch[i][j], k]
for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
else:
num_discarded_labels += 1
if mc.DEBUG_MODE:
print(
'Warning: Discarded {}/({}) labels that are assigned to the same'
'anchor'.format(num_discarded_labels, num_labels))
feed_dict = {
model.image_input:
image_per_batch,
model.keep_prob:
mc.KEEP_PROB,
model.input_mask:
np.reshape(
sparse_to_dense(mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0] * len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
model.box_delta_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
model.box_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
model.labels:
sparse_to_dense(label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0] * len(label_indices)),
}
if step % FLAGS.summary_step == 0:
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, conf_loss, \
bbox_loss, class_loss = sess.run(op_list, feed_dict=feed_dict)
_viz_prediction_result(model, image_per_batch, bbox_per_batch,
label_per_batch, det_boxes, det_class,
det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op,
feed_dict={model.image_to_show: image_per_batch})
num_discarded_labels_op = tf.scalar_summary(
'counter/num_discarded_labels', num_discarded_labels)
num_labels_op = tf.scalar_summary('counter/num_labels',
num_labels)
counter_summary_str = sess.run(
[num_discarded_labels_op, num_labels_op])
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
for sum_str in counter_summary_str:
summary_writer.add_summary(sum_str, step)
print('conf_loss: {}, bbox_loss: {}, class_loss: {}'.format(
conf_loss, bbox_loss, class_loss))
else:
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.class_loss
],
feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = (
'%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print(format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step +
1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
def train():
"""Train SqueezeDet model"""
assert FLAGS.dataset == 'KITTI', \
'Currently only support KITTI dataset'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default():
assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
mc = kitti_vgg16_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = VGG16ConvDet(mc)
elif FLAGS.net == 'resnet50':
mc = kitti_res50_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = ResNet50ConvDet(mc)
elif FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDet(mc)
elif FLAGS.net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlus(mc)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'), 'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print ('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
def _load_data(load_to_placeholder=True):
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend(
[[i, aidx_per_batch[i][j], k] for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
else:
num_discarded_labels += 1
if mc.DEBUG_MODE:
print ('Warning: Discarded {}/({}) labels that are assigned to the same '
'anchor'.format(num_discarded_labels, num_labels))
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
feed_dict = {
image_input: image_per_batch,
input_mask: np.reshape(
sparse_to_dense(
mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0]*len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
box_delta_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
box_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
labels: sparse_to_dense(
label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0]*len(label_indices)),
}
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _ = _load_data()
sess.run(model.enqueue_op, feed_dict=feed_dict)
if mc.DEBUG_MODE:
print ("added to the queue")
if mc.DEBUG_MODE:
print ("Finished enqueue")
except Exception, e:
coord.request_stop(e)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver(tf.global_variables())
summary_op = tf.summary.merge_all()
ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
init = tf.global_variables_initializer()
sess.run(init)
coord = tf.train.Coordinator()
if mc.NUM_THREAD > 0:
enq_threads = []
for _ in range(mc.NUM_THREAD):
enq_thread = threading.Thread(target=_enqueue, args=[sess, coord])
# enq_thread.isDaemon()
enq_thread.start()
enq_threads.append(enq_thread)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
run_options = tf.RunOptions(timeout_in_ms=60000)
# try:
for step in xrange(FLAGS.max_steps):
if coord.should_stop():
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
break
start_time = time.time()
if step % FLAGS.summary_step == 0:
feed_dict, image_per_batch, label_per_batch, bbox_per_batch = \
_load_data(load_to_placeholder=False)
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, \
conf_loss, bbox_loss, class_loss = sess.run(
op_list, feed_dict=feed_dict)
_viz_prediction_result(
model, image_per_batch, bbox_per_batch, label_per_batch, det_boxes,
det_class, det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op, feed_dict={model.image_to_show: image_per_batch})
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
summary_writer.flush()
print ('conf_loss: {}, bbox_loss: {}, class_loss: {}'.
format(conf_loss, bbox_loss, class_loss))
else:
if mc.NUM_THREAD > 0:
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], options=run_options)
else:
feed_dict, _, _, _ = _load_data(load_to_placeholder=False)
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = ('%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print (format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step + 1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
"""Can be train, trainval, val, or test""")
tf.app.flags.DEFINE_string('pretrained_model_path',
'../path/to/pretrained/weights.pkl',
"""Directory where to write event logs """)
tf.app.flags.DEFINE_string('data_path', '../data/KITTI',
"""Root directory of data""")
with tf.Graph().as_default() as g:
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = True
print(FLAGS.pretrained_model_path)
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# add summary ops and placeholders
ap_names = []
for cls in imdb.classes:
ap_names.append(cls + '_easy')
ap_names.append(cls + '_medium')
ap_names.append(cls + '_hard')
eval_summary_ops = []
eval_summary_phs = {}
for ap_name in ap_names:
ph = tf.placeholder(tf.float32)
eval_summary_phs['APs/' + ap_name] = ph
eval_summary_ops.append(tf.summary.scalar('APs/' + ap_name, ph))
def evaluate():
"""Evaluate."""
assert FLAGS.dataset == 'KITTI', \
'Currently only supports KITTI dataset'
with tf.Graph().as_default() as g:
assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
mc = kitti_vgg16_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = VGG16ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'resnet50':
mc = kitti_res50_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = ResNet50ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
saver = tf.train.Saver(model.model_params)
summary_writer = tf.summary.FileWriter(FLAGS.eval_dir, g)
ckpts = set()
while True:
if FLAGS.run_once:
# When run_once is true, checkpoint_path should point to the exact
# checkpoint file.
eval_once(saver, FLAGS.checkpoint_path, summary_writer, imdb,
model)
return
else:
# When run_once is false, checkpoint_path should point to the directory
# that stores checkpoint files.
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
if ckpt and ckpt.model_checkpoint_path:
if ckpt.model_checkpoint_path in ckpts:
# Do not evaluate on the same checkpoint
print(
'Wait {:d}s for new checkpoints to be saved ... '.
format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
else:
ckpts.add(ckpt.model_checkpoint_path)
print('Evaluating {}...'.format(
ckpt.model_checkpoint_path))
eval_once(saver, ckpt.model_checkpoint_path,
summary_writer, imdb, model)
else:
print('No checkpoint file found')
if not FLAGS.run_once:
print(
'Wait {:d}s for new checkpoints to be saved ... '.
format(FLAGS.eval_interval_secs))
time.sleep(FLAGS.eval_interval_secs)
def test_read_batch():
"""Test read batch function"""
assert FLAGS.dataset in ['KITTI', 'PASCAL_VOC', 'VID', 'ILSVRC2013'], \
"""
Invalid dataset {}
Either KITTI / PASCAL_VOC / VID / ILSVRC2013""".format(FLAGS.dataset)
if FLAGS.dataset == 'KITTI':
mc = kitti_vgg16_config()
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
elif FLAGS.dataset == 'PASCAL_VOC':
mc = pascal_voc_vgg16_config()
imdb = pascal_voc(FLAGS.image_set, FLAGS.year, FLAGS.data_path, mc)
elif FLAGS.dataset == 'VID':
mc = vid_vgg16_config()
imdb = vid(FLAGS.image_set, FLAGS.data_path, mc)
elif FLAGS.dataset == 'ILSVRC2013':
mc = imagenet_config()
imdb = imagenet(FLAGS.image_set, FLAGS.data_path, mc)
if FLAGS.dataset != 'ILSVRC2013':
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
#joblib.dump(bbox_per_batch, '/tmp3/jeff/bbox.pkl')
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend([[i, aidx_per_batch[i][j], k]
for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
else:
num_discarded_labels += 1
print('Warning: Discarded {}/({}) labels that are assigned to the same'
'anchor'.format(num_discarded_labels, num_labels))
# Visualize detections
_viz_gt_bboxes(mc, image_per_batch, bbox_per_batch, label_per_batch)
elif FLAGS.dataset == 'ILSVRC2013':
image_per_batch, label_per_batch, _ = imdb.read_cls_batch()
label_per_batch = map(str, label_per_batch)
# TODO(jeff): visualize classification image
_viz_cls_labels(mc, image_per_batch, label_per_batch)
# Save the images
for i, im in enumerate(image_per_batch):
fname = os.path.join(FLAGS.output_dir, '{}.jpg'.format(i))
cv2.imwrite(fname, im)
def main(argv):
config_dir = os.path.abspath(argv[1])
assert config_dir.startswith(
Models), 'Invalid config directory %s' % config_dir
model_name, config_name = config_dir[len(Models):].split('/')[:2]
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
model_dir = Models + model_name + '/'
mc = kitti_squeezeDetPlus_config()
config_dir = model_dir + config_name + '/'
config_path = config_dir + 'config.json'
with open(config_path, 'r+') as f: # load custom params
for key, value in json.load(f).items():
mc[key] = value
mc.IS_TRAINING = FLAGS.train
mc.PRETRAINED_MODEL_PATH = model_dir + 'pretrained.pkl'
train_dir = config_dir + 'train/'
test_dir = config_dir + 'val/'
if FLAGS.save_root:
for dname in ['train', 'val']:
orig_dir = os.path.join(config_dir, dname)
save_dir = os.path.join(FLAGS.save_root, 'models', model_name,
config_name, dname)
make_dir(save_dir)
if not os.path.exists(orig_dir):
os.symlink(save_dir, orig_dir)
print('Creating symlink %s -> %s' % (orig_dir, save_dir))
elif not os.path.islink(orig_dir):
raise RuntimeError(
'%s exists but is not a link. Cannot create new link to %s'
% (orig_dir, save_dir))
if FLAGS.debug:
config.PRINT_STEP = config.SUMMARY_STEP = config.CHECKPOINT_STEP = config.MAX_STEPS = 1
if mc.IS_TRAINING:
kitti_set = 'train'
summary_dir = train_dir
else:
kitti_set = 'val'
summary_dir = test_dir
mc.BATCH_SIZE = 1
imdb = kitti(kitti_set, Root + 'data/KITTI', mc)
summary_writer = tf.summary.FileWriter(summary_dir)
sys.path.append(model_dir)
from load_model import load_model
model = load_model(mc)
def _load_data(load_to_placeholder=True):
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, bbox_per_batch = imdb.read_batch(
)
label_indices, bbox_indices, box_delta_values, mask_indices, box_values = [], [], [], [], []
aidx_set = set()
for i in range(len(label_per_batch)):
for j in range(len(label_per_batch[i])):
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend([[i, aidx_per_batch[i][j], k]
for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
feed_dict = {
image_input:
image_per_batch,
input_mask:
np.reshape(
sparse_to_dense(mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0] * len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
box_delta_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
box_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
labels:
sparse_to_dense(label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0] * len(label_indices)),
}
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _ = _load_data()
sess.run(model.enqueue_op, feed_dict=feed_dict)
except Exception, e:
if not sess.run(model.FIFOQueue.is_closed()):
coord.request_stop(e)
def train():
"""Train SqueezeDet model"""
assert FLAGS.dataset == 'KITTI' or FLAGS.dataset == 'CITYSCAPE', \
'Currently only support KITTI and CITYSCAPE datasets'
assert FLAGS.mask_parameterization in [4,8], 'Values other than 4 and 8 are not supported !'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default():
assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
if FLAGS.dataset == 'KITTI':
mc = kitti_vgg16_config(FLAGS.mask_parameterization, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
elif FLAGS.dataset == 'CITYSCAPE':
mc = cityscape_vgg16_config(FLAGS.mask_parameterization, FLAGS.log_anchors, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
mc.IS_TRAINING = True
# mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
print("Not using pretrained model for VGG, uncomment above line and comment below line to use pretrained model !")
mc.LOAD_PRETRAINED_MODEL = False
if FLAGS.warm_restart_lr != -1.0:
print("Updating the learning rate for warm restart to", FLAGS.warm_restart_lr)
mc.LEARNING_RATE = FLAGS.warm_restart_lr
model = VGG16ConvDet(mc)
elif FLAGS.net == 'resnet50':
if FLAGS.dataset == 'KITTI':
mc = kitti_res50_config(FLAGS.mask_parameterization, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
elif FLAGS.dataset == 'CITYSCAPE':
mc = cityscape_res50_config(FLAGS.mask_parameterization, FLAGS.log_anchors, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
if FLAGS.warm_restart_lr != -1.0:
print("Updating the learning rate for warm restart to", FLAGS.warm_restart_lr)
mc.LEARNING_RATE = FLAGS.warm_restart_lr
model = ResNet50ConvDet(mc)
elif FLAGS.net == 'squeezeDet':
if FLAGS.dataset == 'KITTI':
mc = kitti_squeezeDet_config(FLAGS.mask_parameterization, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
elif FLAGS.dataset == 'CITYSCAPE':
mc = cityscape_squeezeDet_config(FLAGS.mask_parameterization, FLAGS.log_anchors, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
if FLAGS.warm_restart_lr != -1.0:
print("Updating the learning rate for warm restart to", FLAGS.warm_restart_lr)
mc.LEARNING_RATE = FLAGS.warm_restart_lr
model = SqueezeDet(mc)
elif FLAGS.net == 'squeezeDet+':
if FLAGS.dataset == 'KITTI':
mc = kitti_squeezeDetPlus_config(FLAGS.mask_parameterization, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
elif FLAGS.dataset == 'CITYSCAPE':
mc = cityscape_squeezeDetPlus_config(FLAGS.mask_parameterization, FLAGS.log_anchors, FLAGS.only_tune_last_layer, FLAGS.encoding_type)
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
if FLAGS.warm_restart_lr != -1.0:
print("Updating the learning rate for warm restart to", FLAGS.warm_restart_lr)
mc.LEARNING_RATE = FLAGS.warm_restart_lr
model = SqueezeDetPlus(mc)
imdb_valid = None
if FLAGS.dataset == 'KITTI':
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
if FLAGS.eval_valid:
imdb_valid = kitti('val', FLAGS.data_path, mc)
imdb_valid.mc.DATA_AUGMENTATION = False
elif FLAGS.dataset == 'CITYSCAPE':
imdb = cityscape(FLAGS.image_set, FLAGS.data_path, mc)
print("Margins for Training data:", imdb.left_margin, imdb.top_margin, imdb.right_margin, imdb.bottom_margin)
if FLAGS.eval_valid:
imdb_valid = cityscape('val', FLAGS.data_path, mc)
imdb_valid.mc.DATA_AUGMENTATION = False
print("Margins for Validation data:", imdb_valid.left_margin, imdb_valid.top_margin, imdb_valid.right_margin, imdb_valid.bottom_margin)
print("Training model data augmentation:", imdb.mc.DATA_AUGMENTATION)
if imdb_valid != None:
print("Validation model data augmentation:", imdb_valid.mc.DATA_AUGMENTATION)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'), 'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print ('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
def _load_data(load_to_placeholder=True, eval_valid=False):
# read batch input
if eval_valid:
# Only for validation set
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch, edge_adhesions_per_batch = imdb_valid.read_batch(shuffle=False, wrap_around=False)
keep_prob_value = 1.0
else:
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch, edge_adhesions_per_batch = imdb.read_batch()
keep_prob_value = mc.DROP_OUT_PROB
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, edge_adhesions, edge_indices\
= [], [], [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend(
[[i, aidx_per_batch[i][j], k] for k in range(FLAGS.mask_parameterization)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
edge_adhesions.extend(edge_adhesions_per_batch[i][j])
edge_indices.extend(
[[i, aidx_per_batch[i][j], k] for k in range(FLAGS.mask_parameterization)])
else:
num_discarded_labels += 1
if mc.DEBUG_MODE:
print ('Warning: Discarded {}/({}) labels that are assigned to the same '
'anchor'.format(num_discarded_labels, num_labels))
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
edge_scenarios = model.ph_edge_adhesions
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
edge_scenarios = model.edge_adhesions
feed_dict = {
image_input: image_per_batch,
input_mask: np.reshape(
sparse_to_dense(
mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0]*len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
box_delta_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, FLAGS.mask_parameterization],
box_delta_values),
box_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, FLAGS.mask_parameterization],
box_values),
labels: sparse_to_dense(
label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0]*len(label_indices)),
model.keep_prob: keep_prob_value,
edge_scenarios: sparse_to_dense(
edge_indices, [mc.BATCH_SIZE, mc.ANCHORS, FLAGS.mask_parameterization],
edge_adhesions).astype(np.bool, copy=False),
}
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch, edge_indices
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _, _ = _load_data()
sess.run(model.enqueue_op, feed_dict=feed_dict)
if mc.DEBUG_MODE:
print ("added to the queue")
if mc.DEBUG_MODE:
print ("Finished enqueue")
except tf.errors.CancelledError:
coord.request_stop()
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver(tf.global_variables())
summary_op = tf.summary.merge_all()
init = tf.global_variables_initializer()
sess.run(init)
glb_step = sess.run(model.global_step)
print("Global step before restore:", glb_step)
print("Kernels before restore")
for v in tf.trainable_variables():
if 'kernels' in v.name:
print("First few weights of ", v.name, " are ", sess.run(v)[0,0,0,0:5])
print("Learning rate before restore", sess.run(model.lr))
ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
if ckpt and ckpt.model_checkpoint_path:
print("Found checkpoint at step: ", int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]))
last_layer_name = model.preds.name.split('/')[0]
if FLAGS.mask_parameterization == 8 and FLAGS.bounding_box_checkpoint:
print("Loading only partial weights (except last layer", last_layer_name, ")")
saver_partial_weights = tf.train.Saver([v for v in tf.global_variables() if last_layer_name not in v.name])
saver_partial_weights.restore(sess, ckpt.model_checkpoint_path)
if FLAGS.warm_restart_lr != -1.0:
print("Resetting global step")
sess.run([model.global_step.assign(0)])
else:
print("Loading all weights (including the last layer", last_layer_name, ")")
saver.restore(sess, ckpt.model_checkpoint_path)
else:
print("Checkpoint not found !")
glb_step = sess.run(model.global_step)
print("Global step after restore:", glb_step)
print("Kernels after restore")
for v in tf.trainable_variables():
if 'kernels' in v.name:
print("First few weights of ", v.name, " are ", sess.run(v)[0,0,0,0:5])
print("Learning rate after restore", sess.run(model.lr))
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
with open(os.path.join(FLAGS.train_dir, 'training_metrics.txt'), 'a') as f:
f.write("Global step after restore: "+str(glb_step)+"\n")
f.close()
if FLAGS.eval_valid:
with open(os.path.join(FLAGS.train_dir, 'validation_metrics.txt'), 'a') as f:
f.write("Global step after restore: "+str(glb_step)+"\n")
f.close()
coord = tf.train.Coordinator()
if mc.NUM_THREAD > 0:
enq_threads = []
for _ in range(mc.NUM_THREAD):
enq_thread = threading.Thread(target=_enqueue, args=[sess, coord])
# enq_thread.isDaemon()
enq_thread.start()
enq_threads.append(enq_thread)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
run_options = tf.RunOptions(timeout_in_ms=60000)
try:
for step in xrange(glb_step, FLAGS.max_steps):
if coord.should_stop():
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
break
start_time = time.time()
if step % FLAGS.summary_step == 0:
feed_dict, image_per_batch, label_per_batch, bbox_per_batch, edge_ids = \
_load_data(load_to_placeholder=False)
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss, model.edge_adhesions,
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, \
conf_loss, bbox_loss, class_loss, edge_adhesions_pre_filtered = sess.run(
op_list, feed_dict=feed_dict)
summary_writer.add_summary(summary_str, step)
# Visualize the training examples only if validation is not enabled
if not FLAGS.eval_valid:
visualize_gt_masks = False
visualize_pred_masks = False
if mc.EIGHT_POINT_REGRESSION:
visualize_gt_masks = True
visualize_pred_masks = True
assert np.array_equal(feed_dict[model.edge_adhesions], edge_adhesions_pre_filtered), \
"Training Gt edge adhesion not matching edge adhesion tensor"
edge_adhesions_per_batch = [[0]]*mc.BATCH_SIZE
for id_val in range(mc.BATCH_SIZE):
selected_ids = np.where(np.asarray(edge_ids)[:,0] == id_val)[0]
# print("Before",np.asarray(edge_ids)[selected_ids][:,1])
indexes_int = np.unique(np.asarray(edge_ids)[selected_ids][:,1], return_index=True)[1]
anchors_ids = np.asarray([np.asarray(edge_ids)[selected_ids][:,1][index] for index in sorted(indexes_int)])
# print("After",anchors_ids)
batch_id = [id_val]*len(anchors_ids)
edge_adhesions_per_batch[id_val] = edge_adhesions_pre_filtered[batch_id, anchors_ids]
_viz_prediction_result(
model, image_per_batch, bbox_per_batch, label_per_batch, det_boxes,
det_class, det_probs, visualize_gt_masks, visualize_pred_masks)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op, feed_dict={model.image_to_show: image_per_batch})
summary_writer.add_summary(viz_summary, step)
print ('total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}'.\
format(loss_value, conf_loss, bbox_loss, class_loss))
with open(os.path.join(FLAGS.train_dir, 'training_metrics.txt'), 'a') as f:
f.write('step: {}, total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}\n'.\
format(step, loss_value, conf_loss, bbox_loss, class_loss))
f.close()
if FLAGS.eval_valid:
print ('\n!! Validation Set evaluation at step ', step, ' !!')
with open(os.path.join(FLAGS.train_dir, 'validation_metrics.txt'), 'a') as f:
f.write('\n!! Validation Set evaluation at step '+str(step)+' !!\n')
loss_list = []
batch_nr = 0
while True:
batch_nr += 1
if len(imdb_valid._image_idx) % mc.BATCH_SIZE > 0:
# if batch_size unevenly divides the number of samples.
# then number of batches is one more than the actual num of batches
num_of_batches = (len(imdb_valid._image_idx) // mc.BATCH_SIZE) + 1
else:
num_of_batches = (len(imdb_valid._image_idx) // mc.BATCH_SIZE)
if batch_nr > num_of_batches:
break
feed_dict_val, image_per_batch_val, label_per_batch_val, bbox_per_batch_val, edge_ids_val = \
_load_data(load_to_placeholder=False, eval_valid=True)
op_list_val = [
model.loss, model.conf_loss, model.bbox_loss, \
model.class_loss, model.det_boxes, \
model.det_probs, model.det_class,
model.edge_adhesions,
]
loss_value_val, conf_loss_val, bbox_loss_val, class_loss_val, det_boxes_val, \
det_probs_val, det_class_val, edge_adhesions_pre_filtered_val = sess.run(op_list_val, feed_dict=feed_dict_val)
if batch_nr == 1:
# Sample the first batch for visualization
visualize_gt_masks = False
visualize_pred_masks = False
if mc.EIGHT_POINT_REGRESSION:
visualize_gt_masks = True
visualize_pred_masks = True
assert np.array_equal(feed_dict_val[model.edge_adhesions], edge_adhesions_pre_filtered_val), \
"Validation Gt edge adhesion not matching edge adhesion tensor"
edge_adhesions_per_batch_val = [[0]]*mc.BATCH_SIZE
for id_val in range(mc.BATCH_SIZE):
selected_ids = np.where(np.asarray(edge_ids_val)[:,0] == id_val)[0]
indexes_int = np.unique(np.asarray(edge_ids_val)[selected_ids][:,1], return_index=True)[1]
anchors_ids = np.asarray([np.asarray(edge_ids_val)[selected_ids][:,1][index] for index in sorted(indexes_int)])
batch_id = [id_val]*len(anchors_ids)
edge_adhesions_per_batch_val[id_val] = edge_adhesions_pre_filtered_val[batch_id, anchors_ids]
_viz_prediction_result(
model, image_per_batch_val, bbox_per_batch_val, label_per_batch_val, det_boxes_val,
det_class_val, det_probs_val, visualize_gt_masks, visualize_pred_masks)
image_per_batch_visualize = bgr_to_rgb(image_per_batch_val)
loss_list.append([loss_value_val, conf_loss_val, bbox_loss_val, class_loss_val])
f.write('Batch: {}, total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}\n'.\
format(batch_nr, loss_value_val, conf_loss_val, bbox_loss_val, class_loss_val))
loss_list = np.asarray(loss_list)
loss_means = [np.mean(loss_list[:,0]), np.mean(loss_list[:,1]), np.mean(loss_list[:,2]), np.mean(loss_list[:,3])]
loss_stds = [np.std(loss_list[:,0]), np.std(loss_list[:,1]), np.std(loss_list[:,2]), np.std(loss_list[:,3])]
print ('Mean values : total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}'.\
format(loss_means[0], loss_means[1], loss_means[2], loss_means[3]))
print ('Standard Deviation values : total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}'.\
format(loss_stds[0], loss_stds[1], loss_stds[2], loss_stds[3]))
f.write('Mean values : total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}\n'.\
format(loss_means[0], loss_means[1], loss_means[2], loss_means[3]))
f.write('Standard Deviation values : total_loss: {}, conf_loss: {}, bbox_loss: {}, class_loss: {}\n'.\
format(loss_stds[0], loss_stds[1], loss_stds[2], loss_stds[3]))
# Visualize the validation examples
if len(image_per_batch_visualize) != 0:
viz_summary = sess.run(
model.viz_op, feed_dict={model.image_to_show: image_per_batch_visualize})
summary_writer.add_summary(viz_summary, step)
f.close()
summary_writer.flush()
else:
if mc.NUM_THREAD > 0:
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], options=run_options)
else:
feed_dict, _, _, _, _ = _load_data(load_to_placeholder=False)
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = ('%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print (format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))
with open(os.path.join(FLAGS.train_dir, 'training_metrics.txt'), 'a') as f:
f.write(format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch) + '\n')
f.close()
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step + 1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
print("Checkpointing at ", step)
saver.save(sess, checkpoint_path, global_step=step)
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
except KeyboardInterrupt:
print("Keyboard interrupt caught ! Terminating..")
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
sys.exit(0)
except:
print("Unexpected error:", sys.exc_info()[0])
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
sys.exit(0)
def train():
"""Train SqueezeDet model"""
assert FLAGS.dataset == 'KITTI', \
'Currently only support KITTI dataset'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default():
# assert FLAGS.net == 'vgg16' or FLAGS.net == 'resnet50' \
# or FLAGS.net == 'squeezeDet' or FLAGS.net=="zynqDet" or FLAGS.net == 'squeezeDet+', \
# 'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'vgg16':
mc = kitti_vgg16_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = VGG16ConvDet(mc)
elif FLAGS.net == 'resnet50':
mc = kitti_res50_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = ResNet50ConvDet(mc)
elif FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDet(mc)
elif FLAGS.net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlus(mc)
elif FLAGS.net == "zynqDet":
mc = kitti_zynqDet_FPN_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = ZynqDet_FPN(mc)
elif FLAGS.net == "squeezeDet_FPN":
mc = kitti_squeezeDet_FPN_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDet_FPN(mc)
elif FLAGS.net == "yolo":
mc = kitti_vgg16_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = tinyDarkNet_FPN(mc)
elif FLAGS.net == "ZynqDet_Quant":
mc = kitti_zynqDet_FPN_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = ZynqDet_FPN_Quant(mc)
else:
assert (0)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'),
'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
def _load_data_per_scale(label_per_batch, box_delta_per_batch,
aidx_per_batch, bbox_per_batch, s, ANCHORS):
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
for i in range(len(label_per_batch)): # batch_size
for j in range(
len(aidx_per_batch[i]
[s])): # number of annotations with IOU > 0.5
aidx, gt_id = aidx_per_batch[i][s][j]
if (i, aidx) not in aidx_set:
aidx_set.add((i, aidx))
label_indices.append(
[i, aidx, label_per_batch[i][gt_id]])
mask_indices.append([i, aidx])
bbox_indices.extend([[i, aidx, k] for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][s][j])
box_values.extend(bbox_per_batch[i][gt_id])
else:
num_discarded_labels += 1
input_mask = np.reshape(
sparse_to_dense(mask_indices, [mc.BATCH_SIZE, ANCHORS],
[1.0] * len(mask_indices)),
[mc.BATCH_SIZE, ANCHORS, 1])
box_delta_input = sparse_to_dense(bbox_indices,
[mc.BATCH_SIZE, ANCHORS, 4],
box_delta_values)
box_input = sparse_to_dense(bbox_indices,
[mc.BATCH_SIZE, ANCHORS, 4],
box_values)
labels = sparse_to_dense(label_indices,
[mc.BATCH_SIZE, ANCHORS, mc.CLASSES],
[1.0] * len(label_indices))
return input_mask, box_delta_input, box_input, labels
def _load_data(load_to_placeholder=True):
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
# print("label", label_per_batch[0])
# print(aidx_per_batch[0])
input_mask, box_delta_input, box_input, labels =\
_load_data_per_scale(label_per_batch, box_delta_per_batch, aidx_per_batch, bbox_per_batch, 0, mc.ANCHORS)
input_mask2, box_delta_input2, box_input2, labels2 =\
_load_data_per_scale(label_per_batch, box_delta_per_batch, aidx_per_batch, bbox_per_batch, 1, mc.ANCHORS2)
input_mask3, box_delta_input3, box_input3, labels3 =\
_load_data_per_scale(label_per_batch, box_delta_per_batch, aidx_per_batch, bbox_per_batch, 2, mc.ANCHORS3)
input_mask_total = np.concatenate(
[input_mask, input_mask2, input_mask3], 1)
box_delta_input_total = np.concatenate(
[box_delta_input, box_delta_input2, box_delta_input3], 1)
box_input_total = np.concatenate(
[box_input, box_input2, box_input3], 1)
labels_total = np.concatenate([labels, labels2, labels3], 1)
# print("input_mask_shape", input_mask_total.shape)
# if mc.DEBUG_MODE:
# print ('Warning: Discarded {}/({}) labels that are assigned to the same '
# 'anchor'.format(num_discarded_labels, num_labels))
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
feed_dict = {
image_input: image_per_batch,
input_mask: input_mask_total,
box_delta_input: box_delta_input_total,
box_input: box_input_total,
labels: labels_total
}
# feed_dict = {
# image_input: image_per_batch,
# input_mask: np.reshape(
# sparse_to_dense(
# mask_indices, [mc.BATCH_SIZE, mc.ANCHOR_TOTAL],
# [1.0]*len(mask_indices)),
# [mc.BATCH_SIZE, mc.ANCHOR_TOTAL, 1]),
# box_delta_input: sparse_to_dense(
# bbox_indices, [mc.BATCH_SIZE, mc.ANCHOR_TOTAL, 4],
# box_delta_values),
# box_input: sparse_to_dense(
# bbox_indices, [mc.BATCH_SIZE, mc.ANCHOR_TOTAL, 4],
# box_values),
# labels: sparse_to_dense(
# label_indices,
# [mc.BATCH_SIZE, mc.ANCHOR_TOTAL, mc.CLASSES],
# [1.0]*len(label_indices)),
# }
# feed_dict = {
# image_input: image_per_batch,
# input_mask: np.reshape(
# sparse_to_dense(
# mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
# [1.0]*len(mask_indices)),
# [mc.BATCH_SIZE, mc.ANCHORS, 1]),
# box_delta_input: sparse_to_dense(
# bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
# box_delta_values),
# box_input: sparse_to_dense(
# bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
# box_values),
# labels: sparse_to_dense(
# label_indices,
# [mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
# [1.0]*len(label_indices)),
# }
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _ = _load_data()
# print("input.shape", feed_dict[model.ph_labels].shape)
sess.run(model.enqueue_op, feed_dict=feed_dict)
if mc.DEBUG_MODE:
print("added to the queue")
if mc.DEBUG_MODE:
print("Finished enqueue")
except Exception, e:
coord.request_stop(e)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
if FLAGS.resume:
if ckpt and ckpt.model_checkpoint_path:
print("restoring...", FLAGS.train_dir,
ckpt.model_checkpoint_path)
saver.restore(sess, ckpt.model_checkpoint_path)
else:
init = tf.global_variables_initializer()
sess.run(init)
coord = tf.train.Coordinator()
if mc.NUM_THREAD > 0:
enq_threads = []
for _ in range(mc.NUM_THREAD):
enq_thread = threading.Thread(target=_enqueue,
args=[sess, coord])
# enq_thread.isDaemon()
enq_thread.start()
enq_threads.append(enq_thread)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
run_options = tf.RunOptions(timeout_in_ms=60000)
best_map = 0
# try:
for step in xrange(FLAGS.max_steps):
# for step in xrange(1):
if coord.should_stop():
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
break
start_time = time.time()
if step % FLAGS.summary_step == 0:
feed_dict, image_per_batch, label_per_batch, bbox_per_batch = \
_load_data(load_to_placeholder=False)
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, \
conf_loss, bbox_loss, class_loss = sess.run(
op_list, feed_dict=feed_dict)
_viz_prediction_result(model, image_per_batch, bbox_per_batch,
label_per_batch, det_boxes, det_class,
det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op,
feed_dict={model.image_to_show: image_per_batch})
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
summary_writer.flush()
# if step > 10000:
# test.main()
# opts = options.parse_opts()
# opts.eval_dir = './data/out/tmp/bbox'
# scene_list = []
# for i in test_set:
# scene_list.append('exp{:03d}_B'.format(i))
# mAP = eval_det.eval_batch(scene_list,opts)
# if (mAP > best_map):
# best_map = mAP
# print("best_map at step ", step)
# checkpoint_path = os.path.join(FLAGS.train_dir, 'best_model.ckpt')
# saver.save(sess, checkpoint_path, global_step=step)
print('conf_loss: {}, bbox_loss: {}, class_loss: {}'.format(
conf_loss, bbox_loss, class_loss))
else:
# print("input", feed_dict[model.labels], label_per_batch)
# print(label_per_batch[0])
# print(feed_dict[model.labels][0,0,:])
# print(feed_dict[model.labels][0,1,:])
# print("input.shape", feed_dict[model.labels].shape)
if mc.NUM_THREAD > 0:
_, loss_value, conf_loss, bbox_loss, class_loss, num_objects = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.class_loss,
model.num_objects
],
options=run_options)
else:
feed_dict, _, label_per_batch, bbox_per_batch = _load_data(
load_to_placeholder=False)
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.class_loss
],
feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = (
'%s: step %d, loss = %.2f, conf_loss = %.2f, bbox_loss = %.2f, class_loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print(format_str %
(datetime.now(), step, loss_value, conf_loss, bbox_loss,
class_loss, images_per_sec, sec_per_batch))
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step +
1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir,
'model_{}.ckpt'.format(step))
saver.save(sess, checkpoint_path, global_step=step)
tf.train.write_graph(sess.graph.as_graph_def(),
FLAGS.train_dir,
'tensorflowModel.pbtxt',
as_text=True)
def train():
"""Train SqueezeDet model"""
assert FLAGS.dataset == 'KITTI', \
'Currently only support KITTI dataset'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
with tf.Graph().as_default():
assert FLAGS.net == 'shuffleDet' or FLAGS.net == 'squeezeDet' or FLAGS.net == 'squeezeDet+', \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
if FLAGS.net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDet(mc)
elif FLAGS.net == 'shuffleDet':
mc = kitti_shuffleDet_config()
mc.IS_TRAINING = True
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = shuffleDet(mc)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# save model size, flops, activations by layers
_build_model_metrics(model)
def _load_data(load_to_placeholder=True):
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend(
[[i, aidx_per_batch[i][j], k] for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
else:
num_discarded_labels += 1
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
feed_dict = {
image_input: image_per_batch,
input_mask: np.reshape(
sparse_to_dense(
mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0]*len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
box_delta_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
box_input: sparse_to_dense(
bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
labels: sparse_to_dense(
label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0]*len(label_indices)),
}
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _ = _load_data()
sess.run(model.enqueue_op, feed_dict=feed_dict)
if mc.DEBUG_MODE:
print ("added to the queue")
if mc.DEBUG_MODE:
print ("Finished enqueue")
except Exception, e:
coord.request_stop(e)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
# print(tf.global_variables())
# variables = tf.trainable_variables()
variables = tf.contrib.framework.get_variables_to_restore()
variables_to_resotre = [v for v in variables if 'Momentum' not in v.name and 'iou' not in v.name and 'global_step' not in v.name and 'preds' not in v.name]
saver = tf.train.Saver(variables_to_resotre)
summary_op = tf.summary.merge_all()
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
# saver.restore(sess, checkpoint_path)
# ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
# print('ckpt is:'+str(ckpt))
# if ckpt and ckpt.model_checkpoint_path:
# saver.restore(sess, ckpt.model_checkpoint_path)
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
init = tf.global_variables_initializer()
init = tf.initialize_all_variables()
sess.run(init)
coord = tf.train.Coordinator()
if mc.NUM_THREAD > 0:
enq_threads = []
for _ in range(mc.NUM_THREAD):
enq_thread = threading.Thread(target=_enqueue, args=[sess, coord])
# enq_thread.isDaemon()
enq_thread.start()
enq_threads.append(enq_thread)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
run_options = tf.RunOptions(timeout_in_ms=60000)
if FLAGS.reload:
checkpoint_path = os.path.join(FLAGS.train_dir, '/')
print('load the parameters from: '+str(FLAGS.train_dir))
ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
if ckpt and ckpt.model_checkpoint_path:
print('Restores from checkpoint....')
# Restores from checkpoint
saver.restore(sess, ckpt.model_checkpoint_path)
# Assuming model_checkpoint_path looks something like:
# /my-favorite-path/cifar10_train/model.ckpt-0,
# extract global_step from it.
# step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])
step = 1
else:
step = 1
print('No checkpoint found!')
else:
tf.gfile.DeleteRecursively(FLAGS.train_dir)
tf.gfile.MakeDirs(FLAGS.train_dir)
step = 1
# try:
for step in xrange(FLAGS.max_steps):
if coord.should_stop():
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
break
start_time = time.time()
if step % FLAGS.summary_step == 0:
feed_dict, image_per_batch, label_per_batch, bbox_per_batch = \
_load_data(load_to_placeholder=False)
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, \
conf_loss, bbox_loss, class_loss = sess.run(
op_list, feed_dict=feed_dict)
_viz_prediction_result(
model, image_per_batch, bbox_per_batch, label_per_batch, det_boxes,
det_class, det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op, feed_dict={model.image_to_show: image_per_batch})
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
summary_writer.flush()
print ('conf_loss: {}, bbox_loss: {}, class_loss: {}'.
format(conf_loss, bbox_loss, class_loss))
else:
if mc.NUM_THREAD > 0:
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], options=run_options)
else:
feed_dict, _, _, _ = _load_data(load_to_placeholder=False)
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[model.train_op, model.loss, model.conf_loss, model.bbox_loss,
model.class_loss], feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = ('%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print (format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step + 1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step,write_state = True)
def train():
"""Train SqueezeDetPlus model"""
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
if not os.path.exists(FLAGS.train_dir):
os.makedirs(FLAGS.train_dir)
with tf.Graph().as_default():
# Below logic is implemented to choose between pruning and regular training
# Also the type of network is chosen (depending on pruned structure)
# Learning rates are dependent on number of to be updated parameters
# Similarly as lambdas for loss functions (see paper)
if FLAGS.pruning:
if FLAGS.net == 'squeezeDet+PruneFilter':
mc = kitti_squeezeDetPlus_config()
mc.LEARNING_RATE = 0.1
mc.IS_TRAINING = True
mc.IS_PRUNING = True
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilter(mc)
elif FLAGS.net == 'squeezeDet+PruneFilterShape':
mc = kitti_squeezeDetPlus_config()
mc.LEARNING_RATE = 0.001
mc.IS_TRAINING = True
mc.IS_PRUNING = True
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilterShape(mc)
elif FLAGS.net == 'squeezeDet+PruneLayer':
mc = kitti_squeezeDetPlus_config()
mc.LEARNING_RATE = 0.0003
mc.IS_TRAINING = True
mc.IS_PRUNING = True
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneLayer(mc)
else:
if FLAGS.net == 'squeezeDet+PruneFilter':
mc = kitti_squeezeDetPlus_config()
mc.IS_TRAINING = True
mc.IS_PRUNING = False
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilter(mc)
elif FLAGS.net == 'squeezeDet+PruneFilterShape':
mc = kitti_squeezeDetPlus_config()
mc.IS_TRAINING = True
mc.IS_PRUNING = False
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneFilterShape(mc)
elif FLAGS.net == 'squeezeDet+PruneLayer':
mc = kitti_squeezeDetPlus_config()
mc.IS_TRAINING = True
mc.IS_PRUNING = False
mc.LITE_MODE = False
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
model = SqueezeDetPlusPruneLayer(mc)
imdb = kitti(FLAGS.image_set, FLAGS.data_path, mc)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'),
'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
def _load_data(load_to_placeholder=True):
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = imdb.read_batch()
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
= [], [], [], [], []
aidx_set = set()
num_discarded_labels = 0
num_labels = 0
for i in range(len(label_per_batch)): # batch_size
for j in range(len(
label_per_batch[i])): # number of annotations
num_labels += 1
if (i, aidx_per_batch[i][j]) not in aidx_set:
aidx_set.add((i, aidx_per_batch[i][j]))
label_indices.append(
[i, aidx_per_batch[i][j], label_per_batch[i][j]])
mask_indices.append([i, aidx_per_batch[i][j]])
bbox_indices.extend([[i, aidx_per_batch[i][j], k]
for k in range(4)])
box_delta_values.extend(box_delta_per_batch[i][j])
box_values.extend(bbox_per_batch[i][j])
else:
num_discarded_labels += 1
if mc.DEBUG_MODE:
print(
'Warning: Discarded {}/({}) labels that are assigned to the same '
'anchor'.format(num_discarded_labels, num_labels))
if load_to_placeholder:
image_input = model.ph_image_input
input_mask = model.ph_input_mask
box_delta_input = model.ph_box_delta_input
box_input = model.ph_box_input
labels = model.ph_labels
else:
image_input = model.image_input
input_mask = model.input_mask
box_delta_input = model.box_delta_input
box_input = model.box_input
labels = model.labels
feed_dict = {
image_input:
image_per_batch,
input_mask:
np.reshape(
sparse_to_dense(mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0] * len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
box_delta_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
box_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
labels:
sparse_to_dense(label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0] * len(label_indices)),
}
return feed_dict, image_per_batch, label_per_batch, bbox_per_batch
def _enqueue(sess, coord):
try:
while not coord.should_stop():
feed_dict, _, _, _ = _load_data()
sess.run(model.enqueue_op, feed_dict=feed_dict)
if mc.DEBUG_MODE:
print("added to the queue")
if mc.DEBUG_MODE:
print("Finished enqueue")
except Exception as e:
coord.request_stop(e)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
summary_op = tf.summary.merge_all()
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
if ckpt and ckpt.model_checkpoint_path:
print("restoring checkpoint")
saver.restore(sess, ckpt.model_checkpoint_path)
else:
uninitialized_vars = []
for var in tf.all_variables():
try:
sess.run(var)
except tf.errors.FailedPreconditionError:
uninitialized_vars.append(var)
init_new_vars_op = tf.initialize_variables(uninitialized_vars)
sess.run(init_new_vars_op)
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
coord = tf.train.Coordinator()
if mc.NUM_THREAD > 0:
enq_threads = []
for _ in range(mc.NUM_THREAD):
enq_thread = threading.Thread(target=_enqueue,
args=[sess, coord])
# enq_thread.isDaemon()
enq_thread.start()
enq_threads.append(enq_thread)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
run_options = tf.RunOptions(timeout_in_ms=60000)
for step in range(FLAGS.max_steps):
if coord.should_stop():
sess.run(model.FIFOQueue.close(cancel_pending_enqueues=True))
coord.request_stop()
coord.join(threads)
break
start_time = time.time()
if step % FLAGS.summary_step == 0:
feed_dict, image_per_batch, label_per_batch, bbox_per_batch = \
_load_data(load_to_placeholder=False)
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.gamma_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, \
conf_loss, bbox_loss, gamma_loss, class_loss = sess.run(
op_list, feed_dict=feed_dict)
_viz_prediction_result(model, image_per_batch, bbox_per_batch,
label_per_batch, det_boxes, det_class,
det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op,
feed_dict={model.image_to_show: image_per_batch})
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
summary_writer.flush()
print(
'conf_loss: {}, bbox_loss: {}, gamma_loss: {}, class_loss: {}'
.format(conf_loss, bbox_loss, gamma_loss, class_loss))
else:
if mc.NUM_THREAD > 0:
_, loss_value, conf_loss, bbox_loss, gamma_loss, class_loss = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.gamma_loss, model.class_loss
],
options=run_options)
else:
feed_dict, _, _, _ = _load_data(load_to_placeholder=False)
_, loss_value, conf_loss, bbox_loss, gamma_loss, class_loss = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.gamma_loss, model.class_loss
],
feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, gamma_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, gamma_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = (
'%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print (format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))\
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step +
1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
coord.request_stop()
def train():
"""Train SqueezeDet model"""
assert FLAGS.net in ['vgg16', 'vgg16_v2', 'vgg16_v3', 'yolo_v2'], \
'Selected neural net architecture not supported: {}'.format(FLAGS.net)
assert FLAGS.dataset in ['KITTI', 'PASCAL_VOC', 'VID'], \
'Either KITTI / PASCAL_VOC / VID'
if FLAGS.dataset == 'KITTI':
if FLAGS.net == 'yolo_v2':
raise NotImplementedError
else:
mc = kitti_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
train_imdb = kitti(FLAGS.train_set, FLAGS.data_path, mc)
elif FLAGS.dataset == 'PASCAL_VOC':
if FLAGS.net == 'yolo_v2':
mc = pascal_voc_yolo_config()
else:
mc = pascal_voc_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
train_imdb = pascal_voc(FLAGS.train_set, FLAGS.year, FLAGS.data_path,
mc)
elif FLAGS.dataset == 'VID':
if FLAGS.net == 'yolo_v2':
mc = vid_yolo_config()
else:
mc = vid_vgg16_config()
mc.PRETRAINED_MODEL_PATH = FLAGS.pretrained_model_path
train_imdb = vid(FLAGS.train_set, FLAGS.data_path, mc)
with tf.Graph().as_default():
if FLAGS.net == 'vgg16':
model = VGG16ConvDet(mc, FLAGS.gpu)
elif FLAGS.net == 'vgg16_v2':
model = VGG16ConvDetV2(mc, FLAGS.gpu)
elif FLAGS.net == 'vgg16_v3':
model = VGG16ConvDetV3(mc, FLAGS.gpu)
elif FLAGS.net == 'yolo_v2':
model = YOLO_V2(mc, FLAGS.gpu)
# save model size, flops, activations by layers
with open(os.path.join(FLAGS.train_dir, 'model_metrics.txt'),
'w') as f:
f.write('Number of parameter by layer:\n')
count = 0
for c in model.model_size_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nActivation size by layer:\n')
for c in model.activation_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
count = 0
f.write('\nNumber of flops by layer:\n')
for c in model.flop_counter:
f.write('\t{}: {}\n'.format(c[0], c[1]))
count += c[1]
f.write('\ttotal: {}\n'.format(count))
f.close()
print('Model statistics saved to {}.'.format(
os.path.join(FLAGS.train_dir, 'model_metrics.txt')))
saver = tf.train.Saver(tf.all_variables())
summary_op = tf.summary.merge_all()
init = tf.initialize_all_variables()
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(init)
tf.train.start_queue_runners(sess=sess)
summary_writer = tf.summary.FileWriter(FLAGS.train_dir, sess.graph)
for step in xrange(FLAGS.max_steps):
start_time = time.time()
# read batch input
image_per_batch, label_per_batch, box_delta_per_batch, aidx_per_batch, \
bbox_per_batch = train_imdb.read_batch()
label_indices, bbox_indices, box_delta_values, mask_indices, box_values, \
num_discarded_labels, num_labels = _process_batch(image_per_batch, \
label_per_batch, box_delta_per_batch, \
aidx_per_batch, bbox_per_batch, mc.DEBUG_MODE)
feed_dict = {
model.image_input:
image_per_batch,
model.is_training:
True,
model.keep_prob:
mc.KEEP_PROB,
model.input_mask:
np.reshape(
sparse_to_dense(mask_indices, [mc.BATCH_SIZE, mc.ANCHORS],
[1.0] * len(mask_indices)),
[mc.BATCH_SIZE, mc.ANCHORS, 1]),
model.box_delta_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_delta_values),
model.box_input:
sparse_to_dense(bbox_indices, [mc.BATCH_SIZE, mc.ANCHORS, 4],
box_values),
model.labels:
sparse_to_dense(label_indices,
[mc.BATCH_SIZE, mc.ANCHORS, mc.CLASSES],
[1.0] * len(label_indices)),
}
if step % FLAGS.summary_step == 0:
op_list = [
model.train_op, model.loss, summary_op, model.det_boxes,
model.det_probs, model.det_class, model.conf_loss,
model.bbox_loss, model.class_loss
]
_, loss_value, summary_str, det_boxes, det_probs, det_class, conf_loss, \
bbox_loss, class_loss = sess.run(op_list, feed_dict=feed_dict)
_viz_prediction_result(model, image_per_batch, bbox_per_batch,
label_per_batch, det_boxes, det_class,
det_probs)
image_per_batch = bgr_to_rgb(image_per_batch)
viz_summary = sess.run(
model.viz_op,
feed_dict={model.image_to_show: image_per_batch})
num_discarded_labels_op = tf.summary.scalar(
'counter/num_discarded_labels', num_discarded_labels)
num_labels_op = tf.summary.scalar('counter/num_labels',
num_labels)
counter_summary_str = sess.run(
[num_discarded_labels_op, num_labels_op])
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(viz_summary, step)
for sum_str in counter_summary_str:
summary_writer.add_summary(sum_str, step)
print('conf_loss: {}, bbox_loss: {}, class_loss: {}'.format(
conf_loss, bbox_loss, class_loss))
else:
_, loss_value, conf_loss, bbox_loss, class_loss = sess.run(
[
model.train_op, model.loss, model.conf_loss,
model.bbox_loss, model.class_loss
],
feed_dict=feed_dict)
duration = time.time() - start_time
assert not np.isnan(loss_value), \
'Model diverged. Total loss: {}, conf_loss: {}, bbox_loss: {}, ' \
'class_loss: {}'.format(loss_value, conf_loss, bbox_loss, class_loss)
if step % 10 == 0:
num_images_per_step = mc.BATCH_SIZE
images_per_sec = num_images_per_step / duration
sec_per_batch = float(duration)
format_str = (
'%s: step %d, loss = %.2f (%.1f images/sec; %.3f '
'sec/batch)')
print(format_str % (datetime.now(), step, loss_value,
images_per_sec, sec_per_batch))
sys.stdout.flush()
# Save the model checkpoint periodically.
if step % FLAGS.checkpoint_step == 0 or (step +
1) == FLAGS.max_steps:
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
