tensorflow自定义GPU版本op节点

由于前段时间导师布置了一个任务，要修改损失函数，但是这个损失函数在tensorflow自带的库中又没有，想了很多办法，试来试去找不到一个解决方案，因为tensorflow是把框架和数据分开的，所以直接用python写出来的函数是不能用的，只能定义一个节点来调用才行，所以就自然想到先跑一个gpu版本的kernel例程啦，网上cpu版本的教程很多，但是gpu版本的却比较的少，官网的教程极课学院有讲，但我觉得讲的太复杂，反正我是看了一遍没看懂，好了，开始正文。本次例程实现的是将输入tensor中的数字加一输出。

步骤1:写一个kernel

文件名：cuda_op_kernel.cu.cc 代码如下：

#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"__global__ void AddOneKernel(const int* in, const int N, int* out) {for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N;i += blockDim.x * gridDim.x) {out[i] = in[i] + 1;}
}void AddOneKernelLauncher(const int* in, const int N, int* out) {AddOneKernel<<<32, 256>>>(in, N, out);
}#endif

步骤2:编写cpp程序

文件名：cuda_op_kernel.cc 代码如下：

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"using namespace tensorflow;REGISTER_OP("AddOne").Input("input: int32").Output("output: int32").Doc(R"doc(
Adds 1 to all elements of the tensor.
output: A Tensor.output = input + 1
)doc");void AddOneKernelLauncher(const int* in, const int N, int* out);class AddOneOp : public OpKernel {public:explicit AddOneOp(OpKernelConstruction* context) : OpKernel(context) {}void Compute(OpKernelContext* context) override {// Grab the input tensorconst Tensor& input_tensor = context->input(0);auto input = input_tensor.flat<int32>();// Create an output tensorTensor* output_tensor = NULL;OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),&output_tensor));auto output = output_tensor->template flat<int32>();// Set all but the first element of the output tensor to 0.const int N = input.size();// Call the cuda kernel launcherAddOneKernelLauncher(input.data(), N, output.data());}
};REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU), AddOneOp);

步骤3:编译该节点

在上面两个文件夹下打开终端（ubuntu），输入以下命令：

TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')

g++ -std=c++11 -shared cuda_op_kernel.cc -o cuda_op_kernel.so -fPIC -I $TF_INC -O2 -D_GLIBCXX_USE_CXX11_ABI=0

nvcc -std=c++11 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \

-I $TF_INC -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \

cuda_op_kernel.cu.o -I $TF_INC -fPIC -lcudart

当你执行到第四个命令的时候很可能会碰到这样一个报错：

/usr/bin/ld: 找不到 -lcudart
collect2: error: ld returned 1 exit status

对于这个问题我搞了一天，搞清楚这句话的意思大概是有一个叫libcudart.so的库文件找不到，然后我去/usr/bin/目录下确实没找到这个叫libcudart.so的库文件，那咋办呢，那我就在电脑里找这个库文件呀，我电脑里的库文件就在这个文件夹下：/usr/local/cuda/targets/x86_64-linux/lib，大家可以在cuda这个文件下搜一下，我估计就在这个文件夹底下，知道地址以后呢，我们把原来的第四个命令替换成下面这个命令：

g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc cuda_op_kernel.cu.o -L/usr/local/cuda/targets/x86_64-linux/lib -I $TF_INC -fPIC -lcudart

就编译通过啦。

然后我写了一个测试文件

步骤4:测试

文件名：test.py 代码如下：

import tensorflow as tf
cuda_op_module = tf.load_op_library('./cuda_op_kernel.so')
with tf.Session(''):x=cuda_op_module.add_one([[1, 2], [3, 4]]).eval()print x

输出是：
[[2 3]
[4 5]]

tensorflow自定义GPU版本op节点相关推荐

Ubuntu tensorflow自定义GPU版本op节点
参考:https://blog.csdn.net/qq_27637315/article/details/79114633 windows增加op节点: https://github.com/tens ...
安装tensorflow的GPU版本（详细图文教程）--CUDA11.6的安装
文章目录 TensorFlow简介 TensorFlow是什么 tensorflow版本变迁 tensorflow 2.0 架构安装过程常用IDE安装 python3.9的安装 Anaconda的 ...
当涉及某个项目需要大量使用到tensorflow时，最后不要使用tensorflow的GPU版本，这会出很多毛病，最好使用CPU版本
当涉及某个项目需要大量使用到tensorflow时,最后不要使用tensorflow的GPU版本,这会出很多毛病,最好使用CPU版本
常识——windows的tensorflow安装gpu版本,cuda算力只有3.0也可以
文章目录问题 windows安装tensorflow_gpu版本检查显卡是否支持cuda 查看自己的显卡型号去官网查看是否支持cuda 检查显卡的cuda算力按照GPU算力3.0的电脑配置进行 ...
anaconda3 tensorflow安装踩坑记（WIN10+tensorflow带gpu版本）
我看的网课讲解的视频用的是带gpu版本的tensorflow1.4.0,使用pip install tensorflow-gpu==1.4.0,但是如果想直接用之前那个命令的话要使用anaconda老 ...
Windows7上安装TensorFlow的GPU版本后记
上一篇写了之前在我的Windows7上安装TensorFlow GPU版的一路坑坑洼洼,本以为已经解决了.但实际上在过年的几天内,又一次出现了问题.经过了反复的尝试,到今天为止终于稳定下来,特此记录一 ...
tensorflow 安装GPU版本，CUDA与cuDNN版本对应关系，RTX3050Ti （notebook）
前言安装Tensorflow-gpu 与 keras的时候,一定先要注意版本的对应,不然很容易出错,在看的时候,建议先看完整篇文章再上手. 一.环境+配置本机环境显卡:RTX3050Ti(not ...
win7, gtx750ti 2g tensorflow 安装GPU版本，个人总结，步骤比较详细
安装环境:win7, gtx750ti 2g 参考: https://blog.csdn.net/gangeqian2/article/details/79358543 http://blog.sin ...
RTX 3060 快速配置GPU版本tensorflow
目录一.Anaconda安装 1. Anaconda下载方式一:官网下载方式二:清华镜像下载 2. Anaconda安装 3. Anaconda配置 3.1 配置环境变量 3.2 配置国内镜像源 ...

tensorflow自定义GPU版本op节点

tensorflow自定义GPU版本op节点相关推荐

最新文章

热门文章