TFとPyTorchだけを知っているだけでは不十分です。PyTorchから自動微分ツールJAXに切り替える方法を見てみましょう。

現在のディープラーニング フレームワークに関しては、TensorFlow と PyTorch を避けることができないことがよくあります。しかし、これら 2 つのフレームワークに加えて、いくつかの新しい力も過小評価すべきではありません。その 1 つが JAX です。前方および後方の自動微分機能を備えており、高次導関数の計算に非常に優れています。この注目のフレームワークはどれほど便利なのでしょうか? ニューラル ネットワーク内の複雑な勾配更新とバックプロパゲーションを実証するには、どのように使用できるでしょうか? この記事は、Jax の基礎となるロジックを理解し、PyTorch や他のフレームワークからの移行を容易にするためのチュートリアル記事です。

[[326161]]

Jax は、機械学習と数学計算用に Google が開発した Python ライブラリです。起動すると、Jax は Python + NumPy パッケージとして定義されました。微分化、ベクトル化、TPU および GPU での JIT 言語の使用などの機能があります。つまり、これは自動微分も可能な numpy の GPU バージョンです。 Skye Wanderman-Milne 氏のような研究者も、昨年の NeurlPS 2019 カンファレンスで Jax を紹介しました。

しかし、開発者にとって、すでに使い慣れている PyTorch や TensorFlow 2.X から Jax に切り替えることは間違いなく大きな変化です。この 2 つは計算とバックプロパゲーションの構築方法が根本的に異なるからです。 PyTorch は計算グラフを構築し、順方向パスと逆方向パスを計算します。結果ノードの勾配は、中間ノードの勾配から累積されます。

Jax は違います。計算を Python 関数として表現し、grad() を使用して評価可能な勾配関数に変換します。しかし、結果ではなく、結果の勾配が示されます。両者の比較を以下に示します。

これにより、プログラミングとモデルの構築方法が変わります。したがって、テープベースの自動差別化方法を使用し、ステートフル オブジェクトを使用できます。しかし、Jax は grad() 関数を実行するときに微分処理を関数のように動作させるので、驚くかもしれません。

おそらく、flax、trax、haiku などの Jax ベースのツールを検討してみることにしたでしょう。 ResNet のような例を見ると、他のフレームワークのコードとは異なることがわかります。レイヤーを定義してトレーニングを実行する以外に、基礎となるロジックは何ですか? これらの小さな NumPy プログラムはどのようにして巨大なアーキテクチャをトレーニングするのでしょうか?

この記事は、Jax を使用してモデルを構築する方法に関するチュートリアルです。Machine Heart では、次の 2 つの部分を取り上げました。

• PyTorch での LSTM-LM アプリケーションの簡単なレビュー。
• PyTorch スタイルのコード (mutate 状態に基づく) を見て、純粋な関数がモデルを構築する方法 (Jax) を学びます。

PyTorch 上の LSTM 言語モデル

まず、PyTorch を使用して LSTM 言語モデルを実装します。コードは次のとおりです。

輸入トーチ
クラス LSTMCell(torch.nn.Module):
    def __init__(self, in_dim, out_dim):
        super(LSTMCell、self).__init__()
 self.weight_ih = torch.nn.Parameter (torch.rand(4*out_dim, in_dim))
 self.weight_hh = torch.nn.Parameter (torch.rand(4*out_dim, out_dim))
自己バイアス= torch.nn.パラメータ(torch.zeros(4*out_dim,)) 
 
    def forward(自己、入力、h、c):
 ifgo = self .weight_ih@inputs + self .weight_hh@h + self .bias
        i、f、g、 o = torch.chunk (ifgo、4)
 i =トーチ.シグモイド(i)
 f =トーチシグモイド(f)
 g =トーチ.tanh (g)
 o =トーチ.シグモイド(o)
 new_c = f * c + i * g
 new_h = o * torch.tanh(new_c)
        戻り値 (new_h, new_c)

次に、この LSTM ニューロンに基づいて単層ネットワークを構築します。ここには埋め込みレイヤーがあり、学習可能な (h,c)0 と組み合わせることで、個々のパラメータがどのように変化するかを示します。

クラス LSTMLM(torch.nn.Module):
    def __init__(self, vocab_size, dim = 17 ):
        スーパー().__init__()
セルをLSTMCellで囲みます。
自己埋め込み= torch.nn.パラメータ(torch.rand(vocab_size, dim))
 self.c_0 = torch.nn.パラメータ(torch.zeros(dim)) 
 
    @財産
    hc_0(自分)を定義します。
        戻り値 (torch.tanh(self.c_0), self.c_0) 
 
    def forward(self, seq, hc):
損失= torch.tensor (0.)
          seq内のidxの場合:
              損失- = torch.log_softmax (self.embeddings@hc[0], dim =-1)[idx]
 hc =自己.セル(自己.埋め込み[idx,:], *hc)
          リターンロス、HC 
 
    greedy_argmax(self, hc,長さ= 6 )を定義します:
        torch.no_grad() の場合:
 idxs = []
            i が範囲(長さ)内にある場合:
 idx = torch.argmax (self.embeddings@hc[0])
                idxs.append(idx.item())
 hc =自己.セル(自己.埋め込み[idx,:], *hc)
        idxを返す

