我使用TensorFlow实现了一个神经网络，它可以处理一批图像，并预测出一个形状为(batch_size, 9, 9, 10)的张量

对于批次中的每张图像，目标（和预测）看起来是这样的。我们预测一个宽度为9，高度为9，深度为10的图像。每个像素的深度对应一个类别，因此我们有10个类别。

例如，像素(0,0)的预测在经过softmax处理后看起来像这样 [0, 0.8, 0, 0, 0, 0, 0.2, 0, 0, 0]，而像素(0,0)的目标看起来像这样 [0, 1, 0, 0, 0, 0, 0, 0, 0]

我想在这两个张量之间使用经典的交叉熵损失函数。

目前我的损失函数是这样实现的：

def calculate_loss(prediction, target):
    '''
    prediction and target have shape (batch_size, 9, 9, 10)
    '''
    loss = 0
    for batch in range(0, BATCH_SIZE):
        for width in range(0,9):
            for heigth in range(0,9):
                loss += tf.losses.categorical_crossentropy(target[batch][width][height] , prediction[batch][width][height])
    return loss

但在高批次大小（1024）下，这样的计算需要20秒。这实在是太慢了。TensorFlow/Keras中有已经实现的函数可以提高这种损失函数的性能吗？或者你知道我可以如何修改代码来提高性能吗？

回答：

tf.nn.softmax_cross_entropy_with_logits(target, prediction)

这就是答案。然后我可以对输出的值进行归约求和以得到相同的结果。别忘了在你的模型中移除softmax层。