中間アクティベーション用キャッシュ

推論中、モデルはトークンごとに出力します。このシンプルなキャッシュを使ってキーと値のアテンションレイヤーを格納するので、以前のトークンでそれらを再計算する必要がありません

15from typing import Any

キャッシュ

これにより、キーと値のキャッシュが維持され、値のプッシュとポップが同じ順序でキューに入れられます。キューは複数のアテンションレイヤーがあるので便利です

18class Cache:
26    def __init__(self):
27        self._cache = {}

キャッシュをクリア

29    def clear_all(self):
33        self._cache = {}

値をキューにプッシュ

  • name はキューの名前です
  • value プッシュする値です
35    def push(self, name: str, value: Any):

存在しない場合は空のキューを作成

44        if name not in self._cache:
45            self._cache[name] = []

キューにプッシュ

48        self._cache[name].append(value)

キューのサイズを返す

  • name はキューの名前です
  • キューが存在する場合はキューのサイズを返します。それ以外の場合は None

50    def q_size(self, name):
58        if name not in self._cache:
59            return None
60
61        if type(self._cache[name]) != list:
62            return None
63
64        return len(self._cache[name])

キューからポップする

  • name はキューの名前です
  • 値を返す

66    def pop(self, name: str):
73        return self._cache[name].pop(0)

値をキャッシュする

  • key キャッシュされる値の名前
  • value は値です
75    def set(self, key: str, value: Any):
82        self._cache[key] = value

キャッシュから値を取得

  • key キャッシュ時に使用される名前です
  • default キャッシュが空の場合のデフォルト値です
  • キャッシュされた値を返します

84    def get(self, key: str, default: Any = None):
92        return self._cache.get(key, default)

キャッシュ値をクリアする

  • key キャッシュ時に使用される名前です
94    def clear(self, key: str):
100        del self._cache[key]

キャッシュ用シングルトン

104_INSTANCE = None

キャッシュインスタンスを取得

    キャッシュインスタンスを返します

107def get_cache() -> Cache:
113    global _INSTANCE
114
115    if _INSTANCE is None:
116        _INSTANCE = Cache()
117
118    return _INSTANCE