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

对于您正在执行的操作，您可以简单地将 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

Answer 2

基于 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 方法（以及我的 Theano 版本的 arange sum 方法）以及我使用算术序列结果的标准和的 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，那么通过计算沿矢量化结果。

例如，如果你想要 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。

Theano 扫描用于数组上的快速计算

Theano scan for fast computations on an array

python

numpy

theano