def train(args):
num_points = 2048 # ModelNet40
num_channels = 3
max_epoch = args["max_epoch"]
start_learning_rate = args["start_learning_rate"]
decay_steps = args["decay_steps"]
decay_rate = args["decay_rate"]
momentum = args["momentum"]
sort_cloud = args["sort_cloud"]
sort_method = args["sort_method"]
snapshot_interval = args["snapshot_interval"]
batch_size = args["batch_size"]
if len(sys.argv) > 1:
start_epoch = int(sys.argv[1])
else:
start_epoch = 0
model_name = args["model"]
model = importlib.import_module(model_name)
print('Using model ', model_name)
d = modelnet.DataConsumer(file='data/modelnet40_ply_hdf5_2048/train_files.txt', batch_size=batch_size, num_points=num_points, num_channels=num_channels, sort_cloud=sort_cloud, sort_method=sort_method)
num_class = 40
print('Categories: ', num_class)
print('Sort cloud: ', sort_cloud)
use_gpu = args["use_gpu"]
if use_gpu:
str_device = '/gpu:0'
else:
str_device = '/cpu:0'
print('Using device', str_device)
with tf.Graph().as_default(), tf.device(str_device):
if not use_gpu:
config = tf.ConfigProto(
device_count = {'GPU': 0},
log_device_placement = False
)
else:
config = tf.ConfigProto(
allow_soft_placement = True,
log_device_placement = False
)
with tf.Session(config=config) as session:
network = model.PointConvNet(num_class)
batch_points_placeholder = tf.placeholder(tf.float32, shape=(batch_size, num_points, 3))
batch_input_placeholder = tf.placeholder(tf.float32, shape=(batch_size, num_points, num_channels))
batch_label_placeholder = tf.placeholder(tf.int32, shape=(batch_size))
with tf.variable_scope("pointcnn") as scope:
batch_prediction_placeholder = network.model(batch_points_placeholder, batch_input_placeholder, is_training=True)
loss = network.loss(batch_prediction_placeholder, batch_label_placeholder)
# Track the global training progress
global_step = tf.Variable(0, name='global_step', trainable=False)
# Adaptive learning rate
learning_rate = tf.train.exponential_decay(start_learning_rate, global_step,
decay_steps, decay_rate, staircase=True)
optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)
train_op = optimizer.minimize(loss, global_step=global_step)
# Save to disk
saver = tf.train.Saver(max_to_keep=None)
weight_file = "{}_snapshot_{}.tf".format(model_name, start_epoch)
if start_epoch > 0 and os.path.isfile(weight_file+".index"):
saver.restore(session, weight_file)
print('Weight file {} restored'.format(weight_file))
start_epoch += 1
else:
print('Train from scratch');
init_op = tf.global_variables_initializer()
session.run(init_op)
start_epoch = 0
# Live statistics
tf.summary.scalar('learning_rate', learning_rate)
tf.summary.scalar('loss', loss)
# Write some statistics to file during training
merged = tf.summary.merge_all()
train_writer = tf.summary.FileWriter('log', session.graph)
for epoch in range(start_epoch, max_epoch):
print('# Epoch %d' % epoch)
# one epoch
total_correct = 0
total_seen = 0
total_correct_class = np.zeros((num_class))
total_seen_class = np.zeros((num_class))
total_loss = 0
num_batches = 0
batch_time = 0
while True:
points, points_data, gt_label = d.get_batch_point_cloud()
tic = timeit.default_timer()
feed_dict = {
batch_points_placeholder : points,
batch_input_placeholder : points_data,
batch_label_placeholder : gt_label
}
_, summary, loss_val, pred_val = session.run([train_op, merged, loss, batch_prediction_placeholder], feed_dict=feed_dict)
toc = timeit.default_timer()
batch_time += toc - tic
pred_label = np.argmax(pred_val, axis=1)
correct = np.sum(pred_label == gt_label)
total_correct += correct
total_seen += batch_size
total_loss += loss_val
for i in range(batch_size):
l = gt_label[i]
total_seen_class[l] += 1
total_correct_class[l] += (pred_label[i] == l)
num_batches += 1
print('Current loss : %f' % loss_val)
print('Current accuracy : %f' % (correct / float(batch_size)))
print('Average batch time : ', batch_time / num_batches)
if not d.has_next_batch():
break
d.next_batch()
mean_loss = (total_loss / float(num_batches))
mean_accuracy = (total_correct / float(total_seen))
mean_class_accuracy = (np.mean(total_correct_class / total_seen_class))
print('Mean loss : %f' % mean_loss)
print('Mean accuracy : %f' % mean_accuracy)
print('Avg class accuracy : %f' % mean_class_accuracy)
train_writer.add_summary(summary)
if epoch % snapshot_interval == 0 or epoch == max_epoch - 1:
snapshot_file = "./{}_snapshot_{}.tf".format(model_name, epoch)
saver.save(session, snapshot_file)
print('Model %s saved' % snapshot_file)
global_step_val, learning_rate_val = session.run([global_step, learning_rate])
print('Global step : ', global_step_val)
print('Learning rate : ', learning_rate_val)
with open("train_info_{}.txt".format(model_name), 'a') as f:
f.write(str(datetime.datetime.now().strftime("%c")) + ', ' + \
str(epoch) + ', ' + str(mean_loss) + ', ' + str(mean_accuracy) + ', ' + str(mean_class_accuracy) + '\n')
d.next_epoch()
def eval(args):
num_points = 2048 # ModelNet40
num_channels = 3
model_name = args["model"]
model = importlib.import_module(model_name)
print('Using model ', model_name)
sort_cloud = args['sort_cloud']
sort_method = args['sort_method']
if len(sys.argv) > 1:
epoch = int(sys.argv[1])
else:
epoch = 0
batch_size = args["batch_size"]
d = modelnet.DataConsumer(
file='data/modelnet40_ply_hdf5_2048/test_files.txt',
batch_size=batch_size,
num_points=num_points,
num_channels=num_channels,
test=True,
sort_cloud=sort_cloud,
sort_method=sort_method)
num_class = 40
print('Categories: ', num_class)
print('Sort cloud: ', sort_cloud)
use_gpu = args["use_gpu"]
if use_gpu:
str_device = '/gpu:0'
else:
str_device = '/cpu:0'
print('Using device', str_device)
with tf.Graph().as_default(), tf.device(str_device):
if not use_gpu:
config = tf.ConfigProto(device_count={'GPU': 0},
log_device_placement=False)
else:
config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)
with tf.Session(config=config) as session:
network = model.PointConvNet(num_class)
batch_points_placeholder = tf.placeholder(tf.float32,
shape=(batch_size,
num_points, 3))
batch_input_placeholder = tf.placeholder(tf.float32,
shape=(batch_size,
num_points,
num_channels))
batch_label_placeholder = tf.placeholder(tf.int32,
shape=(batch_size))
with tf.variable_scope("pointcnn") as scope:
batch_prediction_placeholder = network.model(
batch_points_placeholder,
batch_input_placeholder,
is_training=False)
saver = tf.train.Saver(max_to_keep=None)
file = "{}_snapshot_{}.tf".format(model_name, epoch)
print('Testing with model ', file)
saver.restore(session, file)
total_correct = 0
total_seen = 0
total_correct_class = np.zeros((num_class))
total_seen_class = np.zeros((num_class))
num_batches = 0
while True:
points, points_data, gt_label = d.get_batch_point_cloud()
feed_dict = {
batch_points_placeholder: points,
batch_input_placeholder: points_data,
batch_label_placeholder: gt_label
}
[pred_val] = session.run([batch_prediction_placeholder],
feed_dict=feed_dict)
pred_label = np.argmax(pred_val, axis=1)
correct = np.sum(pred_label == gt_label)
total_correct += correct
total_seen += batch_size
for i in range(batch_size):
l = gt_label[i]
total_seen_class[l] += 1
total_correct_class[l] += (pred_label[i] == l)
num_batches += 1
print('Current accuracy : %f' % (correct / float(batch_size)))
if not d.has_next_batch():
break
d.next_batch()
mean_accuracy = (total_correct / float(total_seen))
mean_class_accuracy = (np.mean(total_correct_class /
total_seen_class))
print('Mean accuracy : %f' % mean_accuracy)
print('Avg class accuracy : %f' % mean_class_accuracy)
with open("test_info_{}.csv".format(model_name), 'a') as f:
f.write(str(datetime.datetime.now().strftime("%c")) + ', ' + \
str(epoch) + ', ' + str(mean_accuracy) + ', ' + str(mean_class_accuracy) + '\n')