GoogleのAI学習ツール解説

このチュートリアルでは、Google Researchが開発したJAXトレーニングライブラリであるKauldronを実装します。Kauldronは、研究の速度とモジュール性を最適化すると説明されており、私たちはその言葉を文字通りに受け止め、実際に何をもたらすかを検証します。まずインストールし、ノートブックの前半では、KauldronをFlaxとOptaxのスタックと異なるものにする3つのメカニズムに焦点を当てます。それは、実験をJSONを介して往復できるプレーンな辞書のツリーに変換するkonfig、損失がスコアを付けるモデルをインポートすることなく、文字列キーパスでコンポーネントを接続するkontext、そして、名前付き軸が引数間で結合され、一致しない場合に何に結合されていたかを報告するランタイムシェイプチェッカーです。

次に、フレームワークが期待する形式でカスタム損失とカスタムメトリックを記述し、ダウンロードやアクセラレータなしで、合成のインメモリデータ上で実際のTrainerをトレーニングし、モデルを編集せずにモデルの内部レイヤーを監視します。最後に、すべての実験が1つの設定行で異なる5つのバリアントスイープを実行し、トレーニング実行が自身をチェックポイントし、停止した場所から再開できるようにします。

Copy Code Copied Use a different Browser

import os

import sys

import json

import textwrap

import traceback

import subprocess

RESULTS = {}

def banner(title):

print("\n" + "=" * 78)

print(title)

print("=" * 78)

def section(name):

def wrap(fn):

def run(*a, **kw):

banner(name)

try:

out = fn(*a, **kw)

RESULTS[name] = out if isinstance(out, str) else "ok"

return out

except Exception as e:

RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"

print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")

traceback.print_exc(limit=3)

return None

return run

return wrap

banner("0. Install Kauldron, and the one compatibility patch you need today")

subprocess.run([sys.executable, "-m", "pip", "install", "-q", "kauldron==1.4.2"], check=True)

import jax

from etils.enp import array_spec as _array_spec

# jax >= 0.10.1 moved `jax._src.prng`, but etils <= 1.14.0 s