機械学習・自然言語処理の勉強メモ

学んだことのメモやまとめ

pytorch-transformersを触ってみる①

今更ながら、pytorch-transformersを触ってみます。
このライブラリはドキュメントが充実していて、とても親切です。
なので、今回はドキュメントに基づいて触ってみただけの備忘録です。
以下、有名どころのBERTで試してます。

詳しいことはここなどを参照してください。
huggingface.co

はじめに

以下で、入手できます。簡単です。

pip install pytorch-transformers

インストールしたら、以下でimportします。

import torch
from pytorch_transformers import BertTokenizer, BertModel

pytorch-transformersの基本は以下の3つのクラスで構成されます。

model classes モデル本体
configuration classes モデルのパラメータを設定するクラス(デフォルト設定の場合は不要)
tokenizer classes 入力文のtokenizeや辞書に関連するクラス

model classes

各モデル毎にfine tuningのための以下の7つのクラスが準備されています。
例えば、BERTの場合は以下のクラスがあります。

  • BertModel
  • BertForPreTraining
  • BertForMaskedLM
  • BertForNextSentencePrediction
  • BertForSequenceClassification
  • BertForMultipleChoice
  • BertForTokenClassification
  • BertForQuestionAnswering

理解を深めるため、クラスのコードを追っかけてみます。
modeling_bert.pyを見ると、全てBertPreTrainedModelの派生クラスということが分かります。
なので、基本的な振る舞いはほとんど同じです。
更に、BertPreTrainedModelはPreTrainedModelの派生クラスということが分かります。
PreTrainedModelが定義されているmodeling_utils.pyを見ると、
torch.nn.Modulesを継承しているということが分かります。
torch.nn.Modulesは、モデルを構築する時に継承するあれです。

では、BertForMaskedLMのアーキテクチャを見てみます。

