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

Question

问题1

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

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

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

其中{strong> categorical_crossentropy和my_loss都取决于B 。因此，如果我们为B设置停止梯度，则它们都将B设为常数。

但是如何仅设置my_loss的{​​{1}}停止梯度，而保持B不变？像categorical_crossentropy

我的代码是：

B = tf.stop_gradient(B, myloss)

那行得通吗？或者，如何使它起作用？

---

问题2

好的，伙计们，如果Q1可以解决，那么我的最终任务是如何在自定义层中做到这一点？

具体来说，假设我们有一个自定义层，该层仅具有可训练的权重my_loss = ... B = tf.stop_gradient(B) categorical_crossentropy = ... loss = categorical_crossentropy + my_loss 和A，而自损耗B仅适用于该层。

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) 仅对模型丢失（MSE，交叉熵等）可训练，而w仅对B可训练？

如果我添加了my_loss，它会停止tf.stop_gradient(w)的{​​{1}}还是模型的最终损失？

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允许您为任何不允许梯度流过的张量（不仅仅是变量）创建一个“ identity”节点。

问题2

我不知道如何使用通过字符串指定的模型损失（例如model.compile(loss='categorical_crossentropy', ...)，因为您无法控制其构造。但是，您可以使用{ {1}}或使用模型输出自己构建模型级损失，对于前者，只需使用简单变量创建损失，然后使用add_loss版本创建损失，然后使用*_no_grad添加所有损失，然后编译add_loss()进行建模。