tensorflow - 当只需要第一个元素时，为什么要创建一个新轴？

Question

首先，很抱歉标题含糊

由于我有兴趣了解有关 TensorFlow 和图像分割的更多信息，因此我正在学习他们的教程 ( https://www.tensorflow.org/tutorials/images/segmentation )。然而，我注意到一些我无法完全掌握的东西，在谷歌搜索之后也没有掌握。

在本节中:

def create_mask(pred_mask):
    pred_mask = tf.argmax(pred_mask, axis=-1)
    pred_mask = pred_mask[..., tf.newaxis]
    return pred_mask[0]

首先为 pred_mask 向量创建一个新轴，然后仅选择第一个元素的原因是什么？为什么和我想象的不一样，如下图:

def create_mask(pred_mask):
    pred_mask = tf.argmax(pred_mask, axis=-1)
    return pred_mask

Answer 1

这只是为了让图像保持 3D 张量。例如，如果您的预测形状为 (1, 256, 256, 10)(一批 256x256 图像 10 个类)，则在 tf.argmax() 之后，您可以将收到形状为 (1, 256, 256) 的张量(一批没有 channel 的 256x256 图像)。但通常情况下，如果图像采用 HWC 格式(Height, Width, Channel) 而不仅仅是(Height, Width)，则处理图像会更容易。例如，如果您使用 matplotlib 或 OpenCV，通常需要 HWC 图像。