我在玩TensorFlow中的custom op tutorial。在
< >我在C++中实现了自定义OP并创建了一个共享库。我正在用python加载它tf.load_op_库函数调用。然后我用会话.运行(). 在自定义操作运行良好。但我不知道什么时候调用了自定义的析构函数。在
即使我在自定义操作的析构函数中有一个print语句,它也永远不会被打印出来。自定义操作实例似乎永远不会被破坏。在
这是预期的行为吗?如果是这样的话,有没有办法通知tensorflow我已经使用定制操作完成了?在
请注意,只有当我们使用tf.占位符. 如果输入矩阵为常数,即
x=zero_out_module.zero_out([[1, 2], [3, 4]])
然后调用自定义操作的析构函数。在
<>自定义C++在C++中的实现#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
using namespace tensorflow;
REGISTER_OP("ZeroOut")
.Input("to_zero: int32")
.Output("zeroed: int32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0));
return Status::OK();
});
class ZeroOutOp : public OpKernel {
public:
explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
printf("Constructor: ZeroOutOp\n"); fflush(stdout);
}
~ZeroOutOp() override {
printf("Destructor: ZeroOutOp\n"); fflush(stdout);
}
void Compute(OpKernelContext* context) override {
// Grab the input tensor
const Tensor& input_tensor = context->input(0);
auto input = input_tensor.flat<int32>();
// Create an output tensor
Tensor* output_tensor = NULL;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor));
auto output_flat = output_tensor->flat<int32>();
// Set all but the first element of the output tensor to 0.
const int N = input.size();
for (int i = 1; i < N; i++) {
output_flat(i) = 0;
}
// Preserve the first input value if possible.
if (N > 0) output_flat(0) = input(0);
}
};
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);
测试自定义操作的Python代码
^{pr2}$
目前没有回答
相关问题 更多 >
编程相关推荐