添加自定义稀疏运算(稀疏行列式)

时间:2018-09-23 05:40:05

标签: c++ tensorflow

我正在努力使一些稀疏矩阵运算在Tensorflow中工作。我要解决的第一个问题是通过稀疏的Cholesky分解的稀疏行列式。 Eigen的Cholesky稀疏,所以我想把它包起来。

我一直在进步,但是现在有点卡住了。我知道Tensorflow中的SparseTensors由三部分组成:indicesvaluesshape。复制类似的操作,我去了下面的REGISTER_OP声明:

REGISTER_OP("SparseLogDet")
    .Input("a_indices: int64")
    .Input("a_values: float32")
    .Input("a_shape: int64")
    .Output("determinant: float32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle h;
      c->set_output(0, h);
      return Status::OK();
    });

这可以很好地编译,但是当我使用一些示例代码运行它时:

import tensorflow as tf

log_det_op = tf.load_op_library('./sparse_log_det_op.so')

with tf.Session(''):
  t = tf.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2],
                      dense_shape=[3, 4])
  print(log_det_op.sparse_log_det(t).eval().shape)
  print(log_det_op.sparse_log_det(t).eval())

它抱怨说:

TypeError: sparse_log_det() missing 2 required positional arguments: 'a_values' and 'a_shape'

这对我来说很有意义,因为它期待着其他论点。但是,我真的很想传递稀疏张量,而不是将其分解为分量!有谁知道如何处理其他稀疏操作?

谢谢!

1 个答案:

答案 0 :(得分:1)

如果要传递稀疏张量,然后从中确定indicesvaluesshape,则应该可行。只需修改您的OP即可接受单个Tensor输入,并产生单个float输出。然后通过遍历其元素,从Eigen :: Tensor中提取所需的信息,如下所示:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
#include <Eigen/Dense>

using namespace tensorflow;

REGISTER_OP("SparseDeterminant")
    .Input("sparse_tensor: float")
    .Output("sparse_determinant: float");


class SparseDeterminantOp : public OpKernel {
public:
    explicit SparseDeterminantOp(OpKernelConstruction *context) : OpKernel(context) {}

    void Compute(OpKernelContext *context) override {

   // get the input tesnorflow tensor
   const Tensor& sparse_tensor = context->input(0);  
   // get shape of input
   const TensorShape& sparse_shape = sparse_tensor.shape();

   // get Eigen Tensor for input tensor
   auto eigen_sparse = sparse_tensor.matrix<float>();

   //extract the data you want from the sparse tensor input
   auto a_shape = sparse_tensor.shape();

   // loop over all elements of the input tensor and add to values and indices
   for (int i=0; i<a_shape.dim_size(0); ++i){
    for (int j=0; j<a_shape.dim_size(1); ++j){
        if(eigen_sparse(i,j) != 0){
        /// ***Here add non zero elements to list/tensor of values and their indicies*** 
            std::cout<<eigen_sparse(i,j)<<" at"<<" "<<i<<" "<<j<<" "<<"not zero."<<std::endl;
        }
    }
   }

   // create output tensor
   Tensor *output_tensor = NULL;      
   TensorShape output_shape;
   OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output_tensor));
   auto output = output_tensor->scalar<float>();

   output(0) = 1.; //**asign return value***;

    }
};
REGISTER_KERNEL_BUILDER(Name("SparseDeterminant").Device(DEVICE_CPU), SparseDeterminantOp);

不幸的是,当您将t传递到op中时,它变成了Tensorflow::Tensor,并且丢失了与values相关联的indicestf.sparsetensor方法,因此您可以不要轻易得到它们。

一旦编译后,此代码即可运行:

//run.py
import tensorflow as tf
import numpy as np

my_module = tf.load_op_library('./out.so')

# create sparse matrix
a = np.zeros((10,10))
for i in range(len(a)):
    a[i,i] = i

print(a)

a_t = tf.convert_to_tensor(a, dtype= float)

with tf.Session() as sess:
    sess.run(my_module.sparse_determinant(a_t))