from pytorch_transformers import BertForMaskedLM
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
BertForMaskedLM(
  (bert): BertModel(
    (embeddings): BertEmbeddings(
      (word_embeddings): Embedding(30522, 768, padding_idx=0)
      (position_embeddings): Embedding(512, 768)
      (token_type_embeddings): Embedding(2, 768)
      (LayerNorm): BertLayerNorm()
      (dropout): Dropout(p=0.1)
    )
    (encoder): BertEncoder(
      (layer): ModuleList(
        (0): BertLayer(
        .
        .
        .
  (cls): BertOnlyMLMHead(
    (predictions): BertLMPredictionHead(
      (transform): BertPredictionHeadTransform(
        (dense): Linear(in_features=768, out_features=768, bias=True)
        (LayerNorm): BertLayerNorm()
      )
      (decoder): Linear(in_features=768, out_features=30522, bias=False)
    )
  )
)

最後のdecoderの部分に着目すると、

(decoder): Linear(in_features=768, out_features=30522, bias=False)

30522クラスへの線形変換を行っていることが分かります。
30522は単語の数なので、単語を出力するためのモデルとなっていることが分かります。

同様にBertForSequenceClassificationの出力層を見てみます。

from pytorch_transformers import BertForSequenceClassification
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
(classifier): Linear(in_features=768, out_features=2, bias=True)

2クラスの線形変換であり、分類問題ということが分かります。
このように、タスクに応じて適切なモデルが用意されています。

では、BertModelとは何なんでしょうか。出力層を見てみます。

model = BertModel.from_pretrained('bert-base-uncased')
  (pooler): BertPooler(
    (dense): Linear(in_features=768, out_features=768, bias=True)
    (activation): Tanh()
  )

良く分かりませんが、説明では

use BertModel to encode our inputs in hidden-states

と書いていますので、この出力を別のモデルに連結する場合に使用するのでしょうか。
(用途が良く分かっていません。。)

Pretrained models

上ではモデルをロードする際に以下のように書いていました。

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

bert-base-uncased とは何なんでしょうか。

これは学習済のパラメータをロードするための記述です。
例えば、bert-base-uncasedと指定した場合は、12-layer, 768-hidden, 12-headsで小文字英文での事前学習済のパラメータがロードされます。
*論文中のBERT(BASE)に相当します。

以下のように自身で学習したパラメータのロードも可能です。

model = BertModel.from_pretrained('./test/saved_model/')

from_pretrainedについても、少しだけ追っかけてみました。
先ほど、BertModelはPreTrainedModelを継承していると書きました。
このfrom_pretrainedもPreTrainedModelのclassmethodとして定義されています。
(modeling_utils.pyの375行目あたり)
ロードするモデルについては、以下のように記述されています。

# Load model
if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
    archive_file = cls.pretrained_model_archive_map[pretrained_model_name_or_path]
elif os.path.isdir(pretrained_model_name_or_path):
    if from_tf:
        # Directly load from a TensorFlow checkpoint
        archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
    else:
        archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)

cls.pretrained_model_archive_mapは以下のような辞書ファイルです。

BERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin", 
    # ...
}

なので、引数をbert-base-uncasedとした場合は、Amazon S3 のストレージ上のmodelファイルを見に行くのですね。
(ローカルダウンロード後はそちらを見に行きます。)
model_archive_mapに存在しない文字列の場合は、ローカルのmodelファイルを見に行きます。
その後、指定したmodelファイルを以下のようにロードします。

if state_dict is None and not from_tf:
    state_dict = torch.load(resolved_archive_file, map_location='cpu')

このstate_dictもオプショナルな引数のようで、指定して与えることも可能です。

configuration classes

BertConfigなどのConfigクラスは、PretrainedConfigクラスを継承しています。
PretrainedConfigクラスはclassmethodとして以下のような関数があります。

  • from_dict
  • from_json_file
  • from_pretrained

from_dict, from_json_fileなどはハイパーパラメータを与える形式の違いということが想定できます。
ややこしいのが、from_pretrainedです。Modelクラスにもfrom_pretrainedとの違いについて少し掘り下げてみます。
BertConfigも同じようにPretrainedConfigのclassmethodとして定義されています。
(modeling_utils.pyの109行目あたり)
ここも、BertModelとやっていることはほとんど同じです。

config = BertConfig.from_pretrained('bert-base-uncased') 

のように記述された場合、以下の辞書ファイルに従ってAmazon S3 のストレージ上のjsonファイルを見に行きます。

BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
    # ...

ちなみに、こっちはJSON形式なので、ブラウザから参照可能です。
(bert-base-uncasedは以下のようなJSONファイルです。)

{
  "attention_probs_dropout_prob": 0.1,
  "hidden_act": "gelu",
  "hidden_dropout_prob": 0.1,
  "hidden_size": 768,
  "initializer_range": 0.02,
  "intermediate_size": 3072,
  "max_position_embeddings": 512,
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "type_vocab_size": 2,
  "vocab_size": 30522
}

少し、込み入ったので内容を整理すると、

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

のようにしてmodelクラスのfrom_pretrainedが呼ばれると、引数にconfigがない場合は
1. from_pretrainedの文字列でconfig.from_pretrainedを呼び出し
2. config.from_pretrainedの結果をconfigにセット
という手順でハイパーパラメータの値を得ます。
そのconfigの値で、モデルをインスタンス化します。
その後、読み込んでおいたmodelファイルを_load_from_state_dict()でセットします。

tokenizer classes

tokenizer classesにもfrom_pretrained()を通じてインスタンス化可能です。
tokenizer classのfrom_pretrainedも同様に

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

のように記述された場合、Amazon S3 のストレージ上のtxtファイルを見に行きます。

このtxtファイルはtokenが行単位でズラーっと記述されたものです。

.
skin
foreign
opening
governor
okay
.
引数について

tokenizerクラスには以下のような引数があります。

  • do_lower_case .. Whether to lower case the input. Only has an effect when do_wordpiece_only=False
  • do_basic_tokenize .. Whether to do basic tokenization before wordpiece.

※日本語の場合、do_basic_tokenize =False としおかないと文字単位に分割されてしまいます。

実験

最後に、いくつかのモデルを触ってみます。

bertForMaskedLM

ある文(文章)の一部をMASKして、それを予測する言語モデルです。
以下の文で試してみます。

I like to play football with my friends

ここで、footballをMASKして予測できるか試してみます。

まず、BERT用に先頭[CLS]、末尾に[SEP]を付けます。

text = "[CLS] I like to play football with my friends [SEP] ."

このtextをtokenizerに与えます。

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
tokenized_text = tokenizer.tokenize(text)
print(tokenized_text)
# ['[CLS]', 'i', 'like', 'to', 'play', 'football', 'with', 'my', 'friends', '[SEP]', '.']

次に、footballの箇所をMASKします。

masked_index = 5
tokenized_text[masked_index] = '[MASK]'
print(tokenized_text)
# ['[CLS]', 'i', 'like', 'to', 'play', '[MASK]', 'with', 'my', 'friends', '[SEP]', '.']

最後に単語を対応するidに変換したら入力系列の出来上がりです。

indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
tokens_tensor = torch.tensor([indexed_tokens])

これをモデルに与えます。
このあたりは、通常のニューラルネットと同じ感じでいけます。

model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()

tokens_tensor = tokens_tensor.to('cuda')
segments_tensors = segments_tensors.to('cuda')
model.to('cuda')

with torch.no_grad():
    outputs = model(tokens_tensor)
    predictions = outputs[0]
_, predicted_indexes = torch.topk(predictions[0, masked_index], k=5)
predicted_tokens = tokenizer.convert_ids_to_tokens(predicted_indexes.tolist())
print(predicted_tokens)
# ['golf', 'games', 'football', 'chess', 'tennis']

footballが3番目に来てます。良い感じです。
今回の文では、どれも正解と言えると思います。

BertForNextSentencePrediction

2つの文が連続しているかどうかを判定します。

まずは、ドキュメント通りに試してみます。

### Classify next sentence using ``bertForNextSentencePrediction``
# Going back to our initial input
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = tokenizer.tokenize(text)
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
model.eval()

tokens_tensor = tokens_tensor.to('cuda')
segments_tensors = segments_tensors.to('cuda')
model.to('cuda')

# Predict the next sentence classification logits
with torch.no_grad():
    next_sent_classif_logits = model(tokens_tensor, segments_tensors)

print(torch.softmax(next_sent_classif_logits[0], dim=1))
# tensor([[1.0000e+00, 2.8665e-06]], device='cuda:0')

0が連続していると考えると、正解しています。

次に、連続していない文章を与えてみます。

text = "[CLS] Japan has many traditional areas [SEP] Mike goes to library to study [SEP]"

予測結果を見てみます。

torch.softmax(next_sent_classif_logits[0], dim=1)
# tensor([[1.3438e-05, 9.9999e-01]], device='cuda:0')

さっきとは逆に1の方の確率が高くなっています。これも正解です。

おわりに

pytorch-transformersの簡単なところを触ってみました。
次回は、日本語の学習済モデルを使って同様のことを試してみたいと思います。

*コードはの行数などは記事を書いた時点のものなので、今後変更されるかもしれません。