From 87f702efe6c7e1432ffc87d604086e09e7b8dae8 Mon Sep 17 00:00:00 2001 From: azaganidis Date: Mon, 10 Dec 2018 16:03:31 +0000 Subject: [PATCH] three_nn gpu implementation. Redefinition of three_nn by tensorflow operations instead of c++ on cpu. Significantly faster. --- utils/pointnet_util.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/utils/pointnet_util.py b/utils/pointnet_util.py index 6d326233..50ccb377 100644 --- a/utils/pointnet_util.py +++ b/utils/pointnet_util.py @@ -195,6 +195,17 @@ def pointnet_sa_module_msg(xyz, points, npoint, radius_list, nsample_list, mlp_l new_points_concat = tf.concat(new_points_list, axis=-1) return new_xyz, new_points_concat +def three_nn_gpu(xyz1, xyz2): + s1=tf.shape(xyz1) + s2=tf.shape(xyz2) + xyz1=tf.reshape(xyz1, (s1[0],s1[1],1,-1)) + xyz2=tf.reshape(xyz2, (s2[0],1,s2[1],-1)) + dist=xyz2-xyz1 + dist=dist*dist + dist=tf.reduce_sum(dist,-1) + if dist.shape[-1]==1: + dist=tf.tile(dist,(1,1,3)) + return tf.nn.top_k(-dist, k=3) def pointnet_fp_module(xyz1, xyz2, points1, points2, mlp, is_training, bn_decay, scope, bn=True): ''' PointNet Feature Propogation (FP) Module @@ -208,7 +219,7 @@ def pointnet_fp_module(xyz1, xyz2, points1, points2, mlp, is_training, bn_decay, new_points: (batch_size, ndataset1, mlp[-1]) TF tensor ''' with tf.variable_scope(scope) as sc: - dist, idx = three_nn(xyz1, xyz2) + dist, idx = three_nn_gpu(xyz1, xyz2) dist = tf.maximum(dist, 1e-10) norm = tf.reduce_sum((1.0/dist),axis=2,keep_dims=True) norm = tf.tile(norm,[1,1,3])