stMind

about Tech, Computer vision and Machine learning

GPTを自作して学習済みパラメータでテキスト生成

2024年の最初のエントリーはGPTです。 GPTモデルを自作して、OpenAIが公開している学習済みのパラメータをロード、テキスト生成までの一連の処理を実行します。

モデル

正確にはGPT2のTransformerブロックを自作します。 アーキテクチャの大部分はGPTと同じですが、以下の変更(pre-norm)が行われています。

  • LayerNormはAttentionとMLPの前で適用
  • 追加のLayerNormをTransformerブロックの後で適用

Transformerブロックを除くText & Position埋め込みとNext Token生成は、 picoGPTのコードを利用します(解説ブログは GPT in 60 Lines of NumPy | Jay Mody)。

また、以下で紹介するコードはTensorflowを用いて実装しています(picoGPTの諸々のコードがTensorflowを利用していて、そのまま使いたかったため)。

Transformerブロック

GPT2の論文ではアーキテクチャの図がないので、下記はGPTのアーキテクチャ図ですが、上で書いたように、LayerNormはAttentionとMLPの前で適用します。

これをTransformerDecoderBlockクラスとして用意します(推論だけ行うのでDropoutは不要ですが)。

class TransformerDecoderBlock(tf.keras.Model):
    def __init__(self, h_dim, n_heads, drop_p):
        super().__init__()

        self.attn = MaskedMultiSelfAttention(h_dim, n_heads, drop_p)
        self.mlp = tf.keras.Sequential(
            [
                tf.keras.layers.Dense(units=4 * h_dim, activation="gelu"),
                tf.keras.layers.Dense(units=h_dim),
                tf.keras.layers.Dropout(rate=drop_p),
            ]
        )
        self.ln1 = tf.keras.layers.LayerNormalization()
        self.ln2 = tf.keras.layers.LayerNormalization()

    def call(self, x):
        x = self.attn(self.ln1(x)) + x
        x = self.mlp(self.ln2(x)) + x
        return x

Masked Multi Self Attention

Attentionの計算は複雑なところはなく、以前に作った ゼロから作るVision Transformer (JAX/Flax) - stMindとも同じです。

  • 入力トークン列から、クエリq、キーk、バリューvを作成
  • 複数ヘッド毎のアテンション行列とアテンションの計算
  • ヘッド毎のアテンションを集約

ただし、このままだと現在のトークンが未来のトークンも参照してしまうことになるので、アテンション行列において現在のトークンと未来のトークンの関係性はなし(0)にする必要があります。未来のトークンは、アテンション行列の各行で列方向に並ぶので、下三角行列を作成してアテンション行列をマスク。この時、マスクした後でソフトマックスを計算すると正規化されなくなるので、未来のトークンは非常に小さい値にしておいて、そのあとでソフトマックスを適用します。

class MaskedMultiSelfAttention(tf.keras.layers.Layer):
    def __init__(self, h_dim, n_heads, drop_p):
        super(MaskedMultiSelfAttention, self).__init__()
        self.n_heads = n_heads

        self.c_attn = tf.keras.layers.Dense(3 * h_dim)

        self.c_proj = tf.keras.layers.Dense(h_dim)

        self.attn_drop = tf.keras.layers.Dropout(drop_p)
        self.proj_drop = tf.keras.layers.Dropout(drop_p)

    def call(self, x):
        B, T, C = x.shape
        N, D = self.n_heads, C // self.n_heads

        # Create lower triangle mask
        mask = tf.linalg.band_part(tf.ones((T, T)), -1, 0)
        mask = tf.reshape(mask, (1, 1, T, T))

        qkv = self.c_attn(x)
        q, k, v = tf.split(qkv, 3, axis=-1)
        q = tf.reshape(q, (B, T, N, D))
        k = tf.reshape(k, (B, T, N, D))
        v = tf.reshape(v, (B, T, N, D))

        q = tf.transpose(q, perm=[0, 2, 1, 3])
        k = tf.transpose(k, perm=[0, 2, 1, 3])
        v = tf.transpose(v, perm=[0, 2, 1, 3])

        weights = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(
            tf.cast(D, dtype=tf.float32)
        )

        # Apply mask
        weights += (1 - mask) * -1e9

        normalized_weights = tf.nn.softmax(weights, axis=-1)
        attention = self.attn_drop(tf.matmul(normalized_weights, v))
        attention = tf.transpose(attention, perm=[0, 2, 1, 3])
        attention = tf.reshape(attention, (B, T, C))

        out = self.proj_drop(self.c_proj(attention))
        return out

