對於經常調用的函數,特別是遞歸函數或計算密集的函數,記憶(緩存)返回值可以顯著提高性能。 ...
對於經常調用的函數,特別是遞歸函數或計算密集的函數,記憶(緩存)返回值可以顯著提高性能。而在 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
(本文完)