python - 优化 tf.map_fn 的 GPU 使用率
问题描述
我正在 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对相关数据进行预过滤,并且只迭代结果索引而不是所有索引
因此,我似乎对正在发生的事情缺乏一些基本的了解。也许有人可以澄清为什么我的代码效率如此之低?
或者也许有人有更好的想法如何以矢量化方式汇总我的选票?
解决方案
您可以像这样矢量化您的函数:
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 之间不断地交换数据,这是相当昂贵的。
推荐阅读
- servlets - Micronaut Servlet 支持:发现经典的 HttpServlet 实现
- blazor-client-side - Blazor webAssembly 调用不存在的 api 路由返回 200 成功与 index.html
- ios - UITableViewController 底部多余的空格
- python - 我在 Python 中使用 def 的方式有什么问题
- html - 如何让背景图像覆盖整个 div?
- excel - 通过指定帐户发送电子邮件
- php - 结果需要一个或多个查询
- r - 合并多个 CSV 文件时从某些文件中删除列
- php - 来自 Json 数据的 PHP 中的多维数组挑战
- perl - 有没有办法阻止无效的命令行参数