GPT2モデル全体

Transformerブロックができたので、Text & Position埋め込みと追加LayerNormを含めたGPT2全体を作ります。 GPT2全体コードは長くなるので、callメソッドだけ抜き出すと下のようになります。 埋め込みベクトルは、次で説明するパラメータを使って、input_idsに対する埋め込みを取り出して生成します。

def call(self, input_ids):
    # Text and Position Embedding
    input_ids = tf.cast(input_ids, tf.int32)
    x = tf.gather(self.wte, input_ids) + tf.gather(
        self.wpe, range(input_ids.shape[1])
    )
    # Transformer Block (Decoder only)
    for block in self.blocks:
        x = block(x)
    # Additional LayerNorm
    x = self.layer_norm(x)
    # Linear
    return tf.matmul(x, self.params["wte"].T)

トークン生成テスト

OpenAIの学習済みパラメータを使用して、トークンを生成してみます。

学習済みパラメータ

パラメータには、入力と位置埋め込み、Transformerの各ブロックのパラメータがあり、picoGPTで辞書形式に変換されているものを使用します。

  • blocks : Transformerブロックのパラメータ
  • ln_f : 追加のLayerNormのパラメータ
  • wpe : 位置埋め込みベクトル
  • wte : トークンの埋め込みベクトル

また、blocksは124Mのモデルの場合は12個の要素があり、それぞれが以下の項目を含んでいます。(768は次元数)

  • attn : アテンションブロック(以下のbはバイアス項、wは重み)
    • c_attn : {"b": [2304], "w": [768, 2304]}
    • c_proj : {"b": [768], "w": [768, 768]}
  • ln1 : {"b": [768], "g": [768]}, Attentionの前に適用するLayerNorm。bはbeta、gはgamma
  • ln2 : {"b": [768], "g": [768]}, MLPの前に適用するLayerNorm
  • mlp : MLPブロック
    • c_fc : {"b": [3072], "w": [768, 3072]}
    • c_proj : {"b": [768], "w": [3072, 768]}

Tensorflowにおけるパラメータの割り当て

tf.keras.layers.Layerのset_weigthsを使います。この関数は、numpy の配列からパラメータ値を設定します。 例えば、c_attnの場合だと、これはDense層なのでwとbの順番でset_weightsに指定します。

block.layers[0].c_attn.set_weights(
    [
        self.params["blocks"][layer_idx]["attn"]["c_attn"]["w"],
        self.params["blocks"][layer_idx]["attn"]["c_attn"]["b"],
    ]
)

ln_fはLayerNormなので、gammaとbetaでset_weightsに指定します。

self.layer_norm.set_weights(
    [self.params["ln_f"]["g"], self.params["ln_f"]["b"]]
)

生成結果

GPT2モデルの作成、重みパラメータの設定が出来たので、 picoGPTと同じプロンプトで実験してみます。

$ python tf/gpt_tf.py --prompt "Alan Turing theorized that computers would one day become" --n_tokens_to_generate 8
...
Input text:
 Alan Turing theorized that computers would one day become
Generated:
  the most powerful machines on the planet.

同じ結果が生成されました。生成には、M1 Macで2秒弱くらいかかりました。

別のプロンプトも試してみます。

python tf/gpt_tf.py --prompt "Imagination is more important" --n_tokens_to_generate 6
...
Input text:
 Imagination is more important
Generated:
  than any other skill.

文章としては問題ないものが生成されたように思います。

まとめ

