Python 与itertools和numba并行

我已经在一个项目上工作了一段时间，这个项目需要计算一些非常大的数据集，并且很快就超越了我微薄的Excel知识所能处理的任何事情。在过去的几天里，我开始学习Python，它有助于处理我正在处理的数据的大小，但是这些数据集的估计处理时间看起来非常长（在我的笔记本电脑上可能需要几百年）

这里的瓶颈是一个可能产生数万亿或万亿结果的等式，因为它计算6个不同列表的每一个组合，并通过一个你将在代码中看到的等式来运行它。这段代码工作得很好，但对于比我包含的示例更大的数据集来说是不可行的。一个真正的数据集应该更像Set1S、2S和3S，每个都有50个项目，set12a…每个大约有2500个项目（在本例中是50x50。这些集合的长度总是等于前3个列表的平方，但我在这里保持简短）

我很清楚，结果的数量是绝对巨大的，但我想从尽可能大的数据集开始，这样我就可以看到在绘制累积百分比直方图时，我可以减少多少输入大小而不会对结果造成很大影响

'Calculator'
import numpy as np
Set1S = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
Set2S = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
Set3S = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
Set12A = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
Set23A = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
Set13A = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
'Define an empty array to add results'
BlockVol = []
from itertools import product 
'itertools iterates through all combinations of lists'
for i,j,k,a,b,c in product(Set1S, Set2S, Set3S, Set12A, Set23A, Set13A):
    'This is the bottleneck equation, with large input datasets'
    BlockVol.append((abs(i*j*k*np.sin(a)*np.sin(b)*np.sin(c))))
arr = np.array(BlockVol)
'manipulate the result list a couple ways'
BlockVol = np.cbrt(BlockVol)
BlockVol = BlockVol*12
'quick check to size of results list'
len(BlockVol)

这花了我大约3分钟左右的时间，仅仅是盯着时钟，就得到了1130万个结果

我在最后一天左右了解了@njit，prange，但我有点困在试图将我的作品翻译成这种格式。我有一台桌面电脑，有一个相当好的GPU，所以我想我可以把速度提高很多。我很清楚下面的代码是一场垃圾大火，没有任何作用，但我希望我至少能让大家明白我在做什么

似乎应该用我的6个输入列表定义一个函数，但我不知道如何将itertools产品和njit融合在一起

import numpy as np
from itertools import product
from numba import njit, prange

@njit(parallel = True)
def BlockVolCalc(Set1S, Set2S, Set3S, Set12A, Set23A, Set13A):
    numRows =Len(Set12A)
    BlockVol = np.zeros(numRows)
    for i,j,k,a,b,c in product(Set1S, Set2S, Set3S, Set12A, Set23A, Set13A):
        BlockVol.append((abs(i*j*k*np.sin(a)*np.sin(b)*np.sin(c))))
arr = np.array(BlockVol)
BlockVol = np.cbrt(BlockVol)
BlockVol = BlockVol*12
len(BlockVol)

任何帮助都是非常感谢的，因为这一切都是非常新的和压倒性的

---

谢谢大家!

我只是用NumPy代码解决了你的任务，如果可能的话，使用NumPy而不是笨重的Numba总是更好的。下一个只有NumPy的代码将和使用Numba的相同解决方案一样快

我的代码比您的参考代码快2800倍，时间是在代码末尾测量的

在下一个代码中，

BlockValCalcRef（…）

函数只是作为函数组织的参考代码。而

BlockVolCalc（…）

是我基于NumPy的函数，应该可以提供很多加速。最后，我断言

np.allclose（…）

，以检查两个解决方案是否给出相同的结果

我还简化了位集的创建，使用一个

N

param来生成集，在现实世界中，您只需提供必要的集

为了解决这个任务，我做了几件事：

我没有为相同的值多次计算

np.sin（…）

，而是为

Set12A、Set23A、Set13A预先计算了一次。还为所有集合预计算了
np.abs（…）

为了计算叉积，我使用了numpy数组索引的特殊方式，如

[None，None，：，None，None，None]

这允许我们使用所谓的popular

我还知道如何进一步改进代码，使其速度提高6倍，但我认为即使以目前的高速运行，您也可以在几秒钟内填满机器的整个RAM。下一个改进的想法是，目前交叉积计算

6

数字的每一步积，而不是这个，可以计算

K-1

集的积，然后将该数组乘以

K-th

集，以获得

K

集积。这将使

6

时间加快（因为有

6

集），因为您只需要一次乘法，而不是

6

更新：根据上面的段落，我已经实现了第二个改进版的函数

BlockVolCalc2（…）

。它具有

2800x

加速比，对于较大的

N

它可能会更快

输出：

