tf.cond 未按预期运行

Question

我有一个关于 tensorflow 如何计算 tf.cond 语句中的表达式的问题。

我使用自定义操作为我提供 A 类或 B 类数据的批次。它的签名看起来像这样

REGISTER_OP("GetData")
.Attr("path: string")
.Output("data: int32")
.Output("type: int32")
.SetIsStateful()

type 类型 A 的输出为 1，类型 B 的输出为 2。根据类型，我想运行不同的（自定义）操作，例如 opA 或opB。当然，它们的输出具有相同的类型。为了表达这个数据流，我使用 tf.cond 作为：

(data, type) = get_data(path = "...")

opA = op_a(data)
opB = op_b(dat

def perfromA():  return opA
def perfromB():  return opB

joined_op = tf.cond(tf.equal(type, tf.constant(1, dtype=tf.int32)), perfromA, perfromB)

我在 getData、opA 和 opB 中添加了一些调试语句。对于一个数据序列ABA希望得到

getData: returning type A
opA:     got some data
getData: returning type B
opB:     got some data   
getData: returning type A
opA:     got some data

不过，我确实得到了

getData: returning type A
opA:     got some data
opB:     got some data
getData: returning type B
opA:     got some data
opB:     got some data   
getData: returning type A
opA:     got some data
opB:     got some data

这是预期的行为吗？ joined_op 的结果确实正确地考虑了 if-else，但仍然会计算这两个操作。这不仅是一种计算负担，而且如果 opA 和 opA 执行影响变量的操作（例如优化步骤），也会破坏目的。

正确的解决方案 @Yaroslav

指出

def perfromA():  return op_a(data)
def perfromB():  return op_b(data)

有关更详细的示例，请参阅以下三个操作的虚拟实现：

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/simple_philox.h"

using namespace std;
using namespace tensorflow;


REGISTER_OP("GetData")
.Attr("path: string")
.Output("data: int32")
.Output("type: int32")
.SetIsStateful();


class GetData : public OpKernel {
    public:
        explicit GetData(OpKernelConstruction* ctx) : OpKernel(ctx) {}

        void Compute(OpKernelContext* ctx) override
        {
            Tensor data(DT_INT32, TensorShape({}));
            Tensor type(DT_INT32, TensorShape({}));
            if(rng_.Uniform(2) == 0){
                type.scalar<int32>()() = 1;
                data.scalar<int32>()() = 100;
                LOG(INFO) << "returning type A";
            }
            else{
                type.scalar<int32>()() = 2;
                data.scalar<int32>()() = 200;
                LOG(INFO) << "returning type B";
            }

            ctx->set_output(0, data);
            ctx->set_output(1, type);
        }

    private:
        random::PhiloxRandom philox_ =  random::PhiloxRandom(10) ;
        random::SimplePhilox rng_ = random::SimplePhilox(&philox_);

};
REGISTER_KERNEL_BUILDER(Name("GetData").Device(DEVICE_CPU), GetData);


REGISTER_OP("OpA")
.Input("input: int32")
.Output("output: int32");

class OpA : public OpKernel {
    public:
        explicit OpA(OpKernelConstruction* ctx) : OpKernel(ctx) {}

        void Compute(OpKernelContext* ctx) override
        {
              const Tensor& data = ctx->input(0);
              LOG(INFO) << "A: got some data";
              ctx->set_output(0, data);
        }
};
REGISTER_KERNEL_BUILDER(Name("OpA").Device(DEVICE_CPU), OpA);

REGISTER_OP("OpB")
.Input("input: int32")
.Output("output: int32");

class OpB : public OpKernel {
    public:
        explicit OpB(OpKernelConstruction* ctx) : OpKernel(ctx) {}

        void Compute(OpKernelContext* ctx) override
        {
              const Tensor& data = ctx->input(0);
              LOG(INFO) << "B: got some data";
              ctx->set_output(0, data);
        }
};
REGISTER_KERNEL_BUILDER(Name("OpB").Device(DEVICE_CPU), OpB);

Answer 1

您必须在 tf.cond

中创建 opA/opB

查看讨论

tf.cond 未按预期运行

tf.cond not behaving as expected

c++

tensorflow