記憶（快取）函式返回值：Python 實現

NaN不等於NaN發表於2019-02-19
快取函式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 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]

使用第三方庫 – 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

（本文完）

相關文章