base time 2.7569 sec
improved time 0.0015 sec, speedup 1834.83x
improved2 time 0.001 sec, speedup 2755.09x

嵌入到初始代码中的我的函数将如下所示

此外，我还创建了基于的代码变体，它将使用所有的CPU内核和GPU，此代码需要通过

python一次性安装tensorflow-m pip install-升级numpy tensorflow

：

import numpy as np

N = 18
Set1S, Set2S, Set3S, Set12A, Set23A, Set13A = [np.arange(1 + i, N + 1 + i) for i in range(6)]
dtype = np.float32

def Prepare(Set1S, Set2S, Set3S, Set12A, Set23A, Set13A):
    import numpy as np
    Set12A, Set23A, Set13A = np.sin(Set12A), np.sin(Set23A), np.sin(Set13A)
    return [np.abs(s).astype(dtype) for s in [Set1S, Set2S, Set3S, Set12A, Set23A, Set13A]]
    
sets = Prepare(Set1S, Set2S, Set3S, Set12A, Set23A, Set13A)

def ProcessNP(sets):
    import numpy as np
    res = np.ones((1,), dtype = dtype)
    for s in reversed(sets):
        res = (s[:, None] * res[None, :]).ravel()
    res = np.cbrt(res) * 12
    return res

def ProcessTF(sets, *, state = {}):
    if 'graph' not in state:
        import os
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
        import numpy as np, tensorflow as tf
        tf.compat.v1.disable_eager_execution()
        cpus = tf.config.list_logical_devices('CPU')
        #print(f"CPUs: {[e.name for e in cpus]}")
        gpus = tf.config.list_logical_devices('GPU')
        #print(f"GPUs: {[e.name for e in gpus]}")
        print(f"GPU: {len(gpus) > 0}")
        state['graph'] = tf.Graph()
        state['sess'] = tf.compat.v1.Session(graph = state['graph'])
        #tf.device(cpus[0].name if len(gpus) == 0 else gpus[0].name)
        with state['sess'].as_default(), state['graph'].as_default():
            res = tf.ones((1,), dtype = dtype)
            state['inp'] = []
            for s in reversed(sets):
                sph = tf.compat.v1.placeholder(dtype, s.shape)
                state['inp'].insert(0, sph)
                res = sph[:, None] * res[None, :]
                res = tf.reshape(res, (tf.size(res),))
            res = tf.math.pow(res, 1 / 3) * 12
            state['out'] = res
        def Run(sets):
            with state['sess'].as_default(), state['graph'].as_default():
                return tf.compat.v1.get_default_session().run(
                    state['out'], {ph: s for ph, s in zip(state['inp'], sets)}
                )
        state['run'] = Run

    return state['run'](sets)

# ------------ Testing ------------

npa, tfa = ProcessNP(sets), ProcessTF(sets)
assert np.allclose(npa, tfa)

from timeit import timeit

print('Nums:', round(npa.size / 10 ** 6, 3), 'M')
timeit_num = 2
print('NP:', round(timeit(lambda: ProcessNP(sets), number = timeit_num) / timeit_num, 3), 'sec')
print('TF:', round(timeit(lambda: ProcessTF(sets), number = timeit_num) / timeit_num, 3), 'sec')

在我的2核CPU上，它打印：

GPU: False
Nums: 34.012 M
NP: 3.487 sec
TF: 1.185 sec

---

哇，反应真快。我仍在仔细检查你的代码，以确保我完全理解。所以，看起来我可以移除BlockValCalcRef，然后在BlockValcalc中离开，是吗？请原谅我的无知，但是BlockVolCalc生成的列表的名称是什么？我可能能够访问一些云计算资源，如果这有助于克服RAM问题的话。我的测试台式电脑有128 GB的RAM，所以我会先测试一下。比你已经做的还要快是很了不起的，因此，我必须检查我的系统是否因电流过高而过载method@JCombs根据我在回答中的想法，我刚刚发布了第二个更快两倍的解决方案。@如果不考虑在多核CPU上并行化或使用GPU的可能性，JCombs Now算法可能是最快的，甚至可以使用多台计算机。@JCombs您只需复制

BlockVolCalc2（）

函数的代码并使用它，无需其他操作。看看你问题中的第一个代码，行

arr=np.array（BlockVol）

，所以我代码中的所有函数都会生成这个结果

arr

，然后你需要做额外的操作

BlockVol=np.cbrt（BlockVol）

和

BlockVol=BlockVol*12

。哇，再次感谢。一旦我做了一些测试，我会再次更新，但我认为多核CPU或GPU肯定是未来的选择。

GPU: False
Nums: 34.012 M
NP: 3.487 sec
TF: 1.185 sec