以前のViTと今回のGPTでAttentionの自作をしましたが、処理自体はそれほど複雑ではないので、 Transformerブロックを実装するのは、慣れれば難しくはないように感じました。また、モデルを実装して学習するのはHW制約などもあって大変なことが多いですが、推論であれば公開されているパラメータを使うことで、比較的試してみやすいのではと思います。コアとなるTransformer、Attentionを自作することで、Transformer系の論文の数式やコードの読解力が上がったように感じるので、興味のある方は自作にトライしてみることをオススメします。

参考文献とコード

GPT in 60 Lines of NumPy | Jay Mody

GitHub - satojkovic/gpt-tf-pytorch-jax: GPT from scratch (tensorflow / pytorch / jax)

Pose Estimationとローパスフィルタ(One Euro Filter)の実験

フレーム単位の姿勢推定を動画に適用した場合に発生するジッターを解決するローパスフィルタを実験。

使用したのは、ここで紹介されていたコード。

towardsdatascience.com

アクション認識のデータセットUCF-101の動画を利用。実行したのはM1 Mac

フィルタなし

フィルタあり

処理時間としてはほとんど変化ないが、滑らかな追跡となっている。(代わりに追跡の遅れが気になる)

Lambda Labs GPU CloudでJAX/Flax

MacのMetalを使って、手持ちのM1 MacにもJAX/Flaxの実行環境を作ることは出来るのですが、 実際に学習をしようとしてもエラーで詰まってしまうことが多く、JAX/Flaxを実行できる環境を探していました。

Colabを使っても良いのですが、学習を実行するだけでなくて、JAXのビルド自体も試してみたいと思ったので、Lambda Labs GPU Cloudで実行してみることにしました。

実行したのは、FlaxのチュートリアルにあるMNISTの画像分類モデル(CNN)の学習です。

CPU実行するだけであれば何もする必要はなかったのですが、GPUを使う場合には少しだけ苦労しました。

GPU実行時のエラーとTF_FORCE_GPU_ALLOW_GROWTH

最初にGPUで実行したとき、次のようなエラーが出ました。

...
2023-09-18 05:06:37.640924: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:439] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
...
jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.

今回のチュートリアルでは、tensorflow_datasetsを使ってMNISTのデータセットをロードしていたけれど、tensorflowは何も設定しないと、初期化時にGPUメモリの大部分を割り当ててしまうので、TF_FORCE_GPU_ALLOW_GROWTH=trueとしておかないとエラーになる可能性がある様子。

stackoverflow.com

github.com

github.com

最後のIssueでは、PyTorchと一緒に使った場合にCUDNNのエラーになったようで、こちらはXLA_PYTHON_CLIENT_MEM_FRACTION=.88としてJAXのGPU割り当てを制限する方法で解決していた。

MNISTの学習実行

export TF_FORCE_GPU_ALLOW_GROWTH=trueとして実行すると、Lambda Labs GPU CloudでJAX/Flaxで書いたCNNモデルが学習できました。

...
2023-09-18 05:25:48.968484: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.
...
train epoch: 10, loss: 0.007537598721683025, accuracy: 99.84666442871094
test epoch: 10, loss: 0.032926980406045914, accuracy: 99.0184326171875

論文紹介:Video-LLaMA

MetaがLLaMALLaMA2と公開したことで、どんどん加速しているように見えるLLMに関する研究ですが、LLMに関するOpen Challengesの一つとも考えられている様々なモダリティの取り込みに関する研究で、Video-LLaMAを読んでみました。

arxiv.org

github.com

Video-LLaMAで出来ること

GitHubのREADMEにある図を見ると分かりやすいです。

例えば、下の図では、Video-LLaMAに「何が聞こえるかを述べて」と尋ねると、「足音と家の中で犬が吠えている」と回答してます。また、「サングラスはかけているか?」に対してもYesと回答していて、動画の視覚的な内容と聴覚的な内容に対して、正しく回答を生成できています。

https://user-images.githubusercontent.com/18526640/244575825-7f7bddb2-5cf1-4cf4-bce3-3fa67974cbb3.gif

他にも、動画の時系列な認識が出来ている例として「船はどちらの方向に動いているのか」に対して、正しい方向を回答できています。

https://user-images.githubusercontent.com/18526640/244579143-7304ad6f-1009-46f1-aca4-7f861b636363.gif

