Python 3类型提示用于性能优化

Question

PEP 484说“使用类型提示进行性能优化是留给读者的练习。”这告诉我，就像Common Lisp一样，当我发誓我知道自己在做什么时，类型声明可用于在性能密集型函数中预留类型调度。为了自己尝试这个，我用一个p系列计算了一个基准来计算pi。首先，我以天真的方式做，然后我尝试聪明并利用类型提示的性能：

import math
import time

def baselpi0(n):
    baselsum = 0;
    for i in range(1,n):
        baselsum += 1.0 / (i * i)
    return math.sqrt(6.0 * baselsum)

def baselpi1(n : int) -> float:
    n = float(n)
    baselsum  = 0.0
    i = 1.0
    while i < n:
        baselsum += 1.0 / (i * i)
        i += 1.0
    return math.sqrt(6.0 * baselsum)

start = time.time()
print(baselpi0(1000000000))
end = time.time()
print(end - start)
start = time.time()
print(baselpi1(1000000000))
end = time.time()
print(end - start)

我试图模仿的Common Lisp类比是：

(defun baselpi0 (n)
  (let ((baselsum 0.0d0))
    (loop for i from 1 to n do
      (setf baselsum (+ baselsum (/ 1.0 (* i i)))))
    (sqrt (* 6 baselsum))))

(defun baselpi1 (n)
  (let ((baselsum 0.0d0)
        (n (coerce n 'double-float)))
    (declare (type double-float baselsum n)
         (optimize (speed 3) (safety 0) (debug 0)))
    (loop for i from 1.0d0 to n do
          (setf baselsum (+ baselsum (/ 1.0d0 (* i i)))))
    (sqrt (* 6.0d0 baselsum))))

(time (princ (baselpi0 1000000000)))
(time (princ (baselpi1 1000000000)))
(exit)

在我的机器上，使用sbcl运行的lisp版本对于慢版本需要22秒，对于类型提示版本需要4秒，与C相同。对于朴素版本，CPython需要162秒，对于类型提示版本需要141秒。 Pypy在不到5秒的时间内运行非类型提示版本，但是库支持对我的项目来说还不够好。

有没有办法可以改进我的类型暗示版本，以获得更接近lisp或Pypy的性能？

Answer 1

速度差异不是由于类型提示。 Python目前，在可预见的未来，只会丢弃您提供的任何提示，并继续像往常一样动态执行。

这是因为在一种情况下，您在整个代码中使用浮动算法（这会导致更快的执行），而在另一种情况下则不会。

例证：将baselpi1更改为以下内容：

def baselpi1(n : int) -> float:
    n = float(n)
    baselsum  = 0
    i = 1
    while i < n:
        baselsum += 1.0 / (i * i)
        i += 1
    return math.sqrt(6.0 * baselsum)

现在来看看执行时间：

3.141591698659554
0.2511475086212158
3.141591698659554
0.4525010585784912

是的，它慢了。

Answer 2

如果你需要进行大量的数值计算，那么numpy通常会提供一个很好的选择。 Numpy使用较低级别的数据类型（例如固定宽度的整数 - python是无界的）。这为您提供了您感兴趣的类型提示。由于numpy旨在处理具有已知类型的数组中的大量数据，因此它可以在整个数组上有效地执行相同的操作。这也使numpy能够很好地处理SIMD指令的CPU（我不知道没有SIMD的现代CPU）。

我通常会重写你的功能：

import math
import numpy

def baselpi_numpy(n):
    i = numpy.arange(1, n) # array of 1..n
    baselsum = (1 / (i * i)).sum()
    return math.sqrt(6 * baselsum)

但是，对于大型n，你将没有足够的内存。您必须添加一些额外的代码来为您批量操作。那是：

def baselpi_numpy(n, batch_size=1 << 16):
    basel_sum = 0
    i = 1
    for i in range(1, n, batch_size):
        j = min(n, i + batch_size)
        basel_sum += baselsum_numpy(i, j)
    return math.sqrt(6 * basel_sum)

def baselsum_numpy(start, end):
    # equivalent -> sum(1 / (i * i) for i in range(start, end)) 
    i = numpy.arange(start, end, dtype=float)
    # this line and next are memory optimisations which double speed
    # equivalent to i = 1 / (i * i)
    i *= i 
    i = numpy.divide(1, i, out=i)
    basel_sum = i.sum()
    return basel_sum

我在笔记本电脑上的5.2秒内得到了结果。虽然我没有测试你使用的n的价值，但对于较低的n，numpy版本的速度要快20倍。