Tensorflow 2 中的控制流 - 梯度为 None

Question

我有一个 Tensorflow 2.x 模型，目的是动态选择计算路径。这是该模型的示意图：

唯一可训练的模块是决策模块 (DM)，它本质上是一个完全连接的层，具有单个二进制输出（0 或 1；它可以使用称为改进语义哈希的技术进行区分）。网络 A 和 B 具有相同的网络架构。在训练过程中，我前馈一批图像直到 DM 的输出，然后逐个图像处理决策，将每个图像引导到决策网络（A 或 B）。这些预测被连接成一个张量，用于评估性能。这是训练代码（sigma 是 DM 的输出；model 包括特征提取器和 DM）：

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')


@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        # training=True is only needed if there are custom_layers with different
        # behavior during training versus inference (e.g. Dropout).
        _, sigma = model(images, training=True)
        out = []
        for img, s in zip(images, sigma):
            if s == 0:
                o = binary_classifier_model_a(tf.expand_dims(img, axis=0), training=False)
            else:
                o = binary_classifier_model_b(tf.expand_dims(img, axis=0), training=False)
            out.append(o)

        predictions = tf.concat(out, axis=0)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)

问题-当运行这段代码时，gradientsreturns[None, None]。我现在知道的是：

模型的第一部分（直到 DM 的输出）是可微的；我仅通过运行这一部分对其进行了测试，并应用了损失函数 (MSE)，然后应用了 tape.gradients - 我得到了实际的梯度。
我尝试选择一个（常数）网络 - 例如，网络 A - 并简单地将其输出乘以 s（0 或 1）；这是代替代码中的 if-else 块执行的。在这种情况下，我也得到了渐变。

我担心这样的事情可能无法实现 - 引用 official docs:

x = tf.constant(1.0)

v0 = tf.Variable(2.0)
v1 = tf.Variable(2.0)

with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  if x > 0.0:
    result = v0
  else:
    result = v1**2 
Depending on the value of x in the above example, the tape either records result = v0 or result = v1**2. The gradient with respect to x is always None.
dx = tape.gradient(result, x)
print(dx)
>> None

我不是100%确定这是我的情况，但我想在这里征求专家的意见。我正在尝试做的事情可能吗？如果是 - 我应该改变什么才能让它起作用？谢谢

Answer 1

您正确地识别了问题。条件的控制语句是不可微分的，所以你失去了你的 link 到产生 sigma.

的模型变量

在你的例子中，因为你声明 sigma 是 1 或 0，你可以使用 sigma 的值作为掩码，并跳过条件语句（甚至循环）。

with tf.GradientTape() as tape:
    _, sigma = model(images, training=True)
    predictions = (1.0 - sigma) * binary_classifier_model_a(images, training=False)\
                   + sigma * binary_classifier_model_b(images, training=False)
    loss = loss_object(labels, predictions)

Tensorflow 2 中的控制流 - 梯度为 None

Control flow in Tensorflow 2 - gradients are None

python

tensorflow

keras

tensorflow2.0

control-flow