Video-LLaMAのアーキテクチャ

Video-LLaMAでは、動画における視覚と聴覚の内容を認識できるようにするために、動画における時間変化を捉え、視覚と聴覚のデータを扱う二つのブランチを持つアーキテクチャを提案しています。

Vision Language Branchは、4つの要素で構成。動画フレームを入力として、フレーム毎の特徴を抽出するためのVisual Encoder、位置情報を追加するPosition Embedding Layer、フレーム毎の特徴を統合するVideo Q-former、LLMのテキスト埋め込みと同じ次元のベクトルを出力するためのLinear Layerです。 Visual Encoderは事前学習済みのViT-G/14とQ-formerを利用、Video Q-formerはBLIP2のQ-formerを実装している。

Audio Language Branchも同じように、4つの要素で構成。動画を短いクリップ(2秒間)にして、クリップ毎の音声データから特徴を抽出するAudio Encoder、位置情報を追加するPosition Embedding Layer、クリップ毎の特徴を統合するAudio Q-former、出力を生成するLinear Layerです。 Audio Encoderは事前学習済みのImageBind、Audio Q-formerもVideo Q-formerと同様にBLIP2のQ-formerを実装している。

Multi Branch Cross Modal Training

学習は2ステージで行い、最初のステージでは大規模なキャプションデータセットを使って学習、次のステージではinstruction following データセットでFine-tuningを行う。 二つのブランチは別々に学習を行うが、Audioとテキストが含まれるデータは希少なので、Audio BranchでもVision Branchと同じ動画キャプション(WebVid2M)、画像キャプション(CC595K)データセットを用いて学習する。Audio Branchは、Audio EncoderとしてImageBindをつかったことで、音声データで学習していないにも関わらず、推論時には音声を理解するように学習されたと述べている。 二つ目のステージでは、MiniGPT-4のimage-detail-description dataset、LLaVAのimage-instruction dataset、Video-Chatのvideo-instruction datasetを統合して使用する。

まとめ

以上が、Video-LLaMAのざっくりの内容です。Vision Language BranchとAudio Language Branchを用意して、動画フレーム/音声データに基づいた回答を行えるようにしているのはシンプルな仕組みだと思います。Audio Language Branchと言いながら、学習には音声データが使われていないのは不思議な感じですが。

JAXとcomposable program transformations

https://github.com/google/jaxのAboutは、次のように記述されています。

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Composable transformationsはどういうことなのか? NeurIPS2020: JAX Ecosystem Meetupの動画で、DeepMindのエンジニアの方が解説されていました。

NeurIPS 2020: JAX Ecosystem Meetup - YouTube

例として、次の関数を考えます。

def fn(x, y):
  return x**2 + y

fn(1., 2.) # (1**2 + 2) = 3

これに対して、gradientはどう書けるか?

df_dx = grad(fn)
df_dx(1., 2.) # df_dx = 2*x = 2*1 = 2

ここで、gradは関数を返す関数で、df_dxも関数になる。そして、通常の関数呼び出しで使用することができる。

さらに、second order gradientはどう書けるか?

df2_dx = grad(grad(fn))
df2_dx(1., 2.) # df2_dx = d(2*x)_dx = 2

gradはcomposableなため、gradをもう一つ追加するだけで良い。

composableなのはgradだけでなく、他の変換も使用することができる。compiled second-order gradientは以下のように実行できる。

df2_dx = jit(grad(grad(fn)))
df2_dx(1., 2.) # 2, ここでcompileされる
df2_dx(1., 2.) # 2, XLA pre compileのコードを実行、一回目よりも早い実行ができる

さらに、バッチ計算もcomposableに付け加えることができる。(batched compiled second-order gradient)

df2_dx = vmap(jit(grad(grad(fn))))
xs = jnp.ones((batch_size,))
df2_dx(xs, 2 * xs) # [2, 2] if batch_size=2

複数のアクセラレータ(GPUなど)で実行する場合も、composableに付け加えることができる。(multi-gpu batched compiled second-order gradient)