構築後、トレーニング:

トーチ.マニュアル_シード(0)
 # トレーニングデータとして、単語/単語部分/文字のインデックスを用意します。
 # これらはトークン化され、整数化されていると仮定します (もちろん、これはおもちゃの例です)。
 jax.numpyをjnpとしてインポートする
vocab_size = 43 # プライムトリック! :)
トレーニングデータ= jnp.array ([4, 8, 15, 16, 23, 42]) 
 
 lm = LSTMLM (語彙サイズ語彙サイズ= 語彙サイズ)
 print("前のサンプル:", lm.greedy_argmax(lm.hc_0)) 
 
 bptt_length = 3 # hc.detach-ingを説明するため
 
範囲(101)のエポックの場合:
 hc = lm.hc_0
総損失= 0 。
    開始範囲(0、len(training_data)、bptt_length)の場合:
バッチ=トレーニングデータ[開始:開始+bptt_length]
        損失、(h, c) = lm(バッチ、hc)
 hc = (h.detach()、c.detach()) です。
        エポック% 50 == 0の場合:
            総損失 += 損失.item()
        損失.後方()
        lm.named_pa​​rameters() の name、param の場合:
            param.gradがNoneでない場合:
                パラメータデータ- = 0.1 * パラメータ勾配
                del param.grad
     全損の場合:
         print("損失:", totalloss) 
          
 print("後のサンプル:", lm.greedy_argmax(lm.hc_0))
以前のサンプル: [42, 34, 34, 34, 34, 34]
損失: 25.953862190246582
損失: 3.7642268538475037
損失: 1.9537211656570435
後のサンプル: [4, 8, 15, 16, 23, 42]

ご覧のとおり、PyTorch コードは比較的明確ですが、まだいくつか問題があります。非常に注意していますが、計算グラフ内のノードの数に注意を払うことは依然として重要です。これらの中間ノードは適切なタイミングでクリアする必要があります。

純粋関数

JAX がこれをどのように処理するかを理解するには、まず純粋関数の概念を理解する必要があります。これまでに関数型プログラミングを行ったことがある場合、純粋関数は数学の関数や数式のようなものだという概念に馴染みがあるかもしれません。特定の入力値から出力値を取得する方法を定義します。重要なのは、この関数には「副作用」がないこと、つまり関数のどの部分もグローバル状態にアクセスしたり変更したりしないことです。

Pytorch で記述するコードには中間変数や状態が多数含まれており、これらの状態は頻繁に変化するため、推論と最適化が非常に難しくなります。したがって、JAX はプログラマーを純粋関数のスコープ内に制限し、上記の状況が発生しないようにすることを選択します。

JAX について詳しく説明する前に、純粋関数の例をいくつか見てみましょう。純粋関数は次の条件を満たす必要があります。

• 関数を実行する状況と実行時期は出力に影響を与えません。入力が変わらない限り、出力も変わりません。
• 関数を 0 回実行したか、1 回実行したか、あるいはそれ以上実行したかは、後で区別がつかなくなるはずです。

次の不純な関数は、上記の条件の少なくとも 1 つに違反します。

ランダムにインポート
インポート時間
実行回数= 0   
 
 pure_fn_1(x)を定義します:
    2 * xを返す
     
 pure_fn_2(xs)を定義します:
 ys = []
    x が xs 内にある場合:
        # 関数 *内部* で状態のある変数を変更しても問題ありません。
        ys.append(2 * x)
    ysを返す
 
不純なfn_1(xs)を定義します:
    # 引数を変更すると、関数の外部に永続的な影響が生じます! :(
    xs.append(合計(xs))
    xsを返す
 
不純なfn_2(x)を定義します:
    # 明らかに変異している
    地球の状態は悪いです...地球
    実行回数 nr_executions += 1
    2 * xを返す
 
不純なfn_3(x)を定義します:
    # ...しかし、アクセスするだけでも、関数は
    # 実行コンテキスト!
    nr_executions * x を返す
 
不純なfn_4(x)を定義します:
    # IO のようなものは不純度の典型的な例です。
    # 次の 3 行はすべて純粋性の違反です。
    print("こんにちは!")
ユーザー入力= 入力()
実行時間= time.time()
    2 * xを返す
 
不純なfn_5(x)を定義します:
    # これはどちらの制約に違反しますか? 実は両方です! 現在の
    # ランダム性の状態 *そして* 数値ジェネレーターを進化させます!
 p =ランダム.ランダム()
    p * x を返す
JAX が操作する純粋な関数を見てみましょう。導入図の例です。 
 
 # （ほぼ）1次元線形回帰
f(w, x)を定義します。
    w * xを返す
 
印刷(f(13., 42.))
 546.0

今のところ何も起こっていません。 JAX を使用すると、次の関数を、結果を返す代わりに、関数の最初の引数に対する関数の結果の勾配を返す別の関数に変換できるようになりました。

 jaxをインポートする
