一个额外的内核计算方法调用

时间:2017-07-21 07:42:13

标签: tensorflow

构建一个如下所示的天真内核。

#include "tensorflow/core/framework/op.h"                                                                                                                                                      
#include "tensorflow/core/framework/op_kernel.h"                                                                                                                                               
#include "tensorflow/core/framework/tensor.h"                                                                                                                                                  
#include "tensorflow/core/lib/core/errors.h"                                                                                                                                                   
#include "tensorflow/core/lib/core/status.h"                                                                                                                                                   
#include "tensorflow/core/lib/strings/str_util.h"                                                                                                                                              
#include "tensorflow/core/lib/strings/strcat.h"                                                                                                                                                
#include "tensorflow/core/framework/shape_inference.h"                                                                                                                                         
#include "tensorflow/core/platform/logging.h"                                                                                                                                                  


namespace tensorflow {                                                                                                                                                                         

using namespace std;                                                                                                                                                                           

REGISTER_OP("Test")                                                                                                                                                                            
.Input("in: int32")                                                                                                                                                                            
.Output("out: int32")                                                                                                                                                                          
// .SetIsStateful()                                                                                                                                                                            
.Doc(R"doc(                                                                                                                                                                                    
wierd                                                                                                                                                                                          
)doc");                                                                                                                                                                                        

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

  void Compute(OpKernelContext* context) override {                                                                                                                                            
    LOG(INFO) << "Compute";                                                                                                                                                                    
  }                                                                                                                                                                                            
};                                                                                                                                                                                             

REGISTER_KERNEL_BUILDER(Name("Test").Device(DEVICE_CPU), TestOp);                                                                                                                              

}  // namespace tensorflow

使用以下命令在python中运行:

import tensorflow as tf
t = tf.load_op_library("./test_op.so")
test = t.test(8)
sess = tf.Session()
sess.run(test)

这打印&#34;计算&#34;两次。

2017-07-21 15:34:23.049053: I test_op.cc:45] Compute
2017-07-21 15:34:23.049253: I test_op.cc:45] Compute

如果启用SetIsStateful,或者如果是&#34; 8&#34;用占位符替换,然后只执行一次。 这是故意的吗?为什么有必要?

0 个答案:

没有答案