你如何计算二维张量的每一行的掩码平均值？答案

【问题标题】：How do you compute a masked mean for each row of a 2D tensor?你如何计算二维张量的每一行的掩码平均值？
【发布时间】：2019-06-05 12:19:51
【问题描述】：

我有一个像这样的二维张量：

[[1. 0. 0. 2. 1. 0. 1.]
 [0. 0. 0. 1. 0. 0. 0.]
 [2. 0. 2. 1. 1. 3. 0.]]

我想计算每一行中每个非零元素的平均值，所以结果是：

[1.25 1.   1.8 ]

我如何使用 TensorFlow 做到这一点？

【问题讨论】：

标签： python tensorflow mean

【解决方案1】：

计算每行掩码均值的一种方法是使用tf.math.unsorted_segment_mean。本质上，您可以每行有一个段 id，然后用一个额外的替换被屏蔽元素的段 id。

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    x = tf.constant([[1., 0., 0., 2., 1., 0., 1.],
                     [0., 0., 0., 1., 0., 0., 0.],
                     [2., 0., 2., 1., 1., 3., 0.]], tf.float32)
    s = tf.shape(x)
    num_rows = s[0]
    num_cols = s[1]
    # Mask for selected elements
    mask = tf.not_equal(x, 0)
    # Make per-row ids
    row_id = tf.tile(tf.expand_dims(tf.range(num_rows), 1), [1, num_cols])
    # Id is replaced for masked elements
    seg_id = tf.where(mask, row_id, num_rows * tf.ones_like(row_id))
    # Take segmented mean discarding last value (mean of masked elements)
    out = tf.math.unsorted_segment_mean(tf.reshape(x, [-1]), tf.reshape(seg_id, [-1]),
                                        num_rows + 1)[:-1]
    print(sess.run(out))
    # [1.25 1.   1.8 ]

但是，由于在这种情况下，掩码恰好适用于非零元素，因此您也可以只使用tf.math.count_nonzero：

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    x = tf.constant([[1., 0., 0., 2., 1., 0., 1.],
                     [0., 0., 0., 1., 0., 0., 0.],
                     [2., 0., 2., 1., 1., 3., 0.]], tf.float32)
    x_sum = tf.reduce_sum(x, axis=1)
    x_count = tf.cast(tf.count_nonzero(x, axis=1), x.dtype)
    # Using maximum avoids problems when all elements are zero
    out = x_sum / tf.maximum(x_count, 1)
    print(sess.run(out))
    # [1.25 1.   1.8 ]

【讨论】：

【解决方案2】：

我们可以使用tf.map_fn 来实现：

x = tf.constant([[1., 0., 0., 2., 1., 0., 1.],
                 [0., 0., 0., 1., 0., 0., 0.],
                 [2., 0., 2., 1., 1., 3., 0.]], tf.float32)
def mean(row):
  mask = tf.not_equal(row, 0.0)
  filtered = tf.boolean_mask(row, mask)
  return tf.reduce_mean(filtered)

out = tf.map_fn(mean, x)

【讨论】：

我不知道这与 jdehesa 在计算速度方面的综合（赞成！）答案相比如何（很高兴承认我根本不知道 tf.math.count_nonzero 或 tf.math.unsorted_segment_mean！）。我在这里发布它作为一种更通用的方式来做我觉得有用的“这类事情”
谢谢，通常tf.map_fn 不是很快，但tf.math.unsorted_segment_mean 也不是太快，所以值得比较。无论如何，是的，对于类似问题，这是一个很好的通用方法。