Jac*_*k H 7 c++ python deep-learning tensorflow
我目前正在用C++编写一个自定义的TensorFlow OP,它将被包装在python中.我定义我的OP和其CPU内核仿函数如下,按照该指南.
#ifndef BUILD_COVARIANCE_FUNCTOR_HEADER
#define BUILD_COVARIANCE_FUNCTOR_HEADER
#define ARD_PARAM_BUFFER_SIZE 256
#define EIGEN_USE_THREADS
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow{
REGISTER_OP("CovMatrixGenerator")
.Input("input_features: float32")
.Output("output_latent: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0));
return Status::OK();
});
using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;
template <typename D, typename T>
class CovMatrixGenerator : public OpKernel {
public:
explicit CovMatrixGenerator(OpKernelConstruction* context);
void Compute(OpKernelContext* context) override;
};
// Register the CPU kernels.
#define REGISTER_CPU(T) \
REGISTER_KERNEL_BUILDER( \
Name("CovMatrixGenerator").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
CovMatrixGenerator<CPUDevice, T>);
REGISTER_CPU(float);
REGISTER_CPU(int32);
template<typename T, size_t dim>
struct ExpQuadParams {
T ardTheta[dim];
T theta;
const size_t dimensionality = dim;
EIGEN_DEVICE_FUNC
ExpQuadParams(T thetaInit, T ardThetaInit) {
for (size_t i = 0; i < dim; i++) {
ardTheta[i] = ardThetaInit;
}
theta = thetaInit;
}
ExpQuadParams() : ExpQuadParams(static_cast<T>(1), static_cast<T>(1)) {
//
}
};
template<typename D, typename T>
struct BuildCovarianceFunctor {
void operator()(const D &d, const T *data1, const T *data2, size_t size1,
size_t size_2, size_t dim, ExpQuadParams ¶ms, T *out);
};
template<typename T>
struct BuildCovarianceFunctor<CPUDevice, T> {
void operator()(const CPUDevice &deviceType, const T *data1, const T *data2, size_t size1,
size_t size_2, size_t dim, ExpQuadParams ¶ms, T *out);
};
}
#endif
Run Code Online (Sandbox Code Playgroud)
我的实现是在cpp文件中,并省略.您会注意到我声明了一个ExpQuadParams要传递给我的BuildCovarianceFunctor仿函数的结构.我想要做的是用TensorFlow将这个结构体暴露给Python,这样我就可以在Python环境中填充数据成员并将它们传递给C++实现.
这是可能的,如果是这样,TensorFlow会如何做到这一点?Tensorflow是否具有内置的这种包装功能?
| 归档时间: |
|
| 查看次数: |
128 次 |
| 最近记录: |