Theano 扫描在数组上进行快速计算

Question

我正在尝试使用 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，我将能够优化外循环，并获得比粗暴的 numpy 解决方案更快的解决方案。

Answer 1

以 Divakar 的回答为基础...

Theano 可以胜过 numpy 的情况非常具体。一般来说，只有当计算涉及对大张量的可向量化操作时，Theano 才会与 numpy 相比表现良好。

在这种情况下，该操作可以在 numpy 中非常有效地执行。通过使用sum of an arithmetic sequence 的标准结果，根本不需要使用循环。这里n = kN - k0 + 1 是要求和的项目数。

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

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

以下代码实现了原始的 numpy 和 Theano 方法，并将它们与 Divakar 的 numpy 方法（以及我的 arange sum 方法的 Theano 版本）以及我的使用算术序列结果的标准和的 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 计算的蹩脚台式电脑上，我得到以下时间（以秒为单位，越低越好）：

  numpy1 0.27894366
  numpy2 0.00011240
  numpy3 0.00008502
  numpy4 0.00006357
 theano1 0.99175695
 theano2 0.00040656
 theano4 0.00017563

在这种特殊情况下，在 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 不能向量化或者由于其他原因需要循环，那么 Theano 可能会也可能不会提供更好的性能。您必须尝试一下才能找到答案。当需要显式循环时，只有在循环内部发生的计算涉及非常大的张量运算时，Theano 才有可能击败 numpy。

Answer 2

对于您正在执行的操作，您可以简单地将 k0 到 kN 的所有元素相加得到一个标量，该标量必须用于缩放 x 以获得所需的输出。这样，您将拥有一个留在 NumPy 环境中并使用NumPy's strengths 的矢量化方法。 np.sum() 的实现看起来像这样 -

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

您也可以使用np.einsum 进行求和，这样在性能上可能会稍好一些，就像这样 -

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