Claude CodeでPyTorchのprofilerを読み、DataLoaderの待ちを直す手順

PyTorchの学習が遅いときに最初にやることは、直す前に「GPUが計算している時間」と「データを待っている時間」を分けることです。Claude Codeに任せる場合も同じで、torch.profilerの出力とsamples/secの測定値を渡せば、num_workersやpin_memory、前処理のどこを触るかを数字で決められます。この記事では、計測コードを先に入れ、DataLoaderの待ちを見つけ、設定を一つずつ変えて測り直す流れを、Claude Codeへの頼み方と一緒に示します。

遅い学習は「待ち」と「計算」のどちらか

学習ループは、おおまかに次の2つに時間を使います。

• バッチが届くのを待つ時間(DataLoaderがディスク読み込みや前処理をしている間)
• 届いたバッチで順伝播・逆伝播・更新を計算する時間

num_workersの既定値は0で、この場合データの読み込みはメインプロセスで同期的に行われます。PyTorchのチューニングガイドは、その結果としてメインの学習プロセスがデータの準備を待つことになると説明しています。num_workersを正の整数にすると、読み込みと学習が重なります。

つまり、GPUがあまり働いていないのに1ステップが長いなら、疑う場所はデータ側です。逆に待ちがほぼ無いなら、num_workersをいじっても効きません。Claudeに「速くして」とだけ頼むと、効かない場所から手を付けがちです。まず計測で分けます。

計測コードを先に入れる

Claudeに最適化を頼む前に、再現できる計測スクリプトを用意します。入れる計測は2つです。

1. record_functionでラベルを付け、データ待ちと計算を別の行としてprofilerの表に出す
2. 同じ条件でsamples/secを測る

record_functionは、任意のコード範囲に名前を付けるコンテキストマネージャです。プロファイラの表では、その名前が1行として現れます。next(it)にだけdata_waitというラベルを付ければ、バッチ待ちの時間が表で分離できます。

# bench_loader.py
import contextlib
import time
import torch
from torch.profiler import profile, schedule, ProfilerActivity, record_function
 
def run(loader, model, opt, device, steps=60, profile_steps=False):
    activities = [ProfilerActivity.CPU]
    if device.type == "cuda":
        activities.append(ProfilerActivity.CUDA)
    sch = schedule(skip_first=5, wait=1, warmup=1, active=4, repeat=1)
    it = iter(loader)
    seen = 0
    if device.type == "cuda":
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    ctx = (profile(activities=activities, schedule=sch)
           if profile_steps else contextlib.nullcontext())
    with ctx as prof:
        for _ in range(steps):
            with record_function("data_wait"):
                x, y = next(it)
            with record_function("train_step"):
                x = x.to(device, non_blocking=True)
                y = y.to(device, non_blocking=True)
                loss = torch.nn.functional.cross_entropy(model(x), y)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                opt.step()
            seen += x.shape[0]
            if profile_steps:
                prof.step()
    if device.type == "cuda":
        torch.cuda.synchronize()
    sps = seen / (time.perf_counter() - t0)
    if profile_steps:
        print(prof.key_averages().table(
            sort_by="self_cpu_time_total", row_limit=12))
    print(f"samples/sec: {sps:.1f}")

scheduleは、長時間の学習を全部トレースすると遅くなり、ファイルも巨大になるために用意されている仕組みです。上の例のskip_first=5, wait=1, warmup=1, active=4は、最初の5ステップを飛ばし、1ステップ待ってから1ステップを助走に使い、4ステップだけ記録する指定です。助走のステップは、計測開始時の余計な負荷で値が歪むため、結果が捨てられます。

prof.step()を呼ぶ位置はループの末尾です。呼ばないと、scheduleが進まず記録が始まりません。

profile_stepsは実行の使い分けです。Trueの実行はプロファイラを有効にして表を出し、Falseの実行はプロファイラを一切通さずsamples/secだけを出します。表を読む実行と、前後を比べる数字を取る実行を分けるためです。ループの長さ(steps)がDataLoaderの長さを超えるとnext(it)がStopIterationで止まるので、データが少ないときはstepsを下げます。表のソートはself_cpu_time_totalで、data_waitとtrain_stepの行はこの列で比べます。

表の読み方

出力された表で見る行は、まずdata_waitとtrain_stepの2つです。

この表に具体的な秒数の基準は書きません。判断の材料にするのは、data_waitとtrain_stepの比と、設定を変えたあとのsamples/secの差です。

Claude Codeに頼む: 計測結果を渡して一つずつ直す

計測結果が出たら、それをClaudeに渡します。非対話モード(-p)は標準入力を読めるため、プロファイラの出力をそのままパイプで渡せます。

# profile_steps=True の実行で表を出力する
python bench_loader.py > before.txt 2>&1
cat before.txt | claude -p \
  'data_wait と train_step の比から、DataLoader の待ちが主因か判断し、
   最初に変える設定を1つだけ挙げて' --allowedTools "Read"

標準入力に渡せる量は10MBまでで、超えるとClaude Codeはエラーで終了します。プロファイラの表はrow_limitを付けていれば数KBに収まりますが、Chromeトレースのjsonは大きくなりがちなので、パイプでなくファイルのパスを渡します。

「1つだけ」と頼んでいるのには理由があります。設定を3つ同時に変えると、どれが効いたのか分からなくなります。Claudeは親切に複数案を出しますが、評価の単位を1変更に保つには、頼み方で縛るのが確実です。

