Attention Flowを追実験したい - アルゴリズムとかオーダーとか

最近は自然言語学習モデルのBERTを活用した実験を行っている。その中で、Attentionで何が表現されているのか？について興味が出てきた。
とあるツテから以下の論文を紹介してもらったので、この論文の中に出てきているサンプルコードを動かしたいというのが今回の記事の内容である。

https://arxiv.org/abs/2005.00928

なお、サンプルコードは動かせていないので、現時点までで把握した内容についての作業メモを記事として残す。

動かしたいサンプルコード

github.com

論文中に実験に用いたコードのgithubが紹介されていたので、そこにあるサンプルコードを動かす。今回の対象は↑のJyupiter Notebookファイル。
これをgoogle colabで動かそうとしている最中である。

大まかな問題点

HuggingfaceのTransformersを使っている。が、そのバージョンが古すぎるのが問題。論文作成時点でのバージョンは4.10.0だった。
現在の最新版（この記事執筆時点である2022/04/02)では4.18.0である。また、自分が今使っているバージョンは4.16.2なので、以降は基本的にtransformers==4.16.2で動かす方法である。

特に動いていない部分は、Tensorflowベースのモデルをpytorchベースに変換している部分であった。

修正箇所一覧

ここでは、次の作業のために、現時点で認識できている修正箇所を羅列する。

import部分

from transformers import *

ここの表記は最新版ではできない。

from  transformers import  (AutoTokenizer. AutoModel)

とかに書き換えたらよい。

tokenizerとmodelのマッピングとロード部分

MODELS = [(BertModel,       BertTokenizer,       'bert-base-uncased'),
          (OpenAIGPTModel,  OpenAIGPTTokenizer,  'openai-gpt'),
          (GPT2Model,       GPT2Tokenizer,       'gpt2'),
          (CTRLModel,       CTRLTokenizer,       'ctrl'),
          (TransfoXLModel,  TransfoXLTokenizer,  'transfo-xl-wt103'),
          (XLNetModel,      XLNetTokenizer,      'xlnet-base-cased'),
          (XLMModel,        XLMTokenizer,        'xlm-mlm-enfr-1024'),
          (DistilBertModel, DistilBertTokenizer, 'distilbert-base-uncased'),
          (RobertaModel,    RobertaTokenizer,    'roberta-base')]

# Each architecture is provided with several class for fine-tuning on down-stream tasks, e.g.
BERT_MODEL_CLASSES = [BertModel, BertForPreTraining, BertForMaskedLM, BertForNextSentencePrediction,
                      BertForSequenceClassification, BertForTokenClassification, BertForQuestionAnswering]

このマッピングはそもそも不要。Auto系が全部やってくれるので丸々コメントアウトでよい。

# Each architecture is provided with several class for fine-tuning on down-stream tasks, e.g.
BERT_MODEL_CLASSES = [BertModel, BertForPreTraining, BertForMaskedLM, BertForNextSentencePrediction,
                      BertForSequenceClassification, BertForTokenClassification, BertForQuestionAnswering]

# All the classes for an architecture can be initiated from pretrained weights for this architecture
# Note that additional weights added for fine-tuning are only initialized
# and need to be trained on the down-stream task
pretrained_weights = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(pretrained_weights)

model = BertModel.from_pretrained(pretrained_weights,
                                    output_hidden_states=True,
                                    output_attentions=True)

そして、上記箇所を以下に書き換える

pretrained_weights = 'bert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(pretrained_weights)

model = AutoModel.from_pretrained(pretrained_weights,
                                    output_hidden_states=True,
                                    output_attentions=True)

sst2のデータセットがロードできない問題

data = tfds.load('glue/sst2')

# Prepare dataset for GLUE as a tf.data.Dataset instance
train_dataset = glue_convert_examples_to_features(data['train'], tokenizer, max_length=128, task='sst2')
valid_dataset = glue_convert_examples_to_features(data['validation'], tokenizer, max_length=128, task='sst2')
train_dataset = train_dataset.shuffle(100).batch(32).repeat(10)
valid_dataset = valid_dataset.batch(64)

valid_iter = iter(valid_dataset)

このコードブロックが動かなかった。理由は、tensorflowが公開しているはずの、GLUE/SST-2のデータセットが403エラーでダウンロードできないため。

ただ、よくよくコードを読むと多分ここはAttentionを抽出するためにBERTに流し込むdatasetを取ってきているだけに見える。
なので、huggingface/datasetsのload_datasetで置き換えればよさそう。

TFTFBertForSequenceClassificationをpytorchベースに変換している部分

model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased')

# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule 
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
model.compile(optimizer=optimizer, loss=loss, metrics=[metric])

# Train and evaluate using tf.keras.Model.fit()
history = model.fit(train_dataset, epochs=10, steps_per_epoch=115,
                    validation_data=valid_dataset, validation_steps=7)

x,y = valid_iter.next()
print(x['input_ids'][0])

# Load the TensorFlow model in PyTorch for inspection
model.save_pretrained('save/')

pytorch_model = BertForSequenceClassification.from_pretrained('save/', from_tf=True, output_hidden_states=True,
                                    output_attentions=True)

この辺のコードはTensorflowベースのBERTモデルをpytorchベースに変換してるコードっぽい。
transformers 4.16.2であれば、最初からpytorchベースのBERTモデルをロードできるので多分このコードは不要。

今の段階でコードを整理できているのはここまで。

すでに修正済みにColabがあった。。。

と、この記事を書きつつなにげにgithubのREADMEを見ていたら上記の修正を実施したGoogle Colabのリンクがありました。。。。
colab.research.google.com