python - 如何在 tensorflow 中使用指定掩码对张量轴进行平均

Question

例如:

我有一个输入张量(input)，形状为(?,10) dtype=float32，第一个维度表示批量大小。

还有一个掩码张量(mask)，形状为(?,10)。 mask[sample_number] 就像[True,True,False,...]，表示掩码

标签张量(avg_label)，形状为(?,)，表示每个样本的掩码位置的正确平均值

我想训练模型，但找不到获得输出的好方法。

tf.reduce_...(例如tf.reduce_mean)函数似乎不支持关于屏蔽的论点。

我尝试tf.boolean_mask，但它会将输出形状展平为只有一维，抛出sample_number维度，因此无法区分样本

我考虑过tf.where，例如:

masked=tf.where(mask,input,tf.zeros(tf.shape(input)))
avg_out=tf.reduce_mean(masked,axis=1)
loss=tf.pow(avg_out-avg_label,2)

但是上面的代码肯定不起作用，因为 False 设置为 0 会改变 avg。如果使用 np.nan ，它总是会得到 nan。我想知道在进行reduce操作时是否有一个代表缺席的值。

我该怎么做？

Answer 1

您可以使用tf.ragged.boolean_mask来保持维度。

tf.reduce_mean(tf.ragged.boolean_mask(x, mask=mask), axis=1)