def test(raw_dirs, TEST_DATASET):
log_string(FLAGS.log_dir)
with tf.Graph().as_default():
with tf.device('/gpu:0'):
input_pls = model.placeholder_inputs(scope='inputs_pl',
FLAGS=FLAGS,
num_pnts=NUM_INPUT_POINTS)
is_training_pl = tf.compat.v1.placeholder(tf.bool, shape=())
print(is_training_pl)
batch = tf.Variable(0, name='batch')
print("--- Get model and loss")
# Get model and loss
end_points = model.get_model(input_pls,
is_training_pl,
bn=False,
FLAGS=FLAGS)
loss, end_points = model.get_loss(end_points, FLAGS=FLAGS)
gpu_options = tf.compat.v1.GPUOptions(
) # (per_process_gpu_memory_fraction=0.99)
config = tf.compat.v1.ConfigProto(gpu_options=gpu_options)
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.compat.v1.Session(config=config)
##### all
# Init variables
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
######### Loading Checkpoint ###############
# Overall
saver = tf.compat.v1.train.Saver([
v for v in tf.compat.v1.get_collection_ref(
tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
if ('lr' not in v.name) and ('batch' not in v.name)
])
###########################################
ops = {
'input_pls': input_pls,
'is_training_pl': is_training_pl,
'loss': loss,
'step': batch,
'saver': saver,
'end_points': end_points
}
sys.stdout.flush()
TEST_DATASET.start()
log_string('**** INFERENCE ****')
sys.stdout.flush()
test_one_epoch(sess, ops, raw_dirs, TEST_DATASET)
TEST_DATASET.shutdown()
print("Done!")
def create():
log_string(LOG_DIR)
batch_data = read_img_get_transmat()
input_pls = model.placeholder_inputs(BATCH_SIZE, NUM_POINTS, (IMG_SIZE, IMG_SIZE),
num_sample_pc=NUM_SAMPLE_POINTS, scope='inputs_pl', FLAGS=FLAGS)
is_training_pl = tf.placeholder(tf.bool, shape=())
print(is_training_pl)
batch = tf.Variable(0, name='batch')
print("--- Get model and loss")
# Get model and loss
end_points = model.get_model(input_pls, NUM_POINTS, is_training_pl, bn=False,FLAGS=FLAGS)
loss, end_points = model.get_loss(end_points,
sdf_weight=SDF_WEIGHT, num_sample_points=NUM_SAMPLE_POINTS, FLAGS=FLAGS)
# Create a session
gpu_options = tf.GPUOptions() # per_process_gpu_memory_fraction=0.99
config = tf.ConfigProto(gpu_options=gpu_options)
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
######### Loading Checkpoint ###############
saver = tf.train.Saver([v for v in tf.get_collection_ref(tf.GraphKeys.GLOBAL_VARIABLES) if
('lr' not in v.name) and ('batch' not in v.name)])
ckptstate = tf.train.get_checkpoint_state(PRETRAINED_MODEL_PATH)
if ckptstate is not None:
LOAD_MODEL_FILE = os.path.join(PRETRAINED_MODEL_PATH, os.path.basename(ckptstate.model_checkpoint_path))
try:
# load_model(sess, PRETRAINED_PN_MODEL_FILE, ['refpc_reconstruction','sdfprediction','vgg_16'], strict=True)
with NoStdStreams():
saver.restore(sess, LOAD_MODEL_FILE)
print("Model loaded in file: %s" % LOAD_MODEL_FILE)
except:
print("Fail to load overall modelfile: %s" % PRETRAINED_MODEL_PATH)
###########################################
ops = {'input_pls': input_pls,
'is_training_pl': is_training_pl,
'loss': loss,
'step': batch,
'end_points': end_points}
test_one_epoch(sess, ops, batch_data)
def train():
log_string(LOG_DIR)
with tf.Graph().as_default():
with tf.device('/gpu:0'):
input_pls = model.placeholder_inputs(
BATCH_SIZE,
NUM_POINTS, (IMG_SIZE, IMG_SIZE),
num_sample_pc=NUM_SAMPLE_POINTS,
scope='inputs_pl',
FLAGS=FLAGS)
is_training_pl = tf.placeholder(tf.bool, shape=())
print(is_training_pl)
# 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, name='batch')
bn_decay = get_bn_decay(batch)
# tf.summary.scalar('bn_decay', bn_decay)
print("--- Get model and loss")
# Get model and loss
end_points = model.get_model(input_pls,
NUM_POINTS,
is_training_pl,
bn=False,
FLAGS=FLAGS)
loss, end_points = model.get_loss(
end_points,
sdf_weight=SDF_WEIGHT,
mask_weight=FLAGS.mask_weight,
num_sample_points=FLAGS.num_sample_points,
FLAGS=FLAGS)
# tf.summary.scalar('loss', loss)
print("--- Get training operator")
# Get training operator
learning_rate = get_learning_rate(batch)
if OPTIMIZER == 'momentum':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM)
elif OPTIMIZER == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate,
beta1=FLAGS.beta1)
# Create a session
config = tf.ConfigProto()
gpu_options = tf.GPUOptions(
) #(per_process_gpu_memory_fraction=0.99)
config = tf.ConfigProto(gpu_options=gpu_options)
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.Session(config=config)
merged = None
train_writer = None
##### all
update_variables = [
x for x in tf.get_collection_ref(tf.GraphKeys.GLOBAL_VARIABLES)
]
train_op = optimizer.minimize(loss,
global_step=batch,
var_list=update_variables)
# Init variables
init = tf.global_variables_initializer()
sess.run(init)
######### Loading Checkpoint ###############
# CNN(Pretrained from ImageNet)
if PRETRAINED_CNN_MODEL_FILE is not '':
if not load_model(
sess, PRETRAINED_CNN_MODEL_FILE, 'vgg_16',
strict=True):
return
if PRETRAINED_PN_MODEL_FILE is not '':
if not load_model(sess,
PRETRAINED_PN_MODEL_FILE,
['refpc_reconstruction', 'sdfprediction'],
strict=True):
return
# Overall
saver = tf.train.Saver([
v for v in tf.get_collection_ref(tf.GraphKeys.GLOBAL_VARIABLES)
if ('lr' not in v.name) and ('batch' not in v.name)
])
ckptstate = tf.train.get_checkpoint_state(PRETRAINED_MODEL_PATH)
if ckptstate is not None:
LOAD_MODEL_FILE = os.path.join(
PRETRAINED_MODEL_PATH,
os.path.basename(ckptstate.model_checkpoint_path))
try:
load_model(sess,
LOAD_MODEL_FILE, [
'sdfprediction/fold1',
'sdfprediction/fold2', 'vgg_16'
],
strict=True)
# load_model(sess, LOAD_MODEL_FILE, ['sdfprediction','vgg_16'], strict=True)
with NoStdStreams():
saver.restore(sess, LOAD_MODEL_FILE)
print("Model loaded in file: %s" % LOAD_MODEL_FILE)
except:
print("Fail to load overall modelfile: %s" %
PRETRAINED_MODEL_PATH)
###########################################
ops = {
'input_pls': input_pls,
'is_training_pl': is_training_pl,
'loss': loss,
'train_op': train_op,
'merged': merged,
'step': batch,
'lr': learning_rate,
'end_points': end_points
}
best_loss = 1e20
TRAIN_DATASET.start()
best_acc = 0
for epoch in range(MAX_EPOCH):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
avg_accuracy = train_one_epoch(sess, ops, train_writer, saver)
# Save the variables to disk.
if avg_accuracy > best_acc:
best_acc = avg_accuracy
save_path = saver.save(sess,
os.path.join(LOG_DIR, "model.ckpt"))
log_string("best Model saved in file: %s" % save_path)
elif epoch % 10 == 0:
save_path = saver.save(
sess,
os.path.join(LOG_DIR,
"model_epoch_%03d.ckpt" % (epoch)))
log_string("Model saved in file: %s" % save_path)
TRAIN_DATASET.shutdown()
def train():
log_string(FLAGS.log_dir)
with tf.Graph().as_default():
with tf.device('/gpu:{}'.format(0)):
input_pls = model.placeholder_inputs(scope='inputs_pl',
FLAGS=FLAGS)
is_training_pl = tf.compat.v1.placeholder(tf.bool, shape=())
print(is_training_pl)
# 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, name='batch')
bn_decay = get_bn_decay(batch)
# tf.summary.scalar('bn_decay', bn_decay)
print("--- Get model and loss")
# Get model and loss
end_points = model.get_model(input_pls,
is_training_pl,
bn=FLAGS.bn,
bn_decay=bn_decay,
FLAGS=FLAGS)
loss, end_points = model.get_loss(end_points, FLAGS=FLAGS)
# tf.summary.scalar('loss', loss)
print("--- Get training operator")
# Get training operator
learning_rate = get_learning_rate(batch)
if FLAGS.optimizer == 'momentum':
optimizer = tf.compat.v1.train.MomentumOptimizer(
learning_rate, momentum=FLAGS.momentum)
elif FLAGS.optimizer == 'adam':
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate,
beta1=FLAGS.beta1)
# Create a session
gpu_options = tf.compat.v1.GPUOptions(
) #(per_process_gpu_memory_fraction=0.99)
config = tf.compat.v1.ConfigProto(gpu_options=gpu_options)
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.compat.v1.Session(config=config)
merged = None
train_writer = None
##### all
update_variables = [
x for x in tf.compat.v1.get_collection_ref(
tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
]
train_op = optimizer.minimize(loss,
global_step=batch,
var_list=update_variables)
# Init variables
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
######### Loading Checkpoint ###############
# CNN(Pretrained from ImageNet)
if FLAGS.restore_modelcnn is not '':
if not load_model(
sess, FLAGS.restore_modelcnn, FLAGS.encoder,
strict=True):
return
# Overall
saver = tf.compat.v1.train.Saver([
v for v in tf.compat.v1.get_collection_ref(
tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
if ('lr' not in v.name) and ('batch' not in v.name)
])
ckptstate = tf.train.get_checkpoint_state(FLAGS.restore_model)
if ckptstate is not None:
LOAD_MODEL_FILE = os.path.join(
FLAGS.restore_model,
os.path.basename(ckptstate.all_model_checkpoint_paths[0]))
load_model_all(saver, sess, LOAD_MODEL_FILE)
print("Model loaded in file: %s" % LOAD_MODEL_FILE)
# try:
# load_model(sess, LOAD_MODEL_FILE, ['sdfprediction/fold1', 'sdfprediction/fold2', FLAGS.encoder],
# strict=True)
# # load_model(sess, LOAD_MODEL_FILE, ['sdfprediction','vgg_16'], strict=True)
# with NoStdStreams():
# saver.restore(sess, LOAD_MODEL_FILE)
# print("Model loaded in file: %s" % LOAD_MODEL_FILE)
# except:
# print("Fail to load overall modelfile: %s" % FLAGS.restore_model)
###########################################
ops = {
'input_pls': input_pls,
'is_training_pl': is_training_pl,
'loss': loss,
'train_op': train_op,
'merged': merged,
'step': batch,
'lr': learning_rate,
'end_points': end_points
}
TRAIN_DATASET.start()
TEST_DATASET.start()
best_locnorm_diff, best_dir_diff = 10000, 10000
for epoch in range(FLAGS.max_epoch):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
if epoch == 0 and FLAGS.restore_model:
test_one_epoch(sess, ops, epoch)
# test_one_epoch(sess, ops, epoch)
xyz_avg_diff, _, _ = train_one_epoch(sess, ops, epoch)
if epoch % 5 == 0 and epoch > 1:
locnorm_avg_diff, direction_avg_diff = test_one_epoch(
sess, ops, epoch)
# Save the variables to disk.
if locnorm_avg_diff < best_locnorm_diff:
best_locnorm_diff = locnorm_avg_diff
save_path = saver.save(
sess, os.path.join(FLAGS.log_dir, "model.ckpt"))
log_string(
"best locnorm_avg_diff Model saved in file: %s" %
save_path)
# elif direction_avg_diff < best_dir_diff:
# best_dir_diff = direction_avg_diff
# save_path = saver.save(sess, os.path.join(FLAGS.log_dir, "dir_model.ckpt"))
# log_string("best direction Model saved in file: %s" % save_path)
if epoch % 30 == 0 and epoch > 1:
save_path = saver.save(
sess,
os.path.join(FLAGS.log_dir,
"model_epoch_%03d.ckpt" % (epoch)))
log_string("Model saved in file: %s" % save_path)
TRAIN_DATASET.shutdown()
TEST_DATASET.shutdown()
