用C++代码编写cuda函数,然后用python调用。相比直接用python函数,用tf做矩阵操作会快很多,特别是一些特殊的没有支持的操作。本文主要简单介绍在总体思路,为以后可能的应用做准备,并没有涉及具体的C++函数编写等。
一个比较好的例子是FlowNet和UnFlow两个tensorflow代码。
在Python中运行shell代码
编译的过程需要运行nvcc
,可以写一个单独的shell脚本来运行。而UnFlow中作者直接在Python代码中运行了shell代码,其中使用到了subprocess
库。(文件位置:UnFlow/src/e2eflow/ops.py
)
nvcc_cmd = "nvcc -std=c++11 -c -o {} {} {} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC -I " + cuda_dir + " --expt-relaxed-constexpr"
nvcc_cmd = nvcc_cmd.format(" ".join([fn_cu_o, fn_cu_cc]),
tf_inc, tf_lib)
subprocess.check_output(nvcc_cmd,<