df2_dx = pmap(vmap(jit(grad(grad(fn)))))
xs = jnp.ones((num_gpus, batch_size,))
df2_dx(xs, 2 * xs) # [[2, 2], [2, 2]] if batch_size=2 and num_gpus=2

まとめ

以上が約5分のプレゼンで解説されていた内容ですが、分かりやすくて、 変換の組み合わせってそういうことか!と感動しました。 また動画には、HaikuやOptaxといったEcosystemの話や、他にもGANsなど様々なJAX実装の例があり、勉強になりました。 前回、JAX/FlaxでViTを実装してみましたが、今年はJAXをもっと使っていこうと思います。

ゼロから作るVision Transformer (JAX/Flax)

最近、Jax/Flaxを触るようになりました。ここでは、Jax/Flaxを用いてVision Transformerを実装した方法と、Jax/Flaxによる学習の方法について紹介しようと思います。

Vision Transformerのおさらい

Vision Transformerを実装するにあたって、まずはこの図を頭に入れておきます。

併せて、ViTの処理を論文で把握しておきます。

  1. 入力画像からパッチ画像を切り出し、フラットなベクトルに変換。
  2. Transformer Encoderで扱う隠れベクトルの次元へ射影
  3. Positional embedを追加、CLSトークンの追加
  4. Transformer Encoderで処理
  5. CLS embedを使ってMLPで分類

処理ブロックの実装

ここから、Jax/Flaxを使ってそれぞれの処理ブロックを作っていきます。

1. パッチ画像生成とフラットベクトル化

入力画像をH x W x Cとすると、patch_size x patch_size x Cのパッチ画像を切り出して、フラットなベクトルにします(embed_dim次元)。ここでは、カーネルサイズがpatch_size x patch_size、重複しないようにstridesをpatch_size x patch_sizeにした畳み込みで実装します。

class Patches(nn.Module):
  patch_size: int
  embed_dim: int

  def setup(self):
    self.conv = nn.Conv(
        features=self.embed_dim,
        kernel_size=(self.patch_size, self.patch_size),
        strides=(self.patch_size, self.patch_size),
        padding='VALID'
    )

  def __call__(self, images):
    patches = self.conv(images)
    b, h, w, c = patches.shape
    patches = jnp.reshape(patches, (b, h*w, c))
    return patches

2と3. D次元ベクトルへの射影とCLSトークン及びPosition Embeddingの追加

Transformer Encoderでは、全ての層で同じ次元サイズhidden_dimを使用します。先ほど作ったフラットなパッチ画像のベクトルを、hidden_dim次元ベクトルに射影します。また、BERTと同じように、分類に使用する特別なトークンとして、CLSトークンをパッチ系列の先頭に追加します。さらに、位置情報を保持するため、学習可能な1D position embeddingも追加します。

class PatchEncoder(nn.Module):
  hidden_dim: int

  @nn.compact
  def __call__(self, x):
    assert x.ndim == 3
    n, seq_len, _ = x.shape
    # Hidden dim
    x = nn.Dense(self.hidden_dim)(x)
    # Add cls token
    cls = self.param('cls_token', nn.initializers.zeros, (1, 1, self.hidden_dim))
    cls = jnp.tile(cls, (n, 1, 1))
    x = jnp.concatenate([cls, x], axis=1)
    # Add position embedding
    pos_embed = self.param(
        'position_embedding', 
        nn.initializers.normal(stddev=0.02), # From BERT
        (1, seq_len + 1, self.hidden_dim)
    )
    return x + pos_embed

4. Transformer Encoder

上で貼り付けた図にあるように、Transformer EncoderはMulti Head Self Attention(MSA)とMLPが交互に接続された層で構成され、MSAとMLPブロックの前にLayernorm(LN)、ブロックの後にresidual接続を適用します。

class TransformerEncoder(nn.Module):
  embed_dim: int
  hidden_dim: int
  n_heads: int
  drop_p: float
  mlp_dim: int

  def setup(self):
    self.mha = MultiHeadSelfAttention(self.hidden_dim, self.n_heads, self.drop_p)
    self.mlp = MLP(self.mlp_dim, self.drop_p)
    self.layer_norm = nn.LayerNorm(epsilon=1e-6)
  
  def __call__(self, inputs, train=True):
    # Attention Block
    x = self.layer_norm(inputs)
    x = self.mha(x, train)
    x = inputs + x
    # MLP block
    y = self.layer_norm(x)
    y = self.mlp(y, train)

    return x + y

