メモ帳

python, juliaで機械学習をやっていく

ReformerをTrax (tensor2tensorの後継)で使ってみた

本記事の目的

  • Reformerの著者実装 (Trax) を使ってみる!!!
  • Traxで学習・推論するための一連の流れをまとめる

Reformerとは

  • Transformerの大幅な計算効率の向上・省メモリ化に成功

したモデルです。計算リソースは大きなボトルネックになるので価値が高いです。
また、Transformerは各種BERT系統のベースにもなっています。 BERT系統のモデルは、急速に巨大化しおり、もはや計算リソース的に一般には手がだせなくなってきているので押さえておいて損はないと思います。

既に素晴らしい記事があるので詳しくは割愛します。ぜひ、以下をご覧になって下さい。

参考:

Traxとは

github.com

  • Google Brain teamがメンテナンスしているtensor2tensorの後継
  • Reformerの著者実装が含まれている
  • TensorflowとJAXがベースの独立したMLフレームワーク。tf.kerasっぽく書ける
  • GPU, TPUの使用が容易
  • Document皆無
  • 一部のlayerだけTensorflowやpytorchにもってくることはできない(多分)

Traxを使ってみる

以下の様に順を追ってtraxでの開発の流れをまとめます。

  1. modelの定義
  2. loss, optimizer, learning rate schedulerの定義
  3. datasetの準備
  4. Trainerの定義
  5. 学習
  6. 推論パイプラインの作成

なお、本記事のコードはcolabで動かしました。

1. modelの定義

Transformerやresnetなど主流なモデルがいくつか実装されています。(コード)
特に、Reformer関係だと以下の2つが実装されています。

name 定義
ReformerLM language model (decoderのみ)
Reformer encoder-decoder model

Convolution, pooling, RNNなどの主要なlayerやnormalization, activationあたりも豊富に実装されているので、tf.kerasのように層を重ねていくだけでニューラルネットワークモデルを柔軟に実装できます。

ただし、ドキュメントがほとんど何もないですし、ググっても何も引っかからないのでカスタムlayerをゴリゴリに作って、込み入ったことをしようとすると苦労すると思います。

Traxのmodelの実装

traxは独自のkeras.layersの様なlayerが定義されており(コード)、それらを組み合わせてmodelを実装しています。簡単のために、多層パーセプトロンモデルの実装を紹介します。
tl.Serialkeras.models.Sequentialと読み替えて貰えばだいたいイメージはつかめると思います。 (kerasになれていれば)

from trax import layers as tl

def MLP(d_hidden=512,
        n_hidden_layers=2,
        Activation=tl.Relu,
        n_output_classes=10,mode=''):
    """A multi-layer feedforward (perceptron) network."""
    del mode

    def DensePlusActivation():
        return [tl.Dense(d_hidden), Activation()]

    return tl.Serial(
        tl.Flatten(),
        [DensePlusActivation() for _ in range(n_hidden_layers)],
        tl.Dense(n_output_classes),
        tl.LogSoftmax(),
    )

2. loss, optimizer, learning rate schedulerの定義

基本的なものは関数があります。以下や実装を参考にして使用するものを選択します。

loss - code

name 定義
L2Loss 自乗和誤差
CrossEntropyLoss 交差エントロピー誤差

optimizer - code

  • adafactor
  • adam
  • momentum
  • rms_prop
  • sgd
  • sm3

learning rate scheduling - code

実装を参照して下さい。

3. datasetの準備

datasetはtrax.supervised.Inputs()インスタンスを作る必要があります。 まず、以下のようなgeneratorを用意します。ただし、trax.supervised.Inputs()インスタンスはインターフェースを固定するためだけの存在なので、augmentationなどを行いたい場合は、別途tensorflow.dataなどのエコシステムを使っていく必要があります。
また、datasetのshapeを(batch_size, ...)となるようにあらかじめ変換してから渡す必要があります。
trax特有の処理は特にないので今回は、サンプルコードをそのまま使います。

import numpy as onp
import jax.numpy as np

