大数组需要循环单个赋值时适合在GPU上还是CPU上运算？

原创已于 2023-06-17 11:14:53 修改 · 1.3k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #python #tensor #list

于 2022-10-23 19:49:48 首次发布

深度学习同时被 3 个专栏收录

38 篇文章

订阅专栏

PyTorch

12 篇文章

订阅专栏

Python

9 篇文章

订阅专栏

本文探讨了一位算法工程师在训练模型时遇到的GPU运算效率问题。通过分析发现在一个涉及大量循环赋值的操作中，使用列表而非张量导致了显著的性能下降。优化方案包括将数据结构转换为张量并在GPU上运算，进一步将循环赋值移到CPU上，显著减少了执行时间，从17秒降至100毫秒，提升了170倍效率。实验结果显示，在GPU上进行循环运算的效率较低，而使用Python列表比numpy数组更快。这一发现强调了选择合适的数据结构和运算平台对于优化计算性能的重要性。

训练某个模型时发现每次作validation时总是耗时很长，4000个文件的小数据集，2张GPU卡，一个epoch都需要八九个小时，调查原因最后定位到下面一段简单的循环赋值代码:

score_threshs = torch.zeros(top_scores.shape, device=top_scores.device).type_as(top_scores)
for i, label in enumerate(top_labels):
    score_threshs[i] = self._score_threshold[label]

这个top_labels是torch.tensor类型，size (642816,) ，for循环每次执行完需要17秒多，这样的话每一个文件的推理在这里就需要至少消耗17秒多，两个2GPU上并发推理4000个文件自然需要九个小时以上，慢得不可忍受，检查代码发现self._score_threshold是个普通的list不是tensor，这显然是不合理的，每次都要将数据从CPU读取到GPU不慢才怪呢，于是在类创建时把它创建成torch.tensor实例，让score_threshs和self._score_threshold之间的赋值全在GPU上，修改后发现节省了点时间，但是上面的for循环还是需要12秒多！这里还能不能再想办法节省时间呢？但是循环单个赋值是不可避免的，怎么办？j坚持在GPU上运算究竟能减少到多少时间？试验了一下，什么都不干，只是简单地再top_labels上循环看这样需要多少时间？

for i, label in enumerate(top_labels):
   pass

上面的循环居然还需要800多毫秒！这样看来GPU上作循环的效率就很差，单个值赋值的效率就更差！那把数据全拿到CPU上进行循环赋值运算，赋值完了再放回GPU会怎么样？于是改成下面这样：

top_labels_cpu =  top_labels.tolist()
score_threshs_cpu = []
for i in range(len(top_labels_cpu)):
   score_threshs_cpu.append(self._score_threshold[top_labels_cpu[i]])
score_threshs = torch.tensor(score_threshs_cpu, device=top_scores.device).type_as(top_scores)

结果发现for循环代码每次执行只需100毫秒左右，把score_threshs数据从CPU放回GPU只需几个毫秒，这样算来，所花时间从17秒多减少到了100毫秒左右！性能提升到了原来的170倍！看来这种运算在CPU上运算比GPU上效率成百倍提升！

另外也实验了一下，如果采用numpy的ndarray来进行上面的循环运算比上面采用python的list要慢一点，所需时间为140毫秒左右！这点倒也可以理解，一般规律是数据类型包装越多运算越慢。

所以就上面的循环赋值运算的速度性能来说，我实验得到的结果是: python list > numpy.ndarray > torch.tensor