如何使用Tensorflow并行处理大型(已经是CUDA运行的)循环的这个小numba脚本？

Question

作为第一次尝试在显卡上执行例程，我实现了一个包含一个大循环(10亿步)的小函数。虽然它还没有被并行化，但这个脚本在CUDA上运行得相当快：

%%time
from numba import jit, cuda
import numpy as np
from math import sqrt

@cuda.jit
def find_integer_solutions_cuda(arr):
    i=0
    for x in range(0, 1000000000+1):
        y = float(x**6-4*x**2+4)
        sqr = int(sqrt(y))
        if sqr*sqr == int(y):
            arr[i][0]=x
            arr[i][1]=sqr
            arr[i][2]=y
            i+=1

arr=np.zeros((10,3))
find_integer_solutions_cuda[128, 255](arr)

print(arr)

此脚本工作良好，并在5分钟内使用线程配置[128, 255] (其他配置减缓它)在一台具有128 v4内存、英特尔Xeon E5-2630 v4、2.20GHz处理器和两张特斯拉V100类型的图形卡上完成，每块内存为16 v4。它产生：

[[0.00000000e+00 2.00000000e+00 4.00000000e+00]
 [1.00000000e+00 1.00000000e+00 1.00000000e+00]
 [7.08337220e+07 2.64700090e+09 7.00661374e+18]
 [6.56031067e+08 2.29447517e+09 5.26461630e+18]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]
 [0.00000000e+00 0.00000000e+00 0.00000000e+00]]
CPU times: user 58.5 s, sys: 4min 5s, total: 5min 4s
Wall time: 5min 4s

背景:我正在试验大型循环执行(相当简单)算术/数学任务的运行时行为和性能结果。上面的代码片段中的公式只是一个例子。我观察到的是，使用numba的@JIT-decorator (5秒)，使用这段代码效果最好。使用gmpy2，相同的任务在15分钟内完成。使用不优化(只是纯粹的numpy)，相同的例程几乎需要2个小时。我很好奇，当我通过Tensorflow在GPU驱动的机器上并行化这个例程时，它会是什么样子？简单地说，这个例行公事高达10亿人：

5秒 (numba/JIT) <5 5min (numba/CUDA.JIT) <15 5min (gmpy2) <2小时(普通numpy)

我的问题：，我想通过将这个小脚本转移到tensorflow来扩展这个实验，并加入真正的并行性，如果它能够“用任务填充张量”来进行并行化的话。我在考虑用数字填充张量，然后图形卡将这些数字并行地插入方程y = float(x**6-4*x**2+4)中，进行检查，并填充结果数组。

Answer 1

如果您只想测量GPU的速度，而不是解决高达10亿x的任务，那么我可以建议如下。

由于您的多项式的顺序是x ** 6，这意味着它的值只有在x低于1000时才会低于10^18。众所周知，18 * 10^18是最大uint64的近似值。

但是我们正在对sqrt浮动值进行double操作，它最多是53位精度(尾数为53位)。此外，考虑到平方尺的四舍五入误差，我们可以在200-400以下的x中加入双倍。

如果您溢出双精度，那么sqrt结果是不精确的，您不能再做比较sqr * sqr == int(y)，因为太大的错误。你不会找到正方形，也不会解决你的任务(除非你只想检查GPU的速度)。

不幸的是，Tensorflow没有无限精度的整数或浮动算法。它最多只能浮动64/ It 64。否则我们可以用X来做我们想要的大小。由于精度有限，我们只能检查较小的X值。

为了精确地解决任务，我为x提供了下面的Tensorflow代码，限制在200以下。但是，要使计算量与代码段中的计算量相同，我只需多次重复小值X，直到得到10亿个值。

您可以将param N = 10 ** 6 (值的总数)调整为您想要的任何值，但是它应该是block = 200的倍数，并且块不应该大于200，您可以调整这两个参数。

如果您只想测试GPU的速度，而不是精确地解决任务，您也可以设置block = N。在这种情况下，计算多项式和平方根的误差会很大，但是有了这个误差，您仍然可以检查GPU的速度。对于大型X，你只会得到错误的最终结果。

在我的代码末尾，我通过显示元组(x, y, y^2)列表来输出所有找到的方块。同时输出总运行时间。

在运行代码之前，不要忘记通过python -m pip install tensorflow安装Tensorflow。

在我的1.2 Ghz 2核笔记本电脑上，用下面的Tensorflow代码处理10亿个值需要89秒钟。

import time, tensorflow as tf

N = 10 ** 6
block = 200
nblocks = N // block

tb = time.time()

x = tf.repeat(tf.range(0, block, dtype = tf.int64)[None, :], nblocks, axis = 0)
f = x ** 6 - 4 * x ** 2 + 4
r = tf.cast(tf.math.sqrt(tf.cast(f, tf.float64)) + 0.5, tf.int64)
is_square = tf.where(r * r == f)

vals = [tuple(map(int, [i, j, x[i, j], r[i, j], f[i, j]])) for i, j in is_square]
print('(x, y, y^2):\n', [(x0, y0, y0_sqr) for i, j, x0, y0, y0_sqr in vals if i == 0])
print(f'Time {time.time() - tb:.03f} sec')

输出：

(x, y, y^2):
 [(0, 2, 4), (1, 1, 1)]
Time 88.912 sec