Tensorflow1.15 pointnet sem_seg

早速修正した、コードを掲載します。
sem_seg/train_py
import argparse import math import h5py import numpy as np import tensorflow as tf import socket import os import sys BASE_DIR = os.path.dirname(os.path.abspath(__file__)) ROOT_DIR = os.path.dirname(BASE_DIR) sys.path.append(BASE_DIR) sys.path.append(ROOT_DIR) sys.path.append(os.path.join(ROOT_DIR, 'utils')) import provider import tf_util from model import * parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') parser.add_argument('--log_dir', default='log', help='Log dir [default: log]') parser.add_argument('--num_point', type=int, default=4096, help='Point number [default: 4096]') ##parser.add_argument('--max_epoch', type=int, default=50, help='Epoch to run [default: 50]') parser.add_argument('--max_epoch', type=int, default=4, help='Epoch to run [default: 50]') parser.add_argument('--batch_size', type=int, default=24, help='Batch Size during training [default: 24]') parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]') parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]') parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]') parser.add_argument('--decay_step', type=int, default=300000, help='Decay step for lr decay [default: 300000]') parser.add_argument('--decay_rate', type=float, default=0.5, help='Decay rate for lr decay [default: 0.5]') parser.add_argument('--test_area', type=int, default=6, help='Which area to use for test, option: 1-6 [default: 6]') FLAGS = parser.parse_args() BATCH_SIZE = FLAGS.batch_size NUM_POINT = FLAGS.num_point MAX_EPOCH = FLAGS.max_epoch NUM_POINT = FLAGS.num_point BASE_LEARNING_RATE = FLAGS.learning_rate GPU_INDEX = FLAGS.gpu MOMENTUM = FLAGS.momentum OPTIMIZER = FLAGS.optimizer DECAY_STEP = FLAGS.decay_step DECAY_RATE = FLAGS.decay_rate LOG_DIR = FLAGS.log_dir if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) #os.system('cp model.py %s' % (LOG_DIR)) # bkp of model def #os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure MODEL_STORAGE_PATH = os.path.join('.', 'trained_models') if not os.path.exists(MODEL_STORAGE_PATH): os.mkdir(MODEL_STORAGE_PATH) # changed by nishi if False: LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w') LOG_FOUT.write(str(FLAGS)+'\n') else: LOG_FOUT=None MAX_NUM_POINT = 4096 NUM_CLASSES = 13 BN_INIT_DECAY = 0.5 BN_DECAY_DECAY_RATE = 0.5 #BN_DECAY_DECAY_STEP = float(DECAY_STEP * 2) BN_DECAY_DECAY_STEP = float(DECAY_STEP) BN_DECAY_CLIP = 0.99 HOSTNAME = socket.gethostname() data_dir = 'E:/tmp/pointnet/sem_seg/' #hdf5_data_dir = os.path.join(BASE_DIR, './indoor3d_sem_seg_hdf5_data') hdf5_data_dir = data_dir +'indoor3d_sem_seg_hdf5_data' ALL_FILES = provider.getDataFiles(hdf5_data_dir+'/all_files.txt') room_filelist = [line.rstrip() for line in open(hdf5_data_dir+'/room_filelist.txt')] # Load ALL data data_batch_list = [] label_batch_list = [] for h5_filename in ALL_FILES: print('h5_filename=',h5_filename) data_batch, label_batch = provider.loadDataFile(data_dir+h5_filename) data_batch_list.append(data_batch) label_batch_list.append(label_batch) data_batches = np.concatenate(data_batch_list, 0) label_batches = np.concatenate(label_batch_list, 0) print('>data_batches.shape=',data_batches.shape) print('>label_batches.shape=',label_batches.shape) test_area = 'Area_'+str(FLAGS.test_area) train_idxs = [] test_idxs = [] for i,room_name in enumerate(room_filelist): if test_area in room_name: test_idxs.append(i) else: train_idxs.append(i) train_data = data_batches[train_idxs,...] train_label = label_batches[train_idxs] test_data = data_batches[test_idxs,...] test_label = label_batches[test_idxs] print('>train_data.shape=',train_data.shape,' train_label.shape=', train_label.shape) print('>test_data.shape=',test_data.shape,' test_label.shape=', test_label.shape) # add by nishi class Acc_Saver(): def __init__(self,path,thres=0.6): self.path=path self.acc_m=None self.thres=thres if os.path.exists(path): f = open(path) data = f.read() f.close() if data != None and data != 'None': self.acc_m=float(data) def check(self,acc): ok=False if self.acc_m == None: self.acc_m=acc ok=True elif acc > self.acc_m: self.acc_m = acc ok=True if ok == True and self.acc_m > self.thres: return True else: return False def save(self): f = open(self.path,'w') data = f.write(str(self.acc_m)) f.close() # add by nishi class Next_Epoch(): def __init__(self,path,start=0): self.path=path self.next_epoch_m=start self.steps=0 if os.path.exists(path): f = open(path) data = f.read() f.close() j = data.split(',') self.next_epoch_m=int(j[0]) self.steps=int(j[1]) def set(self,epoch): self.next_epoch_m=epoch def inc(self): self.next_epoch_m +=1 def get(self): return self.next_epoch_m def add_steps(self,step): self.steps += step def get_steps(self): return self.steps def save(self): f = open(self.path,'w') s=str(self.next_epoch_m)+','+str(self.steps) data = f.write(s) f.close() def log_string(out_str): if LOG_FOUT != None: LOG_FOUT.write(out_str+'\n') LOG_FOUT.flush() print(out_str) def get_learning_rate(batch): learning_rate = tf.compat.v1.train.exponential_decay( BASE_LEARNING_RATE, # Base learning rate. batch * BATCH_SIZE, # Current index into the dataset. DECAY_STEP, # Decay step. DECAY_RATE, # Decay rate. staircase=True) learning_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE!! return learning_rate def get_bn_decay(batch): bn_momentum = tf.compat.v1.train.exponential_decay( BN_INIT_DECAY, batch*BATCH_SIZE, BN_DECAY_DECAY_STEP, BN_DECAY_DECAY_RATE, staircase=True) bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum) return bn_decay # add by nishi next_epoch=Next_Epoch('trained_models/next_epoch.txt') START_EPOCH = next_epoch.get() print('### Start epoch: {0}'.format(START_EPOCH)) # make Acc_Saver add by nishi acc_saver=Acc_Saver('trained_models/eval_acc.txt',thres=0.3) def train(): with tf.Graph().as_default(): with tf.device('/gpu:'+str(GPU_INDEX)): pointclouds_pl, labels_pl = placeholder_inputs(BATCH_SIZE, NUM_POINT) is_training_pl = tf.compat.v1.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 is equal with test steps. by nishi # changed by nishi if START_EPOCH > 0: batch = tf.Variable(next_epoch.get_steps(), trainable=False) else: batch = tf.Variable(0, trainable=False) bn_decay = get_bn_decay(batch) # set to tensorboard summary bn_decay_tbs = tf.compat.v1.summary.scalar('bn_decay', bn_decay) # Get model and loss pred = get_model(pointclouds_pl, is_training_pl, bn_decay=bn_decay) loss = get_loss(pred, labels_pl) # set to tensorboard summary loss_tbs = tf.compat.v1.summary.scalar('loss', loss) correct = tf.equal(tf.argmax(pred, 2), tf.to_int64(labels_pl)) accuracy = tf.reduce_sum(tf.cast(correct, tf.float32)) / float(BATCH_SIZE*NUM_POINT) # set to tensorboard summary accuracy_tbs = tf.compat.v1.summary.scalar('accuracy', accuracy) # add by nishi test_acc = tf.compat.v1.placeholder(tf.float32, shape=()) # set to tensorboard summary test_acc_tbs = tf.compat.v1.summary.scalar('test_accuracy', test_acc) test_loss = tf.compat.v1.placeholder(tf.float32, shape=()) # set to tensorboard summary test_loss_tbs = tf.compat.v1.summary.scalar('test_loss', test_loss) # Get training operator learning_rate = get_learning_rate(batch) learning_rate_tbs = tf.compat.v1.summary.scalar('learning_rate', learning_rate) if OPTIMIZER == 'momentum': optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM) elif OPTIMIZER == 'adam': optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate) train_op = optimizer.minimize(loss, global_step=batch) # Add ops to save and restore all the variables. saver = tf.compat.v1.train.Saver() # Create a session config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True #config.log_device_placement = True # changed by nishi config.log_device_placement = False sess = tf.compat.v1.Session(config=config) # Add summary writers #merged = tf.compat.v1.summary.merge_all() train_writer = tf.compat.v1.summary.FileWriter(os.path.join(LOG_DIR, 'train'), sess.graph) test_writer = tf.compat.v1.summary.FileWriter(os.path.join(LOG_DIR, 'test')) # changed by nishi if START_EPOCH > 0: MODEL_PATH=os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(START_EPOCH-1)+'.ckpt') #MODEL_PATH="model_sav/model-%d.ckpt" % (START_EPOCH-1) saver.restore(sess, MODEL_PATH) #s="Model restored. %s" %(MODEL_PATH) log_string("Model restored. %s" %(MODEL_PATH)) else: # Init variables init = tf.compat.v1.global_variables_initializer() sess.run(init, {is_training_pl:True}) ops = {'pointclouds_pl': pointclouds_pl, 'labels_pl': labels_pl, 'is_training_pl': is_training_pl, 'pred': pred, 'loss': loss, 'train_op': train_op, #'merged': merged, 'bn_decay_tbs':bn_decay_tbs, 'loss_tbs':loss_tbs, 'accuracy_tbs':accuracy_tbs, 'learning_rate_tbs':learning_rate_tbs, 'test_acc':test_acc, 'test_acc_tbs':test_acc_tbs, 'test_loss':test_loss, 'test_loss_tbs':test_loss_tbs, 'step': batch} for epoch in range(START_EPOCH,START_EPOCH+MAX_EPOCH): log_string('**** EPOCH %03d ****' % (epoch)) sys.stdout.flush() train_one_epoch(sess, ops, train_writer) eval_acc = eval_one_epoch(sess, ops, test_writer,epoch) # Save the variables to disk. if epoch % 2 == 0 and epoch != START_EPOCH: save_path = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch)+'.ckpt')) log_string("Model saved in file: %s" % save_path) # add by nishi if acc_saver.check(eval_acc) == True: model_sav_path = "eval_acc_model-%d.ckpt" % epoch if not os.path.exists(os.path.join(MODEL_STORAGE_PATH, model_sav_path)): #save_path = saver.save(sess, os.path.join(MODEL_SAV_DIR, "model.ckpt")) save_path = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, model_sav_path)) s="eval_acc Model saved in file: %s" % save_path log_string(s) # add by nishi next_epoch.inc() save_path = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch)+'.ckpt')) log_string("Model saved in file: %s" % save_path) def train_one_epoch(sess, ops, train_writer): """ ops: dict mapping from string to tf ops """ is_training = True log_string('----') current_data, current_label, _ = provider.shuffle_data(train_data[:,0:NUM_POINT,:], train_label) file_size = current_data.shape[0] num_batches = file_size // BATCH_SIZE next_epoch.add_steps(num_batches) total_correct = 0 total_seen = 0 loss_sum = 0 for batch_idx in range(num_batches): if batch_idx % 100 == 0: print('Current batch/total batch num: %d/%d'%(batch_idx,num_batches)) start_idx = batch_idx * BATCH_SIZE end_idx = (batch_idx+1) * BATCH_SIZE feed_dict = {ops['pointclouds_pl']: current_data[start_idx:end_idx, :, :], ops['labels_pl']: current_label[start_idx:end_idx], ops['is_training_pl']: is_training} bn_decay_tb,loss_tb,accuracy_tb,learning_rate_tb, step, _, loss_val, pred_val = sess.run([ops['bn_decay_tbs'], ops['loss_tbs'], ops['accuracy_tbs'], ops['learning_rate_tbs'], ops['step'], ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict) #train_writer.add_summary(summary, step) train_writer.add_summary(bn_decay_tb, step) train_writer.add_summary(loss_tb, step) train_writer.add_summary(accuracy_tb, step) train_writer.add_summary(learning_rate_tb, step) pred_val = np.argmax(pred_val, 2) correct = np.sum(pred_val == current_label[start_idx:end_idx]) total_correct += correct total_seen += (BATCH_SIZE*NUM_POINT) loss_sum += loss_val log_string('mean loss: %f' % (loss_sum / float(num_batches))) log_string('accuracy: %f' % (total_correct / float(total_seen))) def eval_one_epoch(sess, ops, test_writer,epoch): """ ops: dict mapping from string to tf ops """ is_training = False total_correct = 0 total_seen = 0 loss_sum = 0 total_seen_class = [0 for _ in range(NUM_CLASSES)] total_correct_class = [0 for _ in range(NUM_CLASSES)] log_string('----') current_data = test_data[:,0:NUM_POINT,:] current_label = np.squeeze(test_label) file_size = current_data.shape[0] num_batches = file_size // BATCH_SIZE for batch_idx in range(num_batches): start_idx = batch_idx * BATCH_SIZE end_idx = (batch_idx+1) * BATCH_SIZE feed_dict = {ops['pointclouds_pl']: current_data[start_idx:end_idx, :, :], ops['labels_pl']: current_label[start_idx:end_idx], ops['is_training_pl']: is_training} step, loss_val, pred_val = sess.run([ops['step'], ops['loss'], ops['pred']], feed_dict=feed_dict) #test_writer.add_summary(summary, step) pred_val = np.argmax(pred_val, 2) correct = np.sum(pred_val == current_label[start_idx:end_idx]) total_correct += correct total_seen += (BATCH_SIZE*NUM_POINT) loss_sum += (loss_val*BATCH_SIZE) for i in range(start_idx, end_idx): for j in range(NUM_POINT): l = current_label[i, j] total_seen_class[l] += 1 total_correct_class[l] += (pred_val[i-start_idx, j] == l) #log_string('eval mean loss: %f' % (loss_sum / float(total_seen/NUM_POINT))) mean_loss=loss_sum / float(total_seen/NUM_POINT) log_string('eval mean loss: %f' % (mean_loss)) #log_string('eval accuracy: %f'% (total_correct / float(total_seen))) eval_acc = total_correct / float(total_seen) log_string('eval accuracy: %f'% (eval_acc)) avg_eval_acc = np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)) log_string('eval avg class acc: %f' % (avg_eval_acc)) #log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) test_acc_tb,test_loss_tb= sess.run([ops['test_acc_tbs'],ops['test_loss_tbs']], feed_dict={ops['test_acc']:eval_acc,ops['test_loss']:mean_loss}) test_writer.add_summary(test_acc_tb, epoch) test_writer.add_summary(test_loss_tb, epoch) return eval_acc if __name__ == "__main__": train() # add by nishi acc_saver.save() next_epoch.save() if LOG_FOUT != None: LOG_FOUT.close()