Tensorflow中应该如何精确地进行Tensor掩蔽和索引编制？

我已经使用TF已有2年了，在每个项目中，我都会冒出很多无意义的错误来掩盖信息，这些错误通常无济于事，也没有指出实际上是什么错误。或更糟糕的是，结果是错误的，但没有错误。我总是在训练循环之外使用伪数据测试代码，这很好。但是在训练（称合适）中，我不明白TensorFlow到底期望什么。仅举一个例子，有经验的人可以告诉我为什么此代码对于二进制交叉熵不起作用，结果是错误的，并且在这种情况下模型不收敛但没有错误：

class MaskedBXE(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        y_true = tf.squeeze(y_true)
        mask = tf.where(y_true!=2)
        y_true = tf.gather_nd(y_true, mask)
        y_pred = tf.gather_nd(y_pred, mask)
        loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
        return tf.reduce_mean(loss)

虽然这可以正常工作：

class MaskedBXE(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        mask = tf.where(y_true!=2, True, False)
        y_true = y_true[mask]
        y_pred = y_pred[mask]
        loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
        return tf.reduce_mean(loss)

对于一个明确的例子，情况恰恰相反。我无法将遮罩用作y_pred [mask]或y_pred [mask [0]]之类的索引，也不能使用tf.squeeze（）等。但是使用tf.gather_nd（）可以。我总是尝试我认为可能的所有组合，但我不明白为什么这么简单的事情会如此艰巨而痛苦。派托克也是这样吗？如果您知道Pytorch没有类似的烦人细节，我们很乐意切换。

编辑1：它们可以在训练循环或图形模式之外正常工作，更准确地说。

y_pred = tf.random.uniform(shape=[10,], minval=0, maxval=1, dtype='float32')
y_true = tf.random.uniform(shape=[10,], minval=0, maxval=2, dtype='int32')

# first method
class MaskedBXE(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
    def call(self, y_true, y_pred):
        y_true = tf.squeeze(y_true)
        mask = tf.where(y_true!=2)
        y_true = tf.gather_nd(y_true, mask)
        y_pred = tf.gather_nd(y_pred, mask)
        loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
        return tf.reduce_mean(loss)

    def get_config(self):
        base_config = super().get_config()
        return {**base_config}

# instantiate
mbxe = MaskedBXE()
print(f'first snippet: {mbxe(y_true, y_pred).numpy()}')


# second method
class MaskedBXE(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        mask = tf.where(y_true!=2, True, False)
        y_true = y_true[mask]
        y_pred = y_pred[mask]
        loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
        return tf.reduce_mean(loss)
    
    def get_config(self):
        base_config = super().get_config()
        return {**base_config}
    
# instantiate
mbxe = MaskedBXE()
print(f'second snippet: {mbxe(y_true, y_pred).numpy()}')

第一个片段：1.2907861471176147

第二段：1.2907861471176147

编辑2：在以@jdehesa的建议方式在图形模式下打印损失后，它们有所不同，它们不应这样做：

class MaskedBXE(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        # first
        y_t = tf.squeeze(y_true)
        mask = tf.where(y_t!=2)
        y_t = tf.gather_nd(y_t, mask)
        y_p = tf.gather_nd(y_pred, mask)
        loss = tf.keras.losses.binary_crossentropy(y_t, y_p)
        first_loss =  tf.reduce_mean(loss)
        tf.print('first:')
        tf.print(first_loss, summarize=-1)
        # second
        mask = tf.where(y_true!=2, True, False)
        y_t = y_true[mask]
        y_p = y_pred[mask]
        loss = tf.keras.losses.binary_crossentropy(y_t, y_p)
        second_loss = tf.reduce_mean(loss)
        tf.print('second:')
        tf.print(second_loss, summarize=-1)
        return second_loss

第一：

0.814215422

第二：

0.787778914

第一：

0.779697835

第二：

0.802924752

。。。