MLPは2層のネットワークです。活性化関数はGELUを用いています。論文に従い、DropoutをDense層の後に適用しています。

class MLP(nn.Module):
  mlp_dim: int
  drop_p: float
  out_dim: Optional[int] = None

  @nn.compact
  def __call__(self, inputs, train=True):
    actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim
    x = nn.Dense(features=self.mlp_dim)(inputs)
    x = nn.gelu(x)
    x = nn.Dropout(rate=self.drop_p, deterministic=not train)(x)
    x = nn.Dense(features=actual_out_dim)(x)
    x = nn.Dropout(rate=self.drop_p, deterministic=not train)(x)
    return x
Multi Head Self Attention(MSA)

MSAについては、以前のブログMulti Head Attentionの概要を掴む - stMindにもまとめました。 ViTの論文では、qkvを求めるのは \rm{U_{qkv}}となっていますが、ここでは独立したDenseで実装しています。また、qkvは[B, N, T, D]の形にして、Single Headと同じようにWeightとAttentionを計算した後で、元の[B, T, C=N*D]に戻して出力するようにします。

class MultiHeadSelfAttention(nn.Module):
  hidden_dim: int
  n_heads: int
  drop_p: float

  def setup(self):
    self.q_net = nn.Dense(self.hidden_dim)
    self.k_net = nn.Dense(self.hidden_dim)
    self.v_net = nn.Dense(self.hidden_dim)

    self.proj_net = nn.Dense(self.hidden_dim)

    self.att_drop = nn.Dropout(self.drop_p)
    self.proj_drop = nn.Dropout(self.drop_p)

  def __call__(self, x, train=True):
    B, T, C = x.shape # batch_size, seq_length, hidden_dim
    N, D = self.n_heads, C // self.n_heads # num_heads, head_dim
    q = self.q_net(x).reshape(B, T, N, D).transpose(0, 2, 1, 3) # (B, N, T, D)
    k = self.k_net(x).reshape(B, T, N, D).transpose(0, 2, 1, 3)
    v = self.v_net(x).reshape(B, T, N, D).transpose(0, 2, 1, 3)

    # weights (B, N, T, T)
    weights = jnp.matmul(q, jnp.swapaxes(k, -2, -1)) / math.sqrt(D)
    normalized_weights = nn.softmax(weights, axis=-1)

    # attention (B, N, T, D)
    attention = jnp.matmul(normalized_weights, v)
    attention = self.att_drop(attention, deterministic=not train)

    # gather heads
    attention = attention.transpose(0, 2, 1, 3).reshape(B, T, N*D)

    # project
    out = self.proj_drop(self.proj_net(attention), deterministic=not train)

    return out

5. CLS embedを使ってMLPで分類

最後にこれまで作ったブロックをまとめて、MLP head(分類ヘッド)を付け加えます。

class ViT(nn.Module):
    patch_size: int
    embed_dim: int
    hidden_dim: int
    n_heads: int
    drop_p: float
    num_layers: int
    mlp_dim: int
    num_classes: int

    def setup(self):
        self.patch_extracter = Patches(self.patch_size, self.embed_dim)
        self.patch_encoder = PatchEncoder(self.hidden_dim)
        self.transformer_blocks = [
            Transformer(
                self.embed_dim, self.hidden_dim, self.n_heads, self.drop_p, self.mlp_dim
            )
            for _ in range(self.num_layers)
        ]
        self.mlp_head = MLP(self.mlp_dim, self.drop_p)
        self.cls_head = nn.Dense(features=self.num_classes)

    def __call__(self, x, train=True):
        x = self.patch_extracter(x)
        x = self.patch_encoder(x)
        for block in self.transformer_blocks:
            x = block(x, train)
        # MLP head
        x = x[:, 0]  # [CLS] token
        x = self.mlp_head(x, train)
        x = self.cls_head(x)
        return x

