スカイフォーマーを動かすだけ

transformer

セルフアテンションの計算に Nyström 近似を利用したスカイフォーマーのコードが以下に公開されています。現時点で最新のリビジョンをみます。

GitHub - pkuzengqi/Skyformer at cfe8c8cb48a151fd150ff4a87fdb24b288356869

モデルのソースは以下です。Long Range Arena タスクが想定されており、ModelForSC クラスがモデルのクラスですが、2つのドキュメントの類似度を求めるタスクについては ModelForSCDual クラスを使用します。

Skyformer/model_LRA.py at cfe8c8cb48a151fd150ff4a87fdb24b288356869 · pkuzengqi/Skyformer · GitHub

リポジトリを clone した ./src/ 下で適当に以下を実行するとインスタンス化できます（ src/requirements.txt のパッケージが充足していれば）。

from models.model_LRA import ModelForSC, ModelForSCDual
from config import Config

model_config = Config["lra-text"]["model"]
model_config["mixed_precision"] = True
model_config["attn_type"] = "softmax"
model = ModelForSC(model_config)
print(model)

ModelForSC(
  (model): Model(
    (embeddings): Embeddings(
      (word_embeddings): Embedding(512, 64)
      (position_embeddings): Embedding(4000, 64)
      (dropout): Dropout(p=0.1, inplace=False)
    )
    (transformer_0): TransformerLayer(
      (norm1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
      (mha): Attention(
        (W_q): Linear(in_features=64, out_features=64, bias=True)
        (W_k): Linear(in_features=64, out_features=64, bias=True)
        (W_v): Linear(in_features=64, out_features=64, bias=True)
        (attn): SoftmaxAttention(
          (drop_attn): Dropout(p=0.1, inplace=False)
        )
# 以下省略

なのでこれに Long Range Arena のデータを渡せばよいです。データの取得及びプレ処理方法は README に記述されていますが、README には親切にもプレ処理済みのデータへのリンクもあるのでそれを拝借することにします。適当に lra-text.dev.pickle をダウンロードします（ご自身の責任で）。中身を確認すると 25000 の文章が入っており、1つの文章は 4096 文字（おそらく ASCII コード）からなっており、何か正解ラベルが付いていることがわかります。

import pickle
with open('./Downloads/lra-text.dev.pickle', 'rb') as ifile:
    x = pickle.load(ifile)

print(len(x))
print(x[0]['input_ids_0'])
print(len(x[0]['input_ids_0']))
print(x[0]['label'])

25000
[ 85 105 102 ...   0   0   0]
4096
1

でここまでやっておいてなんですがモデルには文章とみせかけて適当な整数列を渡せば動くと思います。以下でモデルが値を返却してくれます。

import torch
x = torch.tensor([[0, 1, 2, 3, 4]])
label = torch.tensor([0])
y = model(x, None, label)
print(y)

{'loss': tensor([0.8154], grad_fn=), 'accu': tensor([0.])}

なお AttributeError: module 'torch.cuda.amp' has no attribute 'autocast' と怒られので怒られなくなるまで with torch.cuda.amp.autocast(enabled = False): をコメントアウトしています。後日 GPU 機で試します。