Tensorflow 2 中的控制流 - 梯度为无答案

【问题标题】：Control flow in Tensorflow 2 - gradients are NoneTensorflow 2 中的控制流 - 梯度为无
【发布时间】：2021-07-01 10:27:04
【问题描述】：

我有一个 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)

问题 - 运行此代码时，gradients 返回[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 
根据上例中 x 的值，磁带要么记录结果 = v0 或结果 = v1**2。 相对于的梯度 x 总是无。
dx = tape.gradient(result, x)
print(dx)
>> None

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

【问题讨论】：

标签： python tensorflow keras tensorflow2.0 control-flow

【解决方案1】：

您正确识别了问题。条件的控制语句不可微，因此您失去了与产生 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)

【讨论】：