使用Numba JIT优化Python脚本的性能

Question

使用Numba JIT优化Python脚本的性能

3

我正在运行一个Python模拟程序，用于预测一种加权和普通骰子。我想使用numba来帮助加快我的脚本速度，但是我收到了一个错误：

<timed exec>:6: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "roll" failed type inference due to: Untyped global name 'sum': cannot determine Numba type of <class 'builtin_function_or_method'>

File "<timed exec>", line 9:
<source missing, REPL/exec in use?>

这是我的原始代码：是否有其他类型的numba表达式可以使用？目前我正在测试使用2500次投掷的输入；希望将其缩短到4秒钟（目前为8.5秒）。

%%time
from numba import jit
import random
import matplotlib.pyplot as plt
import numpy

@jit
def roll(sides, bias_list):
    assert len(bias_list) == sides, "Enter correct number of dice sides"
    number = random.uniform(0, sum(bias_list))
    current = 0
    for i, bias in enumerate(bias_list):
        current += bias
        if number <= current:
            return i + 1

no_of_rolls = 2500
weighted_die = {}
normal_die = {}
#weighted die

for i in range(no_of_rolls):
        weighted_die[i+1]=roll(6,(0.15, 0.15, 0.15, 0.15, 0.15, 0.25))
#regular die  
for i in range(no_of_rolls):
        normal_die[i+1]=roll(6,(0.167, 0.167, 0.167, 0.167, 0.167, 0.165))

plt.bar(*zip(*weighted_die.items()))
plt.show()
plt.bar(*zip(*normal_die.items()))
plt.show()

- Lee Whieldon

如果你只是想要一个更快的模拟，你可以使用random.choices来生成加权骰子和正常骰子的列表。 - DarrylG

你能给我一个例子吗？请随意使用我的代码并进行修改。 - Lee Whieldon

2

@LeeWhieldon提供了一个代码示例。但是，我不明白为什么你说这段代码需要8.5秒才能运行。我得到的是毫秒级别的时间（不包括绘图）。 - DarrylG

@DarrylG，谢谢！这确实提高了性能。非常感谢！ - Lee Whieldon

2个回答

0

你可以使用guvectorize来加速它

%%time
from numba import guvectorize
import matplotlib.pyplot as plt
import numpy as np
import random

sides = 6
bias_list = (0.15, 0.15, 0.15, 0.15, 0.15, 0.25)

@guvectorize(["f8[:,:], uint8[:]"], "(n, k) -> (n)", nopython=True)
def roll(biases, side):
    for i in range(biases.shape[0]):
        number = random.uniform(0, np.sum(biases[i,:]))
        current = 0
        for j, bias in enumerate(biases[i,:]):
            current += bias
            if number <= current:
                side[i] = j + 1
                break

no_of_rolls = 2500
biases = np.zeros((no_of_rolls,len(bias_list)))

biases[:,] = np.array(bias_list)

normal_die = roll(biases)

print(normal_die)

这段代码在我的电脑上只需要 ~200 毫秒，而你的代码大约需要 6 秒。

- maciek97x

我运行了你建议的代码，但在我的机器上似乎没有提高速度（实际上比之前慢了1秒）。 - Lee Whieldon

1

好的，刚刚意识到绘图需要很长时间。 - maciek97x

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- DarrylG · Accepted Answer

使用随机选择

重构后的代码

import random
import matplotlib.pyplot as plt

no_of_rolls = 2500

# weights
normal_weights = (0.167, 0.167, 0.167, 0.167, 0.167, 0.165)
bias_weights = (0.15, 0.15, 0.15, 0.15, 0.15, 0.25)

# Replaced roll function with random.choices 
# Reference: https://www.w3schools.com/python/ref_random_choices.asp
bias_rolls = random.choices(range(1, 7), weights = bias_weights, k = no_of_rolls)
normal_rolls = random.choices(range(1, 7), weights = normal_weights, k = no_of_rolls)

# Create dictionaries with same structure as posted code
weighted_die = dict(zip(range(no_of_rolls), bias_rolls))
normal_die = dict(zip(range(no_of_rolls), normal_rolls))

# Use posted plotting calls
plt.bar(*zip(*weighted_die.items()))
plt.show()
plt.bar(*zip(*normal_die.items()))
plt.show()

性能

*Not including plotting.*
Original code: ~6 ms
Revised code:  ~2 ms
(3x improvement, but not sure why the post mentions 8 seconds to run)