Jax/Flaxによる学習

モデルが作成できたので、Jax/Flaxを使って学習を組み立てていきます。

データセット

ここではtorchvisionのCIFAR10を使います。

def image_to_numpy(img):
  img = np.array(img, dtype=np.float32)
  img = (img / 255. - DATA_MEANS) / DATA_STD
  return img
def numpy_collate(batch):
  if isinstance(batch[0], np.ndarray):
    return np.stack(batch)
  elif isinstance(batch[0], (tuple, list)):
    transposed = zip(*batch)
    return [numpy_collate(samples) for samples in transposed]
  else:
    return np.array(batch)
test_transform = image_to_numpy
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomResizedCrop((IMAGE_SIZE, IMAGE_SIZE), scale=CROP_SCALES, ratio=CROP_RATIO),
    image_to_numpy
])

# Validation set should not use the augmentation.
train_dataset = CIFAR10('data', train=True, transform=train_transform, download=True)
val_dataset = CIFAR10('data', train=True, transform=test_transform, download=True)
train_set, _ = torch.utils.data.random_split(train_dataset, [45000, 5000], generator=torch.Generator().manual_seed(SEED))
_, val_set = torch.utils.data.random_split(val_dataset, [45000, 5000], generator=torch.Generator().manual_seed(SEED))
test_set = CIFAR10('data', train=False, transform=test_transform, download=True)

train_loader = torch.utils.data.DataLoader(
    train_set, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=2, persistent_workers=True, collate_fn=numpy_collate,
)
val_loader = torch.utils.data.DataLoader(
    val_set, batch_size=BATCH_SIZE, shuffle=False, drop_last=False, num_workers=2, persistent_workers=True, collate_fn=numpy_collate,
)
test_loader = torch.utils.data.DataLoader(
    test_set, batch_size=BATCH_SIZE, shuffle=False, drop_last=False, num_workers=2, persistent_workers=True, collate_fn=numpy_collate,
)

モデルの初期化

ViTクラスとして定義したモデルの初期化をします。FlaxのModule.initメソッドでPRNGキーとダミー入力を用いて初期化を実行し、戻り値として得られたパラメータを後で作成するTrainStateで管理、学習のループで更新していくという形で使用します。

def initialize_model(
    seed=42,
    patch_size=16, embed_dim=192, hidden_dim=192,
    n_heads=3, drop_p=0.1, num_layers=12, mlp_dim=768, num_classes=10
):
  main_rng = jax.random.PRNGKey(seed)
  x = jnp.ones(shape=(5, 32, 32, 3))
  # ViT
  model = ViT(
      patch_size=patch_size,
      embed_dim=embed_dim,
      hidden_dim=hidden_dim,
      n_heads=n_heads,
      drop_p=drop_p,
      num_layers=num_layers,
      mlp_dim=mlp_dim,
      num_classes=num_classes
  )
  main_rng, init_rng, drop_rng = random.split(main_rng, 3)
  params = model.init({'params': init_rng, 'dropout': drop_rng}, x, train=True)['params']
  return model, params, main_rng
vit_model, vit_params, vit_rng = initialize_model()

TrainStateの作成

Flaxでは、ステップ数やオプティマイザの状態、モデルのパラメータを含む学習の状態を管理するクラスを作成することが一般的なパターンのようです。そのため、基本的なユースケースに対応するTrainStateクラスが提供されています。また、apply_fnにモデルのforwardに相当するapplyを指定しておくことで、学習ループにおける関数の引数リストを少なくすることができます。

def create_train_state(
    model, params, learning_rate
):
  optimizer = optax.adam(learning_rate)
  return train_state.TrainState.create(
      apply_fn=model.apply,
      tx=optimizer,
      params=params
  )
state = create_train_state(vit_model, vit_params, 3e-4)

学習ループ

学習ループ自体は、Jax/Flaxに特有の記述はありません。

