Python 实现：记忆（缓存）函数返回值

对于经常调用的函数，特别是递归函数或计算密集的函数，记忆（缓存）返回值可以显着提高性能。而在 Python 里，可以使用字典来完成。

有什么python相关报错解答自己不会的、或者源码资料/模块安装/女装大佬精通技巧 都可以来这里：(https://jq.qq.com/?_wv=1027&k=dwzSWSBK)或者+V：python10010问我

---

例子：斐波那契数列

下面这个计算斐波那契数列的函数 fib()具有记忆功能，对于计算过的函数参数可以直接给出答案，不必再计算：

fib_memo = {}
def fib(n):
    if n < 2: return 1
    if not n in fib_memo:
        fib_memo[n] = fib(n-1) + fib(n-2)
    return fib_memo[n]
1
2
3
4
5
6

更进一步：包装类

我们可以把这个操作包装成一个类 Memory，这个类的对象都具有记忆功能：

#python学习交流群：903971231###
class Memoize:
    """Memoize(fn) - 一个和 fn 返回值相同的可调用对象，但它具有额外的记忆功能。
       只适合参数为不可变对象的函数。
    """
    def __init__(self, fn):
        self.fn = fn
        self.memo = {}
    def __call__(self, *args):
        if not args in self.memo:
            self.memo[args] = self.fn(*args)
        return self.memo[args]

# 原始函数
def fib(n):
    print(f'Calculating fib({n})')
    if n < 2: return 1
    return fib(n-1) + fib(n-2)

# 使用方法
fib = Memoize(fib)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

运行测试，计算两次 fib(10)：

Calculating fib(10)
Calculating fib(9)
Calculating fib(8)
Calculating fib(7)
Calculating fib(6)
Calculating fib(5)
Calculating fib(4)
Calculating fib(3)
Calculating fib(2)
Calculating fib(1)
Calculating fib(0)
89
89
1
2
3
4
5
6
7
8
9
10
11
12
13

可以看到第二次直接输出 89，没有经过计算。

再进一步：装饰器

对装饰器熟悉的程序员应该已经想到，这个类可以被当成装饰器使用。在定义fib()的时候可以直接这样：

@Memoize
def fib(n):
    if n < 2: return 1
    return fib(n-1) + fib(n-2)
1
2
3
4

这和之前的代码等价，但是更简洁明了。

最后的完善

之前的 Memory 类只适合包装参数为不可变对象的函数。

原因是我们用到了字典作为存储介质，将参数作为字典的 key；

而在 Python 中的 dict 只能把不可变对象作为 key 2，例如数字、字符串、元组（里面的元素也得是不可变对象）。

所以提高代码通用性，我们只能牺牲运行速度，将函数参数序列化为字符串再作为 key 来存储，如下：

class Memoize:
    """Memoize(fn) - 一个和 fn 返回值相同的可调用对象，但它具有额外的记忆功能。
       此时适合所有函数。
    """
    def __init__(self, fn):
        self.fn = fn
        self.memo = {}
    def __call__(self, *args):
        import pickle
        s = pickle.dumps(args)
        if not s in self.memo:
            self.memo[s] = self.fn(*args)
        return self.memo[s]
1
2
3
4
5
6
7
8
9
10
11
12
13

使用第三方库 - joblib

除了这种手工制作的方法，有一个第三方库 joblib 能实现同样的功能，而且性能更好，适用性更广。因为上文中的方法是缓存在内存中的，每次都要比较传入的参数。对于很大的对象作为参数，如 numpy 数组，这种方法性能很差。而joblib.Memory模块提供了一个存储在硬盘上的 Memory 类，其用法如下：

首先定义缓存目录：

>>> cachedir = 'your_cache_location_directory'
1

以此缓存目录创建一个 memory 对象：

>>> from joblib import Memory
>>> memory = Memory(cachedir, verbose=0)
1
2

使用它和使用装饰器一样：

>>> @memory.cache
... def f(n):
...     print(f'Running f({n})')
...     return x
1
2
3
4

以同样的参数运行这个函数两次，只有第一次会真正计算：

 @memory.cache
... def f(n):
...     print(f'Running f({n})')
...     return x
1
2
3
4