Python functools 模块的 @lru_cache 装饰器介绍

functools.lru_cache 是 Python 标准库 functools 模块中的一个装饰器，用于实现简单的缓存机制。它通过缓存函数的返回值来提高函数的执行效率，特别是对于那些被多次调用且参数相同的函数。

LRU 缓存机制

• LRU 代表 Least Recently Used，即最近最少使用。

• LRU 缓存机制会保存最近使用的缓存项，并在缓存满时丢弃最久未使用的项。

lru_cache 装饰器的参数

• maxsize：指定缓存中可以存储的最大键值对数量，类型：int 或 None。

  • 如果 maxsize=None，则缓存大小无限制。

  • 如果指定了 maxsize，则缓存会限制为该大小。当缓存满时，最久未使用的缓存项会被丢弃（LRU 策略）。

• 默认值 128 表示缓存可以存储最多 128 个键值对。

• 每个键值对包括：

  • 键：函数的参数（包括位置参数和关键字参数）。

  • 值：函数的返回值。

• typed:

  • 类型：bool

  • 默认值：False

  • 描述：如果设置为 True，则会区分不同类型的参数。例如，f(3) 和 f(3.0) 会被视为不同的调用。

使用示例

示例 1：缓存斐波那契数列

import functools
import time

@functools.lru_cache(maxsize=None)
def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

# 记录开始时间
start_time = time.time()

print(fibonacci(30))  # 计算斐波那契数列的第 30 项

# 记录结束时间
end_time = time.time()

# 计算运行时间
run_time = end_time - start_time
print(f"加了@lru_cache 装饰器的fibonacci 运行时间: {run_time} 秒")


def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

print(fibonacci(30))  # 计算斐波那契数列的第 30 项

# 记录开始时间
start_time = time.time()

print(fibonacci(30))  # 计算斐波那契数列的第 30 项

# 记录结束时间
end_time = time.time()

# 计算运行时间
run_time = end_time - start_time
print(f"没有 @lru_cache 装饰器的fibonacci 运行时间: {run_time} 秒")

解释：

• 没有缓存时，计算斐波那契数列的第 30 项会非常慢，因为有大量的重复计算。

• 使用 lru_cache 后，函数的返回值被缓存，重复调用时直接从缓存中获取结果，大大提高了效率。

示例 2：缓存单位四元数

Python复制

import functools
import torch

@functools.lru_cache(maxsize=None)
def identity_quats(batch_dims, dtype=None, device=None, requires_grad=True):
    quat = torch.zeros(
        (*batch_dims, 4), 
        dtype=dtype, 
        device=device, 
        requires_grad=requires_grad
    )
    with torch.no_grad():
        quat[..., 0] = 1
    return quat

# 调用函数
quats = identity_quats((2, 3), dtype=torch.float32, device=torch.device('cpu'), requires_grad=True)
print(quats)

输出：

plaintext复制

tensor([[[1., 0., 0., 0.],
         [1., 0., 0., 0.],
         [1., 0., 0., 0.]],

        [[1., 0., 0., 0.],
         [1., 0., 0., 0.],
         [1., 0., 0., 0.]]], grad_fn=)

解释：

• 第一次调用 identity_quats 时，函数会计算并缓存结果。

• 后续调用时，如果参数相同，直接从缓存中返回结果，避免重复计算。

缓存内容和生命周期

缓存的内容

• 最终返回值：lru_cache 只缓存函数的最终返回值。它不会保存函数执行过程中的中间状态或中间值。

• 键值对：缓存是以键值对的形式存储的，其中键是函数的参数（包括位置参数和关键字参数），值是函数的返回值。

缓存的生命周期

• 内存中保存：缓存的值会一直保留在内存中，直到程序退出，或者缓存被手动清除。

• 手动清除缓存：可以使用 cache_clear 方法手动清除缓存。例如：

  fibonacci.cache_clear()

  这将清除 fibonacci 函数的所有缓存值。

总结

functools.lru_cache 是一个非常有用的装饰器，用于缓存函数的返回值，从而提高函数的执行效率。它特别适用于那些被多次调用且参数相同的函数。通过合理使用 lru_cache，可以显著提高程序的性能。