jax.numpyをjnpとしてインポートする
 
 # 勾配: 重みに関して! JAX はデフォルトで最初の引数を使用します。
 df_dw = jax.grad (f) 
 
 def manual_df_dw(w, x):
    xを返す
     
 df_dw(13., 42.) == manual_df_dw(13., 42.) をアサートする
 
印刷(df_dw(13., 42.))
 42.0

ここまでで、JAX README ドキュメントの内容をすべてご覧になったと思いますが、その内容は妥当なものです。しかし、PyTorch コードのように大きなモジュールにジャンプするにはどうすればよいでしょうか?

まず、バイアス項を追加し、1 次元の線形回帰変数を、使い慣れたオブジェクト (LinearRegressor「レイヤー」) にラップしてみます。

クラス LinearRegressor():
    __init__(self, w, b)を定義します。
自己.w = w
自己.b = b 
     
    def predict(自己, x):
        self.w * x + self.b を返す
         
    定義rms(self, xs: jnp.ndarray, ys: jnp.ndarray):
        jnp.sqrt(jnp.sum(jnp.square(self.w * xs + self.b - ys))) を返します。 
         
 my_regressor =線形回帰(13., 0.) 
 
 # トレーニングに使用される損失関数の一種
xs = jnp .array([42.0])
 ys = jnp.array ([500.0])
印刷(my_regressor.rms(xs, ys)) 
 
 # テストデータの予測
印刷(my_regressor.predict(42.))
 46.0
 546.0

トレーニングに勾配をどのように使用するのでしょうか? モデルの重みを入力パラメータとして受け取る純粋な関数が必要です。次のようなものになります。

 loss_fn(w, b, xs, ys)を定義します。
 my_regressor =線形回帰(w, b)
    my_regressor.rms( xs xs =xs, ys ys =ys)を返します
     
 # argnums = (0, 1) を使ってJAXに渡すよう指示します
# 最初のパラメータと 2 番目のパラメータに関するグラデーション。
 grad_fn = jax.grad (loss_fn、引数=(0, 1)) 
 
印刷(loss_fn(13., 0., xs, ys))
 print(grad_fn(13., 0., xs, ys))
 46.0
 (デバイス配列(42., dtype = float32 ), デバイス配列(1., dtype = float32 ))

これが正しいのだと自分自身を納得させなければなりません。さて、これは機能しますが、明らかに、 loss_fn の定義部分ですべてのパラメーターを列挙することは実現可能ではありません。

幸いなことに、JAX はスカラー、ベクトル、行列だけでなく、多くのツリーのようなデータ構造も区別できます。この構造は pytree と呼ばれ、Python 辞書で構成されています。

 loss_fn(パラメータ、xs、ys)を定義します。
 my_regressor =線形回帰(params['w'], params['b'])
    my_regressor.rms( xs xs =xs, ys ys =ys)を返します
 
 grad_fn = jax.grad (損失_fn) 
 
印刷(loss_fn({'w': 13., 'b': 0.}, xs, ys))
 print(grad_fn({'w': 13., 'b': 0.}, xs, ys))
 46.0
 {'b': DeviceArray(1., dtype = float32 ), 'w': DeviceArray(42., dtype = float32 )} これで見た目も良くなりました! トレーニング ループは次のように記述できます。

これで見た目はずっと良くなりました! トレーニング ループは次のように記述できます。

パラメータ= {'w': 13., 'b': 0.} 
 
 _ が範囲内(15)の場合:
    印刷(loss_fn(params, xs, ys))
 grads = grad_fn (パラメータ、xs、ys)
    params.keys() 内の名前の場合:
        params[名前] - = 0.002 * grads[名前] 
         
 # さて、予測してみましょう:
線形回帰(params['w'], params['b']).predict(42.)
 46.0
 42.47003
 38.940002
 35.410034
 31.880066
 28.350098
 24.820068
 21.2901
 17.760132
 14.230164
 10.700165
 7.170166
 3.6401978
 0.110198975
 3.4197998
デバイス配列(500.1102、 dtype = float32 )

より多くの JAX ヘルパーを使用して自分自身を更新できるようになったことに注意してください。パラメーターとグラデーションには共通の (ツリーのような) 構造があるため、これらを一番上に配置し、次のように、どこでも値が 2 つのツリーの「組み合わせ」である新しいツリーを作成することが考えられます。

 update_combiner を定義します(param, grad, lr = 0.002 ):
    戻り値パラメータ - lr * grad 
     
パラメータ= jax.tree_multimap (update_combiner、パラメータ、グラデーション)
 ＃ の代わりに：
 # params.keys() 内の名前:
 # params[名前] - = 0.1 * grads[名前]

参考リンク: https://sjmielke.com/jax-purify.htm

[この記事は51CTOコラム「Machine Heart」、WeChatパブリックアカウント「Machine Heart（id:almosthuman2014）」によるオリジナル翻訳です]

この著者の他の記事を読むにはここをクリックしてください