CLAUDE.mdに書いておく規約

毎回同じ頼み方をするなら、プロジェクトのCLAUDE.mdに手順として書きます。次は、そのまま置ける最小の例です。

## 学習が遅いときの手順
 
- 変更する前に `python bench_loader.py` を実行し、before の `samples/sec` を記録する
- 1回に変えるのは DataLoader の引数1つか、前処理1か所だけ
- 変更後に同じコマンドで測り直し、after の `samples/sec` を並べて報告する
- 測り直しで下がった変更は元に戻す
- 比較の実行ではプロファイラを無効にする(`profile_steps=False`)

数字が改善したかどうかをClaudeの印象でなく測定値で決める、というのがこの規約の役割です。

自動化するときの許可設定

Claudeに計測と編集を続けて任せるなら、--allowedToolsで使えるツールを明示します。許可を与えた範囲で、確認なしに動きます。

claude -p "bench_loader.py を実行して before を記録し、CLAUDE.md の手順で DataLoader を改善して" \
  --allowedTools "Read,Edit,Bash(python bench_loader.py)" \
  --max-turns 15

Bash(python bench_loader.py)のようにコマンドを絞れば、Claudeが実行できるシェルコマンドを計測スクリプトだけに限れます。--max-turnsは、ターン数の上限を決めるフラグで、上限に達するとエラーで終了します。学習そのものを回すコマンドは、時間もGPUも使うため、許可リストには入れず手元で実行する運用が無難です。

DataLoaderを直す順番

待ちが主因と分かったら、次の順で一つずつ変えて測ります。順番には意味があります。上の段ほど「コードを変えずに済む」変更です。

num_workers: まず0をやめる

num_workersには、これが最適という値がありません。PyTorchのチューニングガイドも、ワークロード、CPU、GPU、学習データの置き場所に応じて調整するよう書いています。Claudeに頼むなら、「2, 4, 8で測って」と探索範囲を渡し、表で比べさせます。

for w in (0, 2, 4, 8):
    loader = torch.utils.data.DataLoader(ds, batch_size=64, num_workers=w)
    print("num_workers =", w)
    run(loader, model, opt, device, profile_steps=False)

ワーカーを増やすと、プロセスが増えます。DataLoaderは反復子を作るたびにnum_workers個のワーカープロセスを起こし、dataset・collate_fn・worker_init_fnを各ワーカーへ渡します。データセットのI/Oと変換(collate_fnを含む)は、ワーカー側で動きます。この構造のため、前処理が重いほどワーカーを増やす効果が出ます。

pin_memory: ホストからGPUへの転送

ホストからGPUへのコピーは、固定(ページロック)メモリから始めたほうが速く、pin_memory=Trueを渡すとDataLoaderがテンソルを固定メモリへコピーしてから返します。GPUを使うならTrueのほうが良い、というのがチューニングガイドの記述です。pin_memoryとnon_blockingの使い分けには、PyTorch公式に専用のガイドがあります。

prefetch_factorとpersistent_workers

prefetch_factorは、各ワーカーがあらかじめ読み込んでおくバッチ数で、num_workers > 0のときの既定値は2です。persistent_workers=Trueにすると、データセットを一巡してもワーカーを終了せず、データセットのインスタンスが維持されます。エポックごとにワーカーを作り直すコストが目立つときだけ試します。

前処理の置き場所

ワーカーを増やしても伸びない場合、前処理そのものが重い可能性があります。Claudeには、__getitem__の中をtime.perf_counter()で区切って測るコードを入れさせ、「読み込み」「デコード」「拡張」のどれが長いかを出させます。重い変換は、事前に1回だけ実行して結果をディスクに保存しておく方法もあります。この変更はデータの作り方が変わるため、効果と引き換えの影響(ストレージ量、再現性)をClaudeに挙げさせてから採用します。

変更前後を並べて報告させる

最後に、before/afterをClaudeに1表で出させます。ここでも数字は測定値だけを使わせます。

報告の表には、変更の内容と前後のsamples/secを1行ずつ書いてもらいます。値が下がった行も残してもらうと、「試して戻した」という経緯が後から読めます。

つまずきやすい点

計算側が重いときの次の候補

data_waitがごくわずかなら、DataLoaderを直しても効きません。チューニングガイドは、そのほかの項目としてoptimizer.zero_grad(set_to_none=True)、検証や推論での勾配計算の無効化、torch.backends.cudnn.benchmark = True、混合精度(AMP)などを挙げています。この記事の手順は、同じ型で使えます。1つ変え、同じ条件で測り、samples/secの差で採否を決めます。

計測を先に置く進め方は、CIの待ち時間を縮めるGitHub Actionsの実行時間の短縮でも同じです。学習後にモデルを書き出す段階では、ONNXへの書き出しと出力の一致テストが次の作業になります。

まとめ

PyTorchの学習が遅いと感じたら、record_functionでdata_waitとtrain_stepを分けて表に出すところから始めます。待ちが主因ならnum_workers、pin_memory、前処理の順に一つずつ変え、プロファイラ無しのsamples/secで採否を決めます。Claude Codeには、計測コードの用意、結果の読み取り、CLAUDE.mdの規約に沿った1変更ずつの試行を任せられます。一方で、学習の本番実行や、データの作り方を変える判断は、人が確認する場所に残しておくと安全です。