python - 提高分类交叉熵在四维预测和目标上的性能

Question

我有一个在 tensorflow 中实现的神经网络，它获取一批图像并预测形状为 (batch_size, 9, 9, 10) 的张量

对于批处理中的每个图像，目标(和预测)如下所示。 我们预测宽度为 9、高度为 9、深度为 10 的图像。每个像素的深度对应一个类，因此我们有 10 个类。

所以例如使用softmax对像素(0,0)的预测如下所示 [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中是否有任何已经实现的功能可以提高这种损失的性能？ 或者你知道我是否可以更改代码以提高性能？

Answer 1

tf.nn.softmax_cross_entropy_with_logits(target, prediction)

就是答案。 然后我可以减少输出值的总和以获得相同的结果。 不要忘记删除模型中的 sofmax 层。