Tensorflow，Keras：如何使用停止渐变在Keras图层中设置add_loss？

Question

问题1

我们知道我们可以使用tf.stop_gradient(B)来防止变量B在反向传播中受到训练。 但我不知道如何在某种损失中阻止B

简单来说，假设我们的损失是：

loss = categorical_crossentropy + my_loss
B = tf.stop_gradient(B)

其中categorical_crossentropy和my_loss都依赖于B 因此，如果我们为B设置停止梯度，它们都将B作为常量。

但是我如何只设置my_loss停止渐变wrt B ，保持categorical_crossentropy不变？ 比如B = tf.stop_gradient(B, myloss)

我的代码是：

my_loss = ...
B = tf.stop_gradient(B)
categorical_crossentropy = ...
loss = categorical_crossentropy + my_loss

那会有用吗？ 或者，如何使这项工作？

---

问题2

好的，伙计们，如果Q1能够解决，我最后的任务是如何在自定义层中做到这一点？

具体来说，假设我们有一个自定义图层，它具有可训练的权重A和B并且my_loss此图层自我损失my_loss 。

class My_Layer(keras.layers.Layer):
    def __init__(self, **kwargs):
        super(My_Layer, self).__init__(**kwargs)
    def build(self, input_shape):
        self.w = self.add_weight(name='w', trainable=True)
        self.B = self.add_weight(name='B', trainable=True)
        my_loss = w * B
        # tf.stop_gradient(w)
        self.add_loss(my_loss)

我如何使w只能训练模型损失（MSE，交叉熵等），而B只能训练my_loss ？

如果我添加tf.stop_gradient(w) ， my_loss仅针对my_loss停止w还是模型的最终丢失？

Answer 1

问题1

运行y = tf.stop_gradient(x) ，将创建一个StopGradient操作，其输入为x ，输出为y 。 此操作的行为类似于标识，即x的值与y的值相同，除了渐变不从y流向x 。

如果你想让梯度仅从一些损失流向B ，你可以简单地做：

B_no_grad = tf.stop_gradient(B)
loss1 = get_loss(B)  # B will be updated because of loss1
loss2 = get_loss(B_no_grad)   # B will not be updated because of loss2 

当您考虑要构建的计算图时，事情应该变得清晰。 stop_gradient允许您为任何不允许渐变流过它的张量（不仅仅是变量）创建“标识”节点。

问题2

我不知道如何使用你使用字符串指定的模型丢失（例如model.compile(loss='categorical_crossentropy', ...)因为你不控制它的结构。但是，你可以做它通过增加损失，使用add_loss或使用模型输出建立模型级的损失你自己。对于前者，只需创建一个使用普通的变量的一些使用一些损失和*_no_grad版本，添加它们全部采用add_loss()并编译模型loss=None 。