def copy_task(batch_size, vocab_size, length):
    """This task is to copy a random string w, so the input is 0w0w."""
    while True:
        assert length % 2 == 0
        w_length = (length // 2) - 1
        w = onp.random.randint(low=1, high=vocab_size-1, size=(batch_size, w_length))
        zero = onp.zeros([batch_size, 1], onp.int32)
        loss_weights = onp.concatenate([onp.zeros((batch_size, w_length)),
                                    onp.ones((batch_size, w_length+2))], axis=1)
        x = onp.concatenate([zero, w, zero, w], axis=1)
        yield (x, x, loss_weights)  # =(Inputs, Targets, Weights)

このgeneratorを渡して以下の様に、trax.supervised.Inputs()インスタンスを生成します。

batch_size = 16
vocab_size = 32
sentence_length = 10
copy_inputs = trax.supervised.Inputs(lambda _: copy_task(batch_size, vocab_size, sentence_length))

1つのデータの中身は、以下の様になっています。 今回の例は、入力した文章のコピーを推論するタスクです。5つ目のtokenまでが原文であり、その後の5つのtokenがコピーした文章になります。weightsはinputの各tokenのlossへの重み付けを行います。つまり、ここでは4文字目までは推論結果がlossに影響を与えません。5文字目、つまり、原文の最後のtokenを入力し、コピー文章を全て生成する所までの誤差を最小にするような学習を行います。

data_stream = copy_inputs.train_stream(1)
inputs, targets, weights = next(data_stream)

# inputs[0] ->  [ 0, 5, 3, 14, 20, 0, 5, 3, 14, 20]
# targets[0] -> [ 0, 5, 3, 14, 20, 0, 5, 3, 14, 20]
# weights[0] -> [0., 0., 0., 0., 1., 1., 1., 1., 1., 1.]

4. Trainerの定義

次に、tl.Serial インスタンスやloss, optimizerの関数定義、datasetを1つのtrax.supervised.Trainerインスタンスにまとめます。

output_dir = os.path.expanduser('~/train_dir/') # モデルの保存先を指定
!rm -f ~/train_dir/model.pkl                    # 過去のモデルを削除

trainer = trax.supervised.Trainer(
    model=tiny_reformer_lm,                   # モデルの指定
    loss_fn=trax.layers.CrossEntropyLoss,     # lossの指定
    optimizer=trax.optimizers.Adam,           # optimizerの指定
    lr_schedule=trax.lr.MultifactorSchedule,  # lr scheduleの指定
    inputs=copy_inputs,                       # datasetの指定
    output_dir=output_dir,                    # モデルの保存先の指定
    has_weights=True)

ここで、モデルは以下の様にReformerLM (Reformerのdecoder部分のみ)を使用することにします。

def tiny_reformer_lm(mode):
    return trax.models.ReformerLM(
      d_model=32, d_ff=128, n_layers=2, vocab_size=32, mode=mode)

5. 学習

trainerを使って、以下の様に学習を行うことができます。

n_epochs  = 3
train_steps = 500
eval_steps = 2
# 1 epoch内で500 batchを学習に使い、2 batchでvalidationを行う。
# それを、計3 epoch行う。
for _ in range(n_epochs):
    trainer.train_epoch(train_steps, eval_steps)

以下のようなレポートが表示されます。

...
Step   1500: Ran 500 train steps in 1.67 secs
Step   1500: Evaluation
Step   1500: train                   accuracy |  0.83854169
Step   1500: train                       loss |  0.57018900
Step   1500: train         neg_log_perplexity |  0.57018900
Step   1500: train weights_per_batch_per_core |  96.00000000
Step   1500: eval                    accuracy |  0.83333337
Step   1500: eval                        loss |  0.58189404
Step   1500: eval          neg_log_perplexity |  0.58189404
Step   1500: eval  weights_per_batch_per_core |  96.00000000
Step   1500: Finished evaluation

余談

TransformerLMで全く同じ条件で学習させた所、1epochで約2.2 secs程度でした。 入力の文章が短いのであまり大きな差は確認できませんでしたが、確かにReformerの方が速度はでていそうです。

6. 推論パイプラインの作成

最後に、学習したモデルを使って推論パイプラインをつくります。

まず、predict modeでreformerのinstanceを生成し、学習した重みを.init_from_file()で読み込みます。ただし、途中でstateの初期値を変数にとっておきます。これは、predictの度にstateをcacheしていくので、1文章の処理毎に都度stateを初期化してあげる必要があるためです。

# Initialize model for inference.
model_infer = tiny_reformer_lm(mode='predict')

# Set up the initial state for sampling.
initial_state = model_infer.new_weights_and_state(
    trax.supervised.trainer_lib.ShapeDtype((1,1), dtype=np.int32))[1]

# load pretrained weight
model_infer.init_from_file(os.path.join(output_dir, "model.pkl"), weights_only=True)

次に、データを入力させて実際に推論するための関数を用意します。ここで注意点として、仕様上、推論時はtokenを1つずつ入力しなければいけません。つまり文章の長さだけ繰り返し推論させる必要があります。 今回、原文のコピーを推論するタスクなので、前半部分は原文を出力し、後半のみモデルの推論結果を出力することにします。

def prediction(prompt, length=2048):
    """Sample from the ReformerLM model"""
    # Token id 0 is the equivalent of a "start" token
    model_infer.state = initial_state             # stateの初期化
    cur_inputs = np.zeros((1, 1), dtype=np.int32) # 初期値=0の挿入
    prompt = np.asarray(prompt)
    all_samples = []

    for iteration in range(length):
        logits = model_infer(cur_inputs)
    
        if iteration < prompt.shape[0]:
            cur_samples = onp.array(prompt[iteration], dtype=int)
        else:
            logits = onp.array(logits)[0,0,:]
            probs = onp.exp(logits)
            cur_samples = onp.random.choice(probs.shape[-1], p=probs[:])
            cur_samples = onp.array(cur_samples, dtype=int)
        all_samples.append(cur_samples)
        cur_inputs = np.array(cur_samples[None,None])
    all_samples = onp.stack(all_samples, -1)

    return all_samples

出力例は次のようになります。prefixが入力なので、コピーできていることが分かります。

# Run inference
prefix = [0, 1, 2, 3, 2, 0] 
prediction(prefix, 10)
# -> array([0, 1, 2, 3, 2, 0, 1, 2, 3, 2])

以上がTrax使用の流れです。

まとめ

  • Reformerの著者実装 (Trax) を使ってみました!
  • Traxで学習・推論するための一連の流れをまとめました!

Reformerを試すというよりかは1つのMLフレームワークさわってみた、という内容になりました。Reformerの実装を見る限り、フレームワークそのものも使いやすそうでした。Githubに再現実装がいくつかありますが、基本的に著者実装以外はメンテナンスが頼りない印象があるので現状は、(reformer関連だけは) Traxをそのまま使っていこかなーと思っています。

別記事で、以前書いたtransformerを使った記事のReformer版を試したり、 Reformerではなくattention layerだけ使ったモデルをTraxで書いてみたりしよーかなと考えております。

如何せん情報が全然ないので人柱になるつもりで色々試していきたいと思います!!!

(Part 2) tensorflow 2 でhugging faceのtransformers公式のBERT日本語学習済みモデルを文書分類モデルにfine-tuningする

概要

以下の記事の続編になります。こちらの記事では、hugging faceのtransformersというライブラリを使用してBERTのfine-tuningを試しました。 transformersでの公開済みモデルを使用したfine-tuningの流れを紹介しているので、サポートされていない学習済みモデル(自分で学習させたものなど)を使って転移学習やfine-tuningをしたい場合は前回の記事を参照して頂いた方がいいかと思います。

tksmml.hatenablog.com

本記事では、以下を扱います。

  • 日本語サポートの拡充についてざっくりまとめる
  • 追加された学習済みモデルを使った、前回いまいちだった日本語文書分類モデルの精度の向上 → 飛躍的に精度上がりました!!!
続きを読む

自然言語処理のData Augmentation手法 (Easy Data Augmentation)

自然言語処理 Advent Calendar 2019の10日目です。昨日は、ktr_wtbさん形態素解析ツールインストール(MeCab,Juman++,Janome,GiNZA)と、plantarumさん自然言語処理を始める準備でした。

本記事では、以下の論文の概要をまとめます。(技術系の記事が多いのでちょっと気がひけますが)

EDA: Easy Data Augmentation Techniques for Boosting Performance on Text Classification Tasks [Jason W. Wei, Kai Zou; 2019]

以下、すべての画像はこの論文からの引用です。

TL;DR

  • 英語の5つの文書分類タスクで以下の処理により、Data Augmentationを行った。
    • (先行研究あり) 同義語で置換
    • (New!) randomな同義語の挿入
    • (New!) randomなワードの移動
    • (New!) randomなワードの削除
  • 半分程度のデータ量だけ学習させて同等な精度をだせた
  • convolutionとRNNの両方で効果が出た
  • 著者の実装はこちら
続きを読む

Pyjuliaによるpythonからjulialangへの型変換まとめ

データ処理の高速化のためにPyjuliaを使ってpythonからjuliaの関数を使っています。本記事ではpyjuliaによって、pythonのデータ型がどのようなjuliaのデータ型に変換されるか、その対応表をまとめます。

続きを読む

cosine similarityの重み付けを使った、パーソナライズできるrecommender systemのベースライン

アイテム情報とユーザー情報を組み合わせた、パーソナライズされた推薦を行う基本的なシステムを紹介します。重み付けしたcosine similarity (コサイン類似度)によるシンプルな手法です。いわゆるcontent-basedなrecommendになっています。 機械学習を使った推薦システムでは、metric learningやautoencoderなどで高尚な特徴量に変換し、類似度の大きさを指標としたものが派手な話題性の高い手法だと思います。しかし、これらの手法は直感的でないので、うまく学習できないと全く使えないものになると思います。そこで、保守的な、ベースラインとして使う想定の手法を紹介したいと思います。

続きを読む

pythonの関数の処理時間計測デコレータ

pythonで実行時間を測定するためにはtime ライブラリ、jupyter上では%time などの使用が有名だと思います。コードの中に余計な処理を書き込みたくない人向けに、decoratorを使って関数の定義部分で宣言するだけで、実行後に処理時間を出力するようにかえる方法を紹介します。 また、少し分かりにくいdecoratorを解説します。

続きを読む