Python Theano扫描用于阵列上的快速计算

我试图使用Theano来加速已经在numpy中实现的代码，该代码对数组中的元素求和。在numpy中，函数如下所示

import numpy as np

def numpy_fn(k0, kN, x):
    output = np.zeros_like(x)
    for k in range(k0, kN+1):
        output += k*x
    return output

用一个样本电话

>>> numpy_fn(1, 3, np.arange(10))
array([  0.,   6.,  12.,  18.,  24.,  30.,  36.,  42.,  48.,  54.])

上述函数的theano等价物为

import theano
import theano.tensor as tt

k  = tt.scalar('k')
k0 = tt.scalar('k0')
kN = tt.scalar('kN')
x  = tt.vector('x')

def fn(k, sumtodate):
    return sumtodate + k*x

rslt, updt = theano.scan(fn=fn, 
                         outputs_info=tt.zeros_like(x),
                         sequences=tt.arange(k0, kN+1))
theano_fn = theano.function(inputs=[k0, kN, x], 
                            outputs=rslt[-1])

调用时，这将提供正确的输出

theano_fn(1, 3, np.arange(10))
array([  0.,   6.,  12.,  18.,  24.,  30.,  36.,  42.,  48.,  54.])

然而，当我对两者进行基准测试时，在我的计算机上，numpy函数的速度比theano快了三倍

%timeit theano_fn(1, 1000, np.ones(10000))
10 loops, best of 3: 21.5 ms per loop

%timeit numpy_fn(1, 1000, np.ones(10000))
100 loops, best of 3: 7.9 ms per loop

既然theano将outerloop转换为C，那么它不应该比Python更快吗？如何提高theano代码的速度

编辑：

我知道numpy中的暴力代码可以使用求和进行优化，但我之所以选择theano路线，是因为我对输出更新可以是

k

和

x

的任何通用函数的情况感兴趣

output += x**k
output += exp(k*x)
output += (x-k)**2

output+=k*x

只是说明这一点的一个具体例子。使用数学符号，我试图实现的是一个快速求和

\sum{k=k0}{kN}f（k，x）

，其中

k0

和

kN

是整数，

x

是向量，

f

可以是

k

和

x

的任何一般函数，就像上面给出的函数一样

import numpy as np

def f(k, x):
    return x**k

def numpy_fn(k0, kN, x):
    output = np.zeros_like(x)
    for k in range(k0, kN+1):
        output += f(k, x)
    return output

---

我希望通过使用theano，我能够优化outter循环，并获得比brute numpy解决方案更快的解决方案。

对于您正在执行的操作，您可以简单地将所有元素从

k0

到

kN

相加，以获得标量，它必须用于缩放

x

，以获得所需的输出。这样，您就可以在NumPy环境中使用矢量化方法。 使用的实现如下所示-

np.arange(k0,kN+1).sum()*x

您还可以使用执行求和，这样性能可能会稍好一些，如下所示-

np.einsum('i->',np.arange(k0,kN+1))*x

运行时测试和输出验证-

In [74]: k0 = 10; kN = 10000

In [75]: x = np.random.rand(20000)

In [76]: np.allclose(numpy_fn(k0,kN,x),np.arange(k0,kN+1).sum()*x)
Out[76]: True

In [77]: np.allclose(numpy_fn(k0,kN,x),np.einsum('i->',np.arange(k0,kN+1))*x)
Out[77]: True

In [78]: %timeit numpy_fn(k0,kN,x)
1 loops, best of 3: 460 ms per loop

In [79]: %timeit np.arange(k0,kN+1).sum()*x
10000 loops, best of 3: 54.9 µs per loop

In [80]: %timeit np.einsum('i->',np.arange(k0,kN+1))*x
10000 loops, best of 3: 49.7 µs per loop

---

基于Divakar的答案

Theano能够超越numpy的情况非常具体。一般来说，只有当计算涉及对大张量的矢量运算时，Theano才会比numpy表现得更好

在这种情况下，可以在numpy中非常有效地执行操作。通过使用循环的标准结果，根本不需要使用循环。此处

n=kN-k0+1

是要求和的项数

numpy.arange(k0, kN + 1).sum() == (kN - k0 + 1) * (k0 + kN) / 2

如果出于性能以外的原因（例如，为了获得梯度，或作为某些较大符号计算的一部分）需要使用THANO，则可以计算相同的结果，而无需使用sum或scan，就像在numpy中一样

下面的代码实现了原始的numpy和Theano方法，并将它们与Divakar的numpy方法（以及arange sum方法的my Theano版本）以及使用算术序列结果标准和的my numpy和Theano方法进行比较

import numpy
import timeit
import itertools
import theano
import theano.tensor as tt


