使用TensorFlow在Python中包装C++结构

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 &params, 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 &params, T *out);
    };
}

#endif
Run Code Online (Sandbox Code Playgroud)

我的实现是在cpp文件中,并省略.您会注意到我声明了一个ExpQuadParams要传递给我的BuildCovarianceFunctor仿函数的结构.我想要做的是用TensorFlow将这个结构体暴露给Python,这样我就可以在Python环境中填充数据成员并将它们传递给C++实现.

这是可能的,如果是这样,TensorFlow会如何做到这一点?Tensorflow是否具有内置的这种包装功能?