构建一个如下所示的天真内核。
#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;用占位符替换,然后只执行一次。
这是故意的吗?为什么有必要?