我正在使用TensorFlow构建一个处理3D数据的神经网络，该网络应预测输入数据中地标的位置。策略是密集地（对于每个体素）预测实际地标周围半径为r的球体内的类别，并预测指向实际地标位置的偏移向量。这种策略已被证明有效地改进了地标预测。
每个类别概率和偏移向量的组合被视为一个投票，我现在正尝试在TensorFlow中高效地聚合这些投票。
对于输入形状为(70,70,70)和3个不同的地标加上背景类，我从网络中获得两个输出：

• 形状为(70,70,70,3+1)的概率张量
• 形状为(70,70,70,3*3)的偏移向量张量

我想生成3个形状为(70,70,70)的输出热图。现在，对于热图中的每个体素，我需要聚合指向该体素的偏移向量的概率。
我尝试仅使用Python并使用3个for循环，这在我的CPU上需要7秒。这本来是可以接受的，但最终的输入形状将更像是300x300x300，3个for循环将是O(N^3)，因此不可行。
因此，我尝试使用TensorFlow和GPU加速来预先过滤所有不相关的数据。不相关的偏移向量例如是那些对应的类别概率低于某个阈值或超出输入形状范围的。我使用tf.map_fn实现了它，如下所示：

def filter_votes(probs, dists, prob_threshold, num_landmarks, sample_shape: tf.Tensor):    f_sample_shape = tf.cast(sample_shape, tf.float32)    probs = probs[:,:,:,1:] # 背景的概率不相关    indices = tf.where(tf.greater_equal(probs, prob_threshold)) # 只取类别预测超过某个阈值的体素的索引    def get_flatvect(idx):        f_idx    = tf.cast(idx, tf.float32)        return tf.stack([            f_idx[3], # 这是地标编号（从0到2）            probs[idx[0], idx[1], idx[2], idx[3]], # 这是体素被预测为地标的概率            f_idx[0] + dists[idx[0], idx[1], idx[2], idx[3]], # 这是x偏移加上体素的实际x位置            f_idx[1] + dists[idx[0], idx[1], idx[2], idx[3]+3], # 这是y偏移加上体素的实际y位置            f_idx[2] + dists[idx[0], idx[1], idx[2], idx[3]+6] # 这是z偏移加上体素的实际z位置        ])    res = tf.map_fn(get_flatvect, indices, dtype=tf.float32, parallel_iterations=6)    def get_mask(idx):        dist = idx[2:]        return tf.reduce_all(tf.logical_and(tf.greater_equal(dist, 0.), tf.less(dist, f_sample_shape)))    mask = tf.map_fn(get_mask, res, dtype=tf.bool, parallel_iterations=6) # 获取一个掩码，用于过滤超出实际张量形状范围的偏移    res = tf.boolean_mask(res, mask)    return res # 我返回一个2D张量，沿着第二轴包含[num_landmark, probability_value, x_pos, y_pos, z_pos]

然后我用普通的Python聚合过滤后的结果，由于输入数据大幅减少（大多数体素的预测类别概率较低），速度快得多。
问题是即使输入形状为(70,70,70)，过滤操作也需要将近一分钟，GPU利用率很低。尽管我有6个并行迭代，但它几乎比在Python中聚合所有东西慢10倍。我尝试研究了map_fn，我读到TensorFlow可能无法将所有操作放置在GPU上。但即便如此，我认为它应该更快，因为：

• 我有6个并行迭代和6个CPU核心
• 我在开始时使用tf.where预先过滤相关数据，并且只遍历结果索引而不是所有索引

所以似乎我缺乏对正在发生的事情的基本理解。或许有人能解释为什么我的代码效率如此低下？
或者或许有人有更好的想法，如何以向量化的方式聚合我的投票？

回答：