Python 如何覆盖Keras、Tensorflow中优化算法的梯度向量计算方法？

所以我试图修改Keras中的两个优化算法，即Adam或SGD。因此，默认情况下，我非常确定参数更新的工作方式是，在批处理中的数据点上平均损失，然后根据此损失值计算梯度向量。另一种考虑方法是，根据批处理中每个数据点的损失值平均梯度。这是我想要改变的计算，而且会很昂贵，所以我尝试在使用GPU的优化框架内进行

因此，对于每个批次，我需要计算批次中每个数据点损失的梯度，然后，我将不取梯度的平均值，而是进行其他平均值或计算。有人知道我如何获得覆盖Adam或SGD功能的权限吗

经过一番精彩的评论后，我发现应该有一种方法可以实现我试图使用

GradientTape

中的

jacobian

方法所做的事情。然而，文档不是很全面，我不知道它是如何融入整体的。在这里，我希望有人能帮助我调整代码，使用

jacobian

而不是

gradient

作为一个hello world的例子，我试图简单地用一些使用雅可比矩阵的代码替换

梯度线

，并产生相同的输出。这将说明如何使用

jacobian

方法以及与

gradient

方法输出的连接

工作代码

class CustomModel(keras.Model):
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars) # <-- line to change
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

类客户模型（keras.Model）：
def系列步骤（自身、数据）：
#打开数据包。其结构取决于您的型号和型号
#传递给“fit（）”的内容。
x、 y=数据
使用tf.GradientTape（）作为磁带：
y_pred=自我（x，训练=正确）#向前传球
#计算损失值
#（loss函数在“compile（）”中配置）
损耗=自编译损耗（y，y，pred，正则化损耗=自编译损耗）
#计算梯度
可训练变量=自可训练变量
梯度=磁带。梯度（损耗、可训练变量）#您应该能够执行以下操作：

class CustomModel(keras.Model):
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.jacobian(loss, trainable_vars)

        new_gradients = []
        for grad in gradients:
            new_grad = do_something_to(grad)
            new_gradients.append(new_grad)

        # Update weights
        self.optimizer.apply_gradients(zip(new_gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

一些重要注意事项：
compiled\u loss
函数返回的
loss
不得在批处理轴上平均，即我假设它是形状的张量
（批处理大小，）
，而不是标量。

这将导致雅可比矩阵返回形状
（批大小，）+变量形状
的梯度，也就是说，现在每个批元素都有梯度。现在，您可以随心所欲地操纵这些渐变，并且应该在某个时候消除额外的批处理轴（例如平均值）。也就是说，
new_grad
应具有与相应变量相同的形状

关于您的最后一条评论：正如我所提到的，损失函数确实需要在每个数据点返回一个损失，即不能在批次中平均。但是，这还不够，因为如果要将此向量赋给
tape.gradient
，则gradient函数只会对损失值求和（因为它只适用于标量）。这就是为什么需要雅可比矩阵

最后，
jacobian
可能非常慢。在最坏的情况下，运行时间可能会乘以批量大小，因为它需要计算那么多单独的梯度。但是，这在某种程度上是并行的，因此减速可能不会那么糟糕。
不确定这是否是您想要的，所以发布一条评论——假设您使用的是
GradientTape
，您可以使用
jacobian
方法（而不是
gradient
）要获得每个批处理元素的单独渐变，然后对其执行任何操作。@xdurch0非常感谢，我正在深入研究内部方法，我已经看到了带有
gradient
的GradientTape对象。好的，我来看看雅各比。现在看看文档中的方法，它确实返回了每个数据点的梯度w.r.t。太好了，谢谢@xdurch0我想知道
梯度
方法和
雅可比方法之间的确切联系是什么？我在问题中发布了一些工作代码，作为第一个“hello world”，我试图简单地用
gradient
删除这一行，并用
jacobian
替换它。你能帮我吗？我在这方面运气不好，我也很难找到一个好的文档，准确地说明我应该从雅可比的输出中得到什么。是的，稍后我将在回答中发布一些代码。@xdurch0查看更多文档。我看到的行为似乎是仅从损失函数返回平均损失的结果（这是默认值），我看到损失函数实际上可以从批处理中的数据点返回所有损失。使用梯度函数的损失数组可能会返回所有的梯度。谢谢你，这很好，这正是我的想法。我尝试了这个，现在在
fit
行中出现了这个错误。我昨天遇到了这个问题，正在查找，但无法找出问题所在。我是TF版本2.2.0。错误是：
UnrecognizedFlagError:Unknown命令行标志“f”
。似乎是一个可以通过更改版本来修复的错误。我要试试。好的，明白了，谢谢！是的，我升级到了nightly（2.5.~），并且能够在批处理轴上使用
K.mean
重现我第一个版本的工作代码所看到的内容。