def get_session_and_models():
''' Define model graph, load model parameters,
create session and return session handle and tensors
'''
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
# image segmentaion
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
img_seg_net = ImgSegNet(BATCH_SIZE, NUM_POINT)
img_seg_net.load_graph(FLAGS.img_seg_model)
sess1 = tf.Session(config=config)
# point RPN
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
rpn_model = RPN(BATCH_SIZE,
NUM_POINT,
num_channel=4,
is_training=False)
pls = rpn_model.placeholders
box_center, box_angle, box_size = rpn_model.box_encoder.tf_decode(
rpn_model.end_points)
box_center = box_center + rpn_model.end_points['fg_points_xyz']
rpn_model.end_points['box_center'] = box_center
rpn_model.end_points['box_angle'] = box_angle
rpn_model.end_points['box_size'] = box_size
sess2 = tf.Session(config=config)
saver = tf.train.Saver()
saver.restore(sess2, FLAGS.rpn_model)
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
rcnn_model = FrustumPointNet(BATCH_SIZE_RCNN, NUM_POINT_RCNN)
sess3 = tf.Session(config=config)
saver = tf.train.Saver()
saver.restore(sess3, FLAGS.rcnn_model)
return sess1, img_seg_net, sess2, rpn_model, sess3, rcnn_model
def get_session_and_model(batch_size, num_point):
''' Define model graph, load model parameters,
create session and return session handle and tensors
'''
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
frustum_pointnet = FrustumPointNet(batch_size, num_point)
placeholders = frustum_pointnet.placeholders
end_points = frustum_pointnet.end_points
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
# Restore variables from disk.
saver.restore(sess, MODEL_PATH)
return sess, frustum_pointnet
def train():
''' Main function for training and simple evaluation. '''
best_val_loss = float('inf')
best_avg_cls_acc = 0
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
# Note the global_step=batch parameter to minimize.
# That tells the optimizer to increment the 'batch' parameter
# for you every time it trains.
batch = tf.get_variable('batch', [],
initializer=tf.constant_initializer(0),
trainable=False)
bn_decay = get_bn_decay(batch)
tf.summary.scalar('bn_decay', bn_decay)
# Get model and losses
frustum_pointnet = FrustumPointNet(BATCH_SIZE, NUM_POINT, bn_decay)
end_points = frustum_pointnet.end_points
pls = frustum_pointnet.placeholders
loss, loss_endpoints = frustum_pointnet.get_loss()
tf.summary.scalar('loss', loss)
losses = tf.get_collection('losses')
total_loss = tf.add_n(losses, name='total_loss')
tf.summary.scalar('total_loss', total_loss)
# Write summaries of bounding box IoU and segmentation accuracies
iou2ds, iou3ds = tf.py_func(provider.compute_box3d_iou, [\
end_points['center'], \
end_points['heading_scores'], end_points['heading_residuals'], \
end_points['size_scores'], end_points['size_residuals'], \
pls['centers'], \
pls['heading_class_label'], pls['heading_residual_label'], \
pls['size_class_label'], pls['size_residual_label']], \
[tf.float32, tf.float32])
end_points['iou2ds'] = iou2ds
end_points['iou3ds'] = iou3ds
tf.summary.scalar('iou_2d', tf.reduce_mean(iou2ds))
tf.summary.scalar('iou_3d', tf.reduce_mean(iou3ds))
correct = tf.equal(tf.argmax(end_points['mask_logits'], 2),
tf.to_int64(pls['seg_labels']))
accuracy = tf.reduce_sum(tf.cast(correct, tf.float32)) / \
float(BATCH_SIZE*NUM_POINT)
tf.summary.scalar('segmentation accuracy', accuracy)
# Get training operator
learning_rate = get_learning_rate(batch)
tf.summary.scalar('learning_rate', learning_rate)
if OPTIMIZER == 'momentum':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM)
elif OPTIMIZER == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
if FLAGS.train_cls_only:
# var_list = tf.trainable_variables('cls_')
var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='proposal_classification')
loss = loss_endpoints['cls_loss']
elif FLAGS.train_reg_only:
# tvars = tf.trainable_variables()
# var_list = [var for var in tvars if 'cls_' not in var.name]
var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='box_regression')
loss = loss_endpoints['box_loss']
else:
var_list = tf.trainable_variables()
# Note: when training, the moving_mean and moving_variance need to be updated.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss,
global_step=batch,
var_list=var_list)
# Add ops to save and restore all the variables.
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.Session(config=config)
# Add summary writers
merged = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'),
sess.graph)
test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'),
sess.graph)
# Init variables
if FLAGS.restore_model_path is None:
init = tf.global_variables_initializer()
sess.run(init)
else:
saver.restore(sess, FLAGS.restore_model_path)
ops = {
'pointclouds_pl': pls['pointclouds'],
'img_seg_map_pl': pls['img_seg_map'],
'prop_box_pl': pls['prop_box'],
'calib_pl': pls['calib'],
'cls_label_pl': pls['cls_label'],
'ious_pl': pls['ious'],
'labels_pl': pls['seg_labels'],
'centers_pl': pls['centers'],
'heading_class_label_pl': pls['heading_class_label'],
'heading_residual_label_pl': pls['heading_residual_label'],
'size_class_label_pl': pls['size_class_label'],
'size_residual_label_pl': pls['size_residual_label'],
'is_training_pl': pls['is_training'],
'logits': end_points['mask_logits'],
'cls_logits': end_points['cls_logits'],
'centers_pred': end_points['center'],
'loss': loss,
'loss_endpoints': loss_endpoints,
'train_op': train_op,
'merged': merged,
'step': batch,
'end_points': end_points
}
for epoch in range(MAX_EPOCH):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
train_one_epoch(sess, ops, train_writer)
if epoch >= 10 or True:
val_loss, avg_cls_acc, estimate_acc = eval_one_epoch(
sess, ops, test_writer)
# Save the variables to disk.
# if val_loss < best_val_loss:
# best_val_loss = val_loss
# save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt"))
# log_string("Model saved in file: {0}, val_loss: {1}".format(save_path, val_loss))
save_path = saver.save(
sess, os.path.join(LOG_DIR, "model.ckpt.%03d" % epoch))
log_string("Model saved in file: {0}".format(save_path))
train_loading_thread.stop()
val_loading_thread.stop()