keras中的自定义损失函数涉及巨大的矩阵乘法

Question

我在 Keras 中编写自定义损失函数时遇到了困难。我有层权重“W”和矩阵“M”。我想做以下操作 trace((W * M) * W') 来计算我的损失函数。 Trace 是对角线元素的总和。在 numpy 中，我会执行以下操作：

np.trace(np.dot(np.dot(W,M),W.T))) or 

def custom_regularizer(W,M):
    sum_reg = 0
    for i in range(W.shape[1]):
        for j in range(i,W.shape[1]):
            vector = W[:,i] - W[:,j]
            sum_reg = sum_reg + M[i,j] * (LA.norm(vector)**2)
    return sum_reg

对于 keras，我写了以下损失函数

def custom_loss(W):

  def lossFunction(y_true,y_pred):    
    loss = tf.trace(K.dot(K.dot(W,K.constant(M)),K.transpose(W)))
    return loss

return lossFunction

问题是 keras 正在计算维度为 200000 * 200000 的整个外部矩阵，导致内存错误。有什么方法可以在不进行整个矩阵计算的情况下获得对角线元素的总和。

和keras损失函数一样怎么做？

Answer 1

如果您遵循一些巧妙的技巧来计算跟踪，您应该不会耗尽内存。例如，您可以参考this。