从unflow学编译tensorflow改善bottleneck

本文介绍了如何利用C++编写CUDA函数并用Python调用来提升TensorFlow的矩阵运算速度,特别是在处理未被直接支持的操作时。通过UnFlow项目作为示例,展示了在Python中运行shell脚本编译CUDA代码,以及如何安全地加载和注册到TensorFlow框架的过程。
摘要由CSDN通过智能技术生成

用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,<
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值