def eval():
with tf.Graph().as_default():
with tf.device('/gpu:' + str(FLAGS.gpu)):
pointclouds_pl, labels_pl = MODEL.placeholder_inputs(
BATCH_SIZE, NUM_POINT, NFEATURES)
batch = tf.Variable(0, trainable=False)
alpha = tf.placeholder(tf.float32, shape=())
is_training_pl = tf.placeholder(tf.bool, shape=())
pred, max_pool = MODEL.get_model(pointclouds_pl,
is_training=is_training_pl,
num_class=NUM_CATEGORIES)
mu = tf.Variable(tf.zeros(shape=(FLAGS.n_clusters, FLAGS.max_dim)),
name="mu",
trainable=False) #k centroids
classify_loss = MODEL.get_focal_loss(pred, labels_pl,
NUM_CATEGORIES)
kmeans_loss, stack_dist = MODEL.get_loss_kmeans(
max_pool, mu, FLAGS.max_dim, FLAGS.n_clusters, alpha)
saver = tf.train.Saver()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
if FULL_TRAINING:
saver.restore(sess, os.path.join(LOG_DIR, 'cluster.ckpt'))
else:
saver.restore(sess, os.path.join(LOG_DIR, 'model.ckpt'))
print('model restored')
ops = {
'pointclouds_pl': pointclouds_pl,
'labels_pl': labels_pl,
'stack_dist': stack_dist,
'kmeans_loss': kmeans_loss,
'pred': pred,
'alpha': alpha,
'max_pool': max_pool,
'is_training_pl': is_training_pl,
'classify_loss': classify_loss,
}
eval_one_epoch(sess, ops)
def train():
with tf.Graph().as_default():
with tf.device('/gpu:'+str(GPU_INDEX)):
pointclouds_pl, labels_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT,NUM_FEAT)
is_training_pl = tf.placeholder(tf.bool, shape=())
# Note the global_step=batch parameter to minimize.
# That tells the optimizer to helpfully increment the 'batch' parameter for you every time it trains.
batch = tf.Variable(0)
alpha = tf.placeholder(dtype=tf.float32, shape=())
bn_decay = get_bn_decay(batch)
tf.summary.scalar('bn_decay', bn_decay)
print("--- Get model and loss")
pred , max_pool = MODEL.get_model(pointclouds_pl, is_training=is_training_pl,
bn_decay=bn_decay,
num_class=NUM_CLASSES, weight_decay=FLAGS.wd,
)
class_loss = MODEL.get_focal_loss(pred, labels_pl,NUM_CLASSES)
mu = tf.Variable(tf.zeros(shape=(FLAGS.n_clusters,FLAGS.max_dim)),name="mu",trainable=True) #k centroids
kmeans_loss, stack_dist= MODEL.get_loss_kmeans(max_pool,mu, FLAGS.max_dim,
FLAGS.n_clusters,alpha)
full_loss = kmeans_loss + class_loss
print("--- Get training operator")
# 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)
train_op_full = optimizer.minimize(full_loss, global_step=batch)
train_op = optimizer.minimize(class_loss, global_step=batch)
# 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)
sess.run(tf.global_variables_initializer())
# 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
print("Total number of weights for the model: ",np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]))
ops = {'pointclouds_pl': pointclouds_pl,
'labels_pl':labels_pl,
'is_training_pl': is_training_pl,
'max_pool':max_pool,
'pred': pred,
'alpha': alpha,
'mu': mu,
'stack_dist':stack_dist,
'class_loss': class_loss,
'kmeans_loss': kmeans_loss,
'train_op': train_op,
'train_op_full': train_op_full,
'merged': merged,
'step': batch,
'learning_rate':learning_rate
}
for epoch in range(MAX_EPOCH):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
is_full_training = epoch > MAX_PRETRAIN
max_pool = train_one_epoch(sess, ops, train_writer,is_full_training)
if epoch == MAX_PRETRAIN:
centers = KMeans(n_clusters=FLAGS.n_clusters).fit(np.squeeze(max_pool))
centers = centers.cluster_centers_
sess.run(tf.assign(mu,centers))
eval_one_epoch(sess, ops, test_writer,is_full_training)
if is_full_training:
save_path = saver.save(sess, os.path.join(LOG_DIR, 'cluster.ckpt'))
else:
save_path = saver.save(sess, os.path.join(LOG_DIR, 'model.ckpt'))
log_string("Model saved in file: %s" % save_path)