Merge pull request NVlabs#6 from arsalan-mousavian/arsalan/easy_tf_setup

Easy script for building tf ops

README.md

@@ -12,7 +12,7 @@ ICRA 2021
 
 ## Installation
 
-This code has been tested with python 3.7, tensorflow 2.2, CUDA 10.1, and CUDNN 7.6.0
+This code has been tested with python 3.7, tensorflow 2.2, CUDA 11.1
 
 Create the conda env
 ```
@@ -21,7 +21,10 @@ conda env create -f contact_graspnet_env.yml
 
 ### Troubleshooting
 
-- Recompile pointnet2 tf_ops, see [here](pointnet2/tf_ops/HowTO.md) 
+- Recompile pointnet2 tf_ops:
+```shell
+sh compile_pointnet_tfops.sh
+```
 
 ### Hardware
 Training: 1x Nvidia GPU >= 24GB VRAM, >=64GB RAM  

compile_pointnet_tfops.sh

@@ -0,0 +1,32 @@
+CUDA_INCLUDE=' -I/usr/local/cuda/include/'
+CUDA_LIB=' -L/usr/local/cuda/lib64/'
+TF_CFLAGS=$(python3 -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))')
+TF_LFLAGS=$(python3 -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))')
+cd pointnet2/tf_ops/sampling
+
+nvcc -std=c++11 -c -o tf_sampling_g.cu.o tf_sampling_g.cu \
+ ${CUDA_INCLUDE} ${TF_CFLAGS} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
+
+g++ -std=c++11 -shared -o tf_sampling_so.so tf_sampling.cpp \
+ tf_sampling_g.cu.o ${CUDA_INCLUDE} ${TF_CFLAGS} -fPIC -lcudart ${TF_LFLAGS} ${CUDA_LIB}
+
+echo 'testing sampling'
+python3 tf_sampling.py
+
+cd ../grouping
+
+nvcc -std=c++11 -c -o tf_grouping_g.cu.o tf_grouping_g.cu \
+ ${CUDA_INCLUDE} ${TF_CFLAGS} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
+
+g++ -std=c++11 -shared -o tf_grouping_so.so tf_grouping.cpp \
+ tf_grouping_g.cu.o ${CUDA_INCLUDE} ${TF_CFLAGS} -fPIC -lcudart ${TF_LFLAGS} ${CUDA_LIB}
+
+echo 'testing grouping'
+python3 tf_grouping_op_test.py
+
+
+cd ../3d_interpolation
+g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -shared -fPIC ${TF_CFLAGS} ${TF_LFLAGS} -O2
+echo 'testing interpolate'
+python3 tf_interpolate_op_test.py
+