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

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

例子：斐波那契数列

下面这个计算斐波那契数列的函数 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]

更进一步：包装类

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

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)

运行测试，计算两次 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

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

再进一步：装饰器

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

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

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

最后的完善

之前的 Memory 类只适合包装参数为不可变对象的函数。原因是我们用到了字典作为存储介质，将参数作为字典的 key；而在 Python 中的 dict 只能把不可变对象作为 key ²，例如数字、字符串、元组（里面的元素也得是不可变对象）。所以提高代码通用性，我们只能牺牲运行速度，将函数参数序列化为字符串再作为 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]

使用第三方库 - joblib

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

首先定义缓存目录：

>>> cachedir = 'your_cache_location_directory'

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

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

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

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

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

>>> print(f(1))
Running f(1)
1
>>> print(f(1))
1

参考

1 http://code.activestate.com/recipes/52201/

2 https://docs.python.org/3/tutorial/datastructures.html#dictionaries

3 https://joblib.readthedocs.io/en/latest/memory.html#use-case

（本文完）