由于前段时间导师布置了一个任务,要修改损失函数,但是这个损失函数在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 tensor
const Tensor& input_tensor = context->input(0);
auto input = input_tensor.flat<int32>();
// Create an output tensor
Tensor* output_tensor = NULL;
OP_