优化tf.map_fn的GPU使用

Question

我正在使用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:] # probability of background is irrelevant
    indices = tf.where(tf.greater_equal(probs, prob_threshold)) # take only the indices of voxels, that have a category prediction over a certain threshold

    def get_flatvect(idx):
        f_idx    = tf.cast(idx, tf.float32)
        return tf.stack([
            f_idx[3], # this is the landmark number (goes from 0 to  2)
            probs[idx[0], idx[1], idx[2], idx[3]], # this is the predicted probability for the voxel to be the landmark
            f_idx[0] + dists[idx[0], idx[1], idx[2], idx[3]], # this is the x offset+ the actual x-position of the voxel
            f_idx[1] + dists[idx[0], idx[1], idx[2], idx[3]+3], # this is the y offset+ the actual y-position of the voxel
            f_idx[2] + dists[idx[0], idx[1], idx[2], idx[3]+6] # this is the z offset+ the actual z-position of the voxel
        ])
    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) # get a mask that filters offsets that went out of bounds of the actual tensor shape
    res = tf.boolean_mask(res, mask)
    return res # I return a 2D-Tensor that contains along the 2nd axis [num_landmark, probability_value, x_pos, y_pos, z_pos]

然后我在纯Python中汇总筛选结果，由于输入数据明显减少（大多数体素具有较低的预测分类概率），所以速度更快。问题是，即使使用输入形状（70,70,70），过滤操作也仅需花费近一分钟，GPU利用率很低。尽管我有6个并行迭代，但速度几乎比Python中汇总所有内容慢10倍。我尝试研究了map_fn，并且读到tf可能无法将所有操作放在GPU上。但即使如此，我认为应该会更快，因为：

• 我有6个并行迭代和6个CPU核心
• 我在开头使用tf.where对相关数据进行预过滤，并仅在结果索引上迭代，而不是在所有索引上迭代

所以似乎我缺乏一些基本的理解。也许有人可以澄清为什么我的代码效率如此低下？或者也许有人有更好的想法如何以矢量化的方式汇总我的投票？

- B4ckup

---

看起来你可以使用 tf.gather_nd 来替代这两个 Python 函数。你能否在代码中添加有关张量形状的注释，就像你为 res 所做的那样？这会让不太熟悉你的问题的人更容易理解。 - McAngus

@McAngus 感谢您的关注！事实证明，您是正确的，应该在这里使用tf.gather_nd。jdehesa在下面提交了一个整洁的解决方案 :) - B4ckup

Answer 1

您可以像这样将函数向量化：

import tensorflow as tf

def filter_votes_vec(probs, dists, prob_threshold, num_landmarks, sample_shape: tf.Tensor):
    probs = probs[:, :, :, 1:]
    indices = tf.where(probs >= prob_threshold)
    landmark = tf.to_float(indices[:, 3])
    p = tf.gather_nd(probs, indices)
    indices_dists = tf.stack([
        indices,
        tf.concat([indices[..., :-1], indices[..., -1:] + 3], axis=-1),
        tf.concat([indices[..., :-1], indices[..., -1:] + 6], axis=-1)
    ], axis=1)
    d = tf.gather_nd(dists, indices_dists) + tf.to_float(indices[:, :3])
    res = tf.concat([tf.expand_dims(landmark, 1), tf.expand_dims(p, 1), d], axis=1)
    mask = tf.reduce_all((d >= 0) & (d < tf.cast(sample_shape, tf.float32)), axis=1)
    res =  tf.boolean_mask(res, mask)
    return res

使用IPython进行快速测试和基准测试：

import tensorflow as tf
import numpy as np

with tf.Graph().as_default(), tf.Session() as sess:
    np.random.seed(100)
    probs = np.random.rand(70, 70, 70, 3 + 1).astype(np.float32)
    probs /= probs.sum(-1, keepdims=True)
    probs = tf.convert_to_tensor(probs, tf.float32)
    dists = tf.convert_to_tensor(100 * np.random.rand(70, 70, 70, 3 * 3), tf.float32)
    prob_threshold = tf.convert_to_tensor(0.5, tf.float32)
    num_landmarks = tf.convert_to_tensor(3, tf.int32)  # This is not actually used in the code
    sample_shape = tf.convert_to_tensor([50, 60, 70], tf.int32)

    result = filter_votes(probs, dists, prob_threshold, num_landmarks, sample_shape)
    result_vec = filter_votes_vec(probs, dists, prob_threshold, num_landmarks, sample_shape)
    value, value_vec = sess.run([result, result_vec])
    print(np.allclose(value, value_vec))
    # True
    %timeit sess.run(result)
    # CPU
    # 2.55 s ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
    # GPU
    # 54 s ± 596 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
    %timeit sess.run(result_vec)
    # CPU
    # 63.2 µs ± 781 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
    # GPU
    # 216 µs ± 2.29 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

据推测，GPU的荒谬时间是由于TensorFlow不断在CPU和GPU之间交换数据，这是一种相当昂贵的操作。

- jdehesa