import argparse import subprocess import tensorflow as tf import numpy as np from datetime import datetime import json import os import sys import datetime import time import collections BASE_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(BASE_DIR) import P2PNET import ioUtil # DEFAULT SETTINGS parser = argparse.ArgumentParser() parser.add_argument('--train_hdf5', default='training data file name(*.hdf5)' ) parser.add_argument('--test_hdf5', default='test data file name(*.hdf5)' ) parser.add_argument('--domain_A', default='skeleton', help='name of domain A') parser.add_argument('--domain_B', default='surface', help='name of domain B') parser.add_argument('--mode', type=str, default='train', help='train or test') parser.add_argument('--gpu', type=int, default=0, help='which GPU to use [default: 0]') parser.add_argument('--batch_size', type=int, default=4, help='Batch Size during training [default: 4]') parser.add_argument('--epoch', type=int, default=200, help='number of epoches to run [default: 200]') parser.add_argument('--decayEpoch', type=int, default=50, help='steps(how many epoches) for decaying learning rate') parser.add_argument("--densityWeight", type=float, default=1.0, help="density weight [default: 1.0]") parser.add_argument("--regularWeight", type=float, default=0.1, help="regularization weight [default: 0.1]") parser.add_argument("--nnk", type=int, default=8, help="density: number of nearest neighbours [default: 8]") parser.add_argument("--range_max", type=float, default=1.0, help="max length of point displacement[default: 1.0]") parser.add_argument("--radiusScal", type=float, default=1.0, help="a constant for scaling radii in pointnet++ [default: 1.0]") parser.add_argument("--noiseLength", type=int, default=32, help="length of point-wise noise vector [default: 32]") parser.add_argument('--checkpoint', default=None, help='epoch_##.ckpt') ### None None None parser.add_argument('--point_num', type=int, default=None, help='do not set the argument') parser.add_argument('--example_num', type=int, default=None, help='do not set the argument') parser.add_argument('--output_dir', type=str, default=None, help='do not set the argument') FLAGS = parser.parse_args() Train_examples = ioUtil.load_examples(FLAGS.train_hdf5, FLAGS.domain_A, FLAGS.domain_B, 'names') Test_examples = ioUtil.load_examples(FLAGS.test_hdf5, FLAGS.domain_A, FLAGS.domain_B, 'names') FLAGS.point_num = Train_examples.pointsets_A.shape[1] POINT_NUM = FLAGS.point_num Example_NUM = Train_examples.pointsets_A.shape[0] FLAGS.example_num = Example_NUM TRAINING_EPOCHES = FLAGS.epoch batch_size = FLAGS.batch_size if Train_examples.pointsets_B.shape[1] != POINT_NUM \ or Test_examples.pointsets_A.shape[1] != POINT_NUM \ or Test_examples.pointsets_B.shape[1] != POINT_NUM : print( 'point number inconsistent in the data set.') exit() ########## create output folders datapath, basefname = os.path.split( FLAGS.train_hdf5 ) output_dir = 'output_' + basefname[0:basefname.index('_')] + '_' + FLAGS.domain_A + '-' + FLAGS.domain_B ## + '_noise' + str(FLAGS.noiseLength) + '_dw' + str(FLAGS.densityWeight)+ '_rw' + str(FLAGS.regularWeight) if not os.path.exists(output_dir): os.mkdir(output_dir) MODEL_STORAGE_PATH = os.path.join(output_dir, 'trained_models') if not os.path.exists(MODEL_STORAGE_PATH): os.mkdir(MODEL_STORAGE_PATH) SUMMARIES_FOLDER = os.path.join(output_dir, 'summaries') if not os.path.exists(SUMMARIES_FOLDER): os.mkdir(SUMMARIES_FOLDER) ########## Save test input ioUtil.output_point_cloud_ply( Test_examples.pointsets_A, Test_examples.names, output_dir, 'gt_'+FLAGS.domain_A) ioUtil.output_point_cloud_ply( Test_examples.pointsets_B, Test_examples.names, output_dir, 'gt_'+FLAGS.domain_B) # print arguments for k, v in FLAGS._get_kwargs(): print(k + ' = ' + str(v) ) def train(): with tf.Graph().as_default(): with tf.device('/gpu:' + str(FLAGS.gpu)): model = P2PNET.create_model(FLAGS) ########## Init and Configuration ########## saver = tf.train.Saver( max_to_keep=5 ) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) # Restore variables from disk. Start_epoch_number = 1 if FLAGS.checkpoint is not None: print('load checkpoint: ' + FLAGS.checkpoint ) saver.restore(sess, FLAGS.checkpoint ) fname = os.path.basename( FLAGS.checkpoint ) Start_epoch_number = int( fname[6:-5] ) + 1 print( 'Start_epoch_number = ' + str(Start_epoch_number) ) train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train', sess.graph) test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test') fcmd = open(os.path.join(output_dir, 'arguments.txt'), 'w') fcmd.write(str(FLAGS)) fcmd.close() ########## Training one epoch ########## def train_one_epoch(epoch_num): now = datetime.datetime.now() print(now.strftime("%Y-%m-%d %H:%M:%S")) start_time = time.time() is_training = True Train_examples_shuffled = ioUtil.shuffle_examples(Train_examples) pointsets_A = Train_examples_shuffled.pointsets_A pointsets_B = Train_examples_shuffled.pointsets_B names = Train_examples_shuffled.names num_data = pointsets_A.shape[0] num_batch = num_data // batch_size total_data_loss_A = 0.0 total_shape_loss_A = 0.0 total_density_loss_A = 0.0 total_data_loss_B = 0.0 total_shape_loss_B = 0.0 total_density_loss_B = 0.0 total_reg_loss = 0.0 for j in range(num_batch): begidx = j * batch_size endidx = (j + 1) * batch_size feed_dict = { model.pointSet_A_ph: pointsets_A[begidx: endidx, ...], model.pointSet_B_ph: pointsets_B[begidx: endidx, ...], model.is_training_ph: is_training, } fetches = { "train": model.total_train, "shapeLoss_A": model.shapeLoss_A, "densityLoss_A": model.densityLoss_A, "shapeLoss_B": model.shapeLoss_B, "densityLoss_B": model.densityLoss_B, "data_loss_A": model.data_loss_A, "data_loss_B": model.data_loss_B, "regul_loss": model.regul_loss, "learning_rate": model.learning_rate, "global_step": model.global_step, } results = sess.run(fetches, feed_dict=feed_dict) total_data_loss_A += results["data_loss_A"] total_shape_loss_A += results["shapeLoss_A"] total_density_loss_A += results["densityLoss_A"] total_data_loss_B += results["data_loss_B"] total_shape_loss_B += results["shapeLoss_B"] total_density_loss_B += results["densityLoss_B"] total_reg_loss += results["regul_loss"] if j % 50 == 0: print(' ' + str(j) + '/' + str(num_batch) + ': ' ) print(' data_loss_A = {:.4f},'.format(results["data_loss_A"] ) + \ ' shape = {:.4f},'.format(results["shapeLoss_A"] ) + \ ' density = {:.4f}'.format(results["densityLoss_A"] ) ) print(' data_loss_B = {:.4f},'.format(results["data_loss_B"] ) + \ ' shape = {:.4f},'.format(results["shapeLoss_B"] ) + \ ' density = {:.4f}'.format(results["densityLoss_B"] ) ) print(' regul_loss = {:.4f}\n'.format(results["regul_loss"] ) ) print(' learning_rate = {:.6f}'.format(results["learning_rate"] ) ) print(' global_step = {0}'.format(results["global_step"] ) ) total_data_loss_A /= num_batch total_shape_loss_A /= num_batch total_density_loss_A /= num_batch total_data_loss_B /= num_batch total_shape_loss_B /= num_batch total_density_loss_B /= num_batch total_reg_loss /= num_batch # evaluate summaries training_sum = sess.run( model.training_sum_ops, \ feed_dict={model.train_dataloss_A_ph: total_data_loss_A, \ model.train_dataloss_B_ph: total_data_loss_B, \ model.train_regul_ph: total_reg_loss,\ }) train_writer.add_summary(training_sum, epoch_num) print( '\tData_loss_A = %.4f,' % total_data_loss_A + \ ' shape = %.4f,' % total_shape_loss_A + \ ' density = %.4f' % total_density_loss_A ) print( '\tData_loss_B = %.4f,' % total_data_loss_B + \ ' shape = %.4f,' % total_shape_loss_B + \ ' density = %.4f' % total_density_loss_B ) print( '\tReg_loss: %.4f\n' % total_reg_loss) elapsed_time = time.time() - start_time print( '\tply/sec:' + str( round(num_data/elapsed_time) ) ) print( '\tduration of this epoch:' + str(round(elapsed_time/60) ) + ' min' ) print( '\testimated finishing time:' + str(round(elapsed_time/60.0 * (TRAINING_EPOCHES-epoch_num-1)) ) + ' min' ) ################## end of train function #################### end of train function ########## def eval_one_epoch(epoch_num, mustSavePly=False): is_training = False pointsets_A = Test_examples.pointsets_A pointsets_B = Test_examples.pointsets_B names = Test_examples.names num_data = pointsets_A.shape[0] num_batch = num_data // batch_size total_data_loss_A = 0.0 total_shape_loss_A = 0.0 total_density_loss_A = 0.0 total_data_loss_B = 0.0 total_shape_loss_B = 0.0 total_density_loss_B = 0.0 total_reg_loss = 0.0 for j in range(num_batch): begidx = j * batch_size endidx = (j + 1) * batch_size feed_dict = { model.pointSet_A_ph: pointsets_A[begidx: endidx, ...], model.pointSet_B_ph: pointsets_B[begidx: endidx, ...], model.is_training_ph: is_training, } fetches = { "shapeLoss_A": model.shapeLoss_A, "densityLoss_A": model.densityLoss_A, "shapeLoss_B": model.shapeLoss_B, "densityLoss_B": model.densityLoss_B, "data_loss_A": model.data_loss_A, "data_loss_B": model.data_loss_B, "regul_loss": model.regul_loss, "Predicted_A": model.Predicted_A, "Predicted_B": model.Predicted_B, } results = sess.run(fetches, feed_dict=feed_dict) total_data_loss_A += results["data_loss_A"] total_shape_loss_A += results["shapeLoss_A"] total_density_loss_A += results["densityLoss_A"] total_data_loss_B += results["data_loss_B"] total_shape_loss_B += results["shapeLoss_B"] total_density_loss_B += results["densityLoss_B"] total_reg_loss += results["regul_loss"] # write test results if epoch_num % 20 == 0 or mustSavePly: # save predicted point sets with 1 single feeding pass nametosave = names[begidx: endidx, ...] Predicted_A_xyz = np.squeeze(np.array(results["Predicted_A"])) Predicted_B_xyz = np.squeeze(np.array(results["Predicted_B"])) ioUtil.output_point_cloud_ply(Predicted_A_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X1') ioUtil.output_point_cloud_ply(Predicted_B_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X1') # save predicted point sets with 4 feeding passes for i in range(3): results = sess.run(fetches, feed_dict=feed_dict) Predicted_A_xyz__ = np.squeeze(np.array(results["Predicted_A"])) Predicted_B_xyz__ = np.squeeze(np.array(results["Predicted_B"])) Predicted_A_xyz = np.concatenate((Predicted_A_xyz, Predicted_A_xyz__), axis=1) Predicted_B_xyz = np.concatenate((Predicted_B_xyz, Predicted_B_xyz__), axis=1) ioUtil.output_point_cloud_ply(Predicted_A_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X4') ioUtil.output_point_cloud_ply(Predicted_B_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X4') # save predicted point sets with 8 feeding passes for i in range(4): results = sess.run(fetches, feed_dict=feed_dict) Predicted_A_xyz__ = np.squeeze(np.array(results["Predicted_A"])) Predicted_B_xyz__ = np.squeeze(np.array(results["Predicted_B"])) Predicted_A_xyz = np.concatenate((Predicted_A_xyz, Predicted_A_xyz__), axis=1) Predicted_B_xyz = np.concatenate((Predicted_B_xyz, Predicted_B_xyz__), axis=1) ioUtil.output_point_cloud_ply( Predicted_A_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X8') ioUtil.output_point_cloud_ply( Predicted_B_xyz, nametosave, output_dir, 'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X8') total_data_loss_A /= num_batch total_shape_loss_A /= num_batch total_density_loss_A /= num_batch total_data_loss_B /= num_batch total_shape_loss_B /= num_batch total_density_loss_B /= num_batch total_reg_loss /= num_batch # evaluate summaries testing_sum = sess.run( model.testing_sum_ops, \ feed_dict={model.test_dataloss_A_ph: total_data_loss_A, \ model.test_dataloss_B_ph: total_data_loss_B, \ model.test_regul_ph: total_reg_loss,\ }) test_writer.add_summary(testing_sum, epoch_num) print('\tData_loss_A = %.4f,' % total_data_loss_A + \ ' shape = %.4f,' % total_shape_loss_A + \ ' density = %.4f' % total_density_loss_A) print('\tData_loss_B = %.4f,' % total_data_loss_B + \ ' shape = %.4f,' % total_shape_loss_B + \ ' density = %.4f' % total_density_loss_B) print('\tReg_loss: %.4f\n' % total_reg_loss) ################## end of test function #################### end of test function ########## if not os.path.exists(MODEL_STORAGE_PATH): os.mkdir(MODEL_STORAGE_PATH) if FLAGS.mode=='train': for epoch in range(Start_epoch_number, TRAINING_EPOCHES+1): print( '\n>>> Training for the epoch %d/%d ...' % (epoch, TRAINING_EPOCHES)) train_one_epoch(epoch) if epoch % 20 == 0: cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch) + '.ckpt')) print( 'Successfully store the checkpoint model into ' + cp_filename) print('\n<<< Testing on the test dataset...') eval_one_epoch(epoch, mustSavePly=True) else: print( '\n<<< Testing on the test dataset ...') eval_one_epoch(Start_epoch_number, mustSavePly=True) if __name__ == '__main__': train()