def numpy1(k0, kN, x):
    output = numpy.zeros_like(x)
    for k in range(k0, kN + 1):
        output += k * x
    return output


def numpy2(k0, kN, x):
    return numpy.arange(k0, kN + 1).sum() * x


def numpy3(k0, kN, x):
    return numpy.einsum('i->', numpy.arange(k0, kN + 1)) * x


def theano1_step(k, s_tm1, x):
    return s_tm1 + k * x


def compile_theano1():
    k0 = tt.lscalar()
    kN = tt.lscalar()
    x = tt.vector()
    outputs, _ = theano.scan(theano1_step, sequences=[tt.arange(k0, kN + 1)], outputs_info=[tt.zeros_like(x)],
                             non_sequences=[x], strict=True)
    return theano.function([k0, kN, x], outputs=outputs[-1])


def compile_theano2():
    k0 = tt.lscalar()
    kN = tt.lscalar()
    x = tt.vector()
    return theano.function([k0, kN, x], outputs=tt.arange(k0, kN + 1).sum() * x)


def numpy4(k0, kN, x):
    return ((kN - k0 + 1) * (k0 + kN) / 2) * x


def compile_theano4():
    k0 = tt.lscalar()
    kN = tt.lscalar()
    x = tt.vector()
    return theano.function([k0, kN, x], outputs=((kN - k0 + 1) * (k0 + kN) / 2) * x)


def main():
    iteration_count = 10
    k0 = 10
    kN = 10000
    x = numpy.random.standard_normal(size=(20000,)).astype(theano.config.floatX)

    functions = [numpy1, numpy2, numpy3, numpy4, compile_theano1(), compile_theano2(), compile_theano4()]
    function_count = len(functions)
    results = numpy.empty((iteration_count * function_count, x.shape[0]), dtype=theano.config.floatX)
    times = numpy.empty((iteration_count * function_count,), dtype=theano.config.floatX)

    for iteration in xrange(iteration_count):
        for function_index, function in enumerate(functions):
            start = timeit.default_timer()
            results[iteration * function_count + function_index] = function(k0, kN, x)
            times[iteration * function_count + function_index] = timeit.default_timer() - start

    for result1, result2 in itertools.izip(results[0::2], results[1::2]):
        assert numpy.allclose(result1, result2)

    for function_name, function_index in itertools.izip(
            ('numpy1', 'numpy2', 'numpy3', 'numpy4', 'theano1', 'theano2', 'theano4'),
            xrange(function_count)):
        time = times[function_index::function_count].mean()
        print '%8s %.8f' % (function_name, float(time))


main()

在我糟糕的台式计算机上，它使用CPU（而不是GPU）进行Theano计算，我得到以下计时（以秒为单位，越低越好）：

在这种特殊情况下，在GPU上运行Theano代码不太可能有用，除非

x

非常大。但即便如此，将

x

复制到GPU内存中的成本也可能会抵消并行元素乘法的任何收益

编辑

要在问题的编辑版本中解决新的部分

Theano不适合显式循环。如果可以对函数

f

进行矢量化，则可以通过沿矢量化结果的

x

轴计算总和，在numpy和Theano中更高效地执行计算（在时间上，但可能不是在空间上）

例如，如果您想要

output+=exp（k*x）

，那么您可以在numpy中实现这一点，而无需像下面这样的显式循环：

k = numpy.arange(k0, kN + 1)
result = numpy.exp(numpy.outer(x, k)).sum(axis=0)

---

如果

f

无法矢量化，或者由于其他原因需要循环，则NO可能提供也可能不提供更好的性能。你必须尝试一下才能找到答案。当需要显式循环时，只有当循环内部的计算涉及到非常大的张量运算时，Theano才有可能击败numpy。

我还没有深入到numpy，但我要说的是，我更喜欢numpy代码，而不是Theano的代码——更容易阅读——也许Python在阅读Theano代码时遇到了困难：-）这是一个明智的举动获取简化版的

np.arange（k0，kN+1）.sum（）

！我给出了那个具体的例子，但我试图解决一个更一般的问题，其中输出的增量可以是

k

和

x

的任何一般函数，比如

output+=x**k

，

output+=exp（k*x）

或

output+=（x-k）**2

。这就是为什么我没有选择在numpy中进行优化，而是决定使用theano，它使我能够进行这样的通用计算。@dzhelil我更新了我的答案，试图解决您更新后的问题。谢谢您的回答。请参阅我的编辑，以了解我为什么对使用theano感兴趣。有没有办法让theano代码运行得更快？

k = numpy.arange(k0, kN + 1)
result = numpy.exp(numpy.outer(x, k)).sum(axis=0)