def train_model(train_loader, val_loader, state, rng, num_epochs=100):
  best_eval = 0.0
  for epoch_idx in tqdm(range(1, num_epochs + 1)):
    state, rng = train_epoch(train_loader, epoch_idx, state, rng)
    if epoch_idx % 1 == 0:
      eval_acc = eval_model(val_loader, state, rng)
      logger.add_scalar('val/acc', eval_acc, global_step=epoch_idx)
      if eval_acc >= best_eval:
        best_eval = eval_acc
        save_model(state, step=epoch_idx)
      logger.flush()
  # Evaluate after training
  test_acc = eval_model(test_loader, state, rng)
  print(f'test_acc: {test_acc}')
def train_epoch(train_loader, epoch_idx, state, rng):
  metrics = defaultdict(list)
  for batch in tqdm(train_loader, desc='Training', leave=False):
    state, rng, loss, acc = train_step(state, rng, batch)
    metrics['loss'].append(loss)
    metrics['acc'].append(acc)
  for key in metrics.keys():
    arg_val = np.stack(jax.device_get(metrics[key])).mean()
    logger.add_scalar('train/' + key, arg_val, global_step=epoch_idx)
    print(f'[epoch {epoch_idx}] {key}: {arg_val}')
  return state, rng
def eval_model(data_loader, state, rng):
  # Test model on all images of a data loader and return avg loss
  correct_class, count = 0, 0
  for batch in data_loader:
    rng, acc = eval_step(state, rng, batch)
    correct_class += acc * batch[0].shape[0]
    count += batch[0].shape[0]
  eval_acc = (correct_class / count).item()
  return eval_acc

Train step

train_stepの中では、ロス関数を定義し、モデルのパラメータに対する勾配を求め、勾配に基づいてパラメータの更新を行います。value_and_gradsでstate.paramsに対する勾配を求め、apply_gradientsでTrainStateを更新します。ロス関数の中では、TrainStateの作成の時に指定したapply_fn(model.applyと同じ)でlogitsを計算して、cross entropy lossを求めます。

@jax.jit
def train_step(state, rng, batch):
  loss_fn = lambda params: calculate_loss(params, state, rng, batch, train=True)
  # Get loss, gradients for loss, and other outputs of loss function
  (loss, (acc, rng)), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
  # Update parameters and batch statistics
  state = state.apply_gradients(grads=grads)
  return state, rng, loss, acc
def calculate_loss(params, state, rng, batch, train):
  imgs, labels = batch
  rng, drop_rng = random.split(rng)
  logits = state.apply_fn({'params': params}, imgs, train=train, rngs={'dropout': drop_rng})
  loss = optax.softmax_cross_entropy_with_integer_labels(logits=logits, labels=labels).mean()
  acc = (logits.argmax(axis=-1) == labels).mean()
  return loss, (acc, rng)

学習結果

モデルの設定とハイパーパラメータはViT-Tinyを参考にして学習した結果です。Colab proの標準GPUで1.5h程度でした。

test_acc: 0.7704000473022461

まとめ

Jax/Flaxを用いてVision Transformerを実装する方法、学習と評価を行う方法を紹介しました。Test精度を上げるには、もう少し手を入れる必要がありそうですが、基本的なところは実現できたように思います。

github.com

Reference

pytorchとtensorflowのチュートリアル写経で見つけた小さなバグをPull Requestしてマージされるまで

正確にいうと、torchvisionとTensorflow Textのチュートリアル

torchvisionの方は、ビデオファイルを読み込んでビデオフレームとオーディオフレームを返すVideo APIのチュートリアルで、ビデオのptsはvideo_ptsにappendして、オーディオのptsはaudio_ptsにappendするべきところが、オーディオの方もvideo_ptsにappendされていた。単にvideo_ptsをaudio_ptsに置き換えるだけのPRを作成して、3日ほどでマージされました。

github.com

Tensorflow Textの方は、Transformerを使った機械翻訳のチュートリアルで、ポルトガル語を処理するTokenizerがtokenizers.enとなっていたので、tokenizers.ptと修正してPRを作成。こちらも2日ほどでマージされました。現在、このチュートリアルは結構書き換えられていて、修正した箇所を含むブロックは無くなっています:)

github.com

今回はチュートリアル写経の副産物としてのPRでしたが、このような貢献の形もあるよという紹介でした。