Kaggleに役立つ高速化・並列化テクニック

2025/8/23 第4回関東Kaggler会の登壇資料です。
https://kanto-kaggler.connpass.com/event/362280/

13:15~ 山口祐 ymg_aq
「Kaggleに役立つ 高速化・並列化テクニック」

Kaggle上のPython実装を、さまざまな手段で高速化・並列化する手法について、実際のコンペで役に立った実例も交えて紹介します。

Transcript

1. Kaggleに役⽴つ ⾼速化‧並列化テクニック 第4回 関東 Kaggler 会 ⼭⼝ 祐 / ymg_aq 2025.8.23
2. ⾃⼰紹介 ⼭⼝ 祐 @ymg_aq • AIST/⽶NISTで研究の傍ら、ゲームAIを開発 • 上場企業執⾏役員を経てチューリングCTOに • 2025年1⽉からKaggleを始め、Masterに 2 Kaggleに入門しました | ymg_aq
3. ⾃⼰紹介 3 • 囲碁のプロ棋⼠にゼロから5時間だけ 強化学習したAIで対局する ABCI 1.0 (https://www.aist.go.jp/aist_j/magazine/20180901.html) • ABCI（2018年当時TOP500 5位）を 最⼤1/4 (1100GPU) 占有して並列分散学習していた → 勝利 → ⾼速化‧並列化によって短時間で超⼈間級に
4. コンペ経歴 4 🥇 🥈 🥈 🥇 ⾼速化と並列化こそKaggleの本質…？🤔 モデル学習 ⾼速化 並列化 ✅ ✅ ✅ ✅ ✅ ✅ ❌ ❌ ❌ ❌ 🔺 ✅
5. • ⾼速化 … プログラムを改良して処理時間を短くすること ◦ アルゴリズムの改善 ◦ プログラミング⾔語‧処理系の選定 ◦ メモリ‧キャッシュ‧I/O最適化 ◦ ∋ 並列化 • 並列化 … 複数の処理を同時に実⾏できる形に分割すること ◦ タスク並列化 ◦ データ並列化 ◦ 複数スレッド/プロセスの利⽤ ⾼速化‧並列化 5
6. ☑ 本⽇のゴール Pythonの処理を 速くする⽅法を知る 5000倍 の ※ Kaggle Notebookで ちゃんと使えます
7. Step1: アルゴリズムを知る 7 • アルゴリズム‧データ構造 ◦ ⼆分探索 ◦ 動的計画法 ◦ Union-Find ◦ セグメント⽊ ◦ etc ... • 探索‧ヒューリスティック ◦ ビームサーチ ◦ 焼きなまし法 ◦ 遺伝的アルゴリズム ◦ etc ... 問題解決力を鍛える！アルゴリズムとデータ構造 計算量削減はどんな⾼速化よりも強い
8. 役に⽴った例(1) （複数GPUへの割り振り） 8 • 複数GPUへのタスク配分を最適化 ◦ 不均衡なシーン（S0, S1, …Sn-1枚）を含むデータセット ◦ 2枚のGPUに処理が均等になるよう割り当てたい 🥇 複数の画像からカメラ姿勢を推定 するコンペ。 今年はデータセット内のシーンを 分離するタスクが追加 GPU0 GPU1 追加の 処理時間 データ量 均等な データ量 GPU0 GPU1 データ量 ✅ ❌
9. 役に⽴った例(1) （複数GPUへの割り振り） 9 • 複数GPUへのタスク配分を最適化 ◦ 不均衡なシーン（S0, S1, …Sn-1枚）を含むデータセット ◦ 2枚のGPUに処理が均等になるよう割り当てたい • 部分和問題を動的計画法で解く ◦ 単純な貪欲法より最⼤20%処理時間を削減（※LPTとでは数%程度） 🥇 def split_scenes_with_dp(data_num_list: list[int], target_scenes: list[str]) -> tuple[list[str], list[str]]: n = len(data_num_list) total = sum(data_num_list) half = total // 2 dp = {0: []} for idx, val in enumerate(data_num_list): for s, subset in list(dp.items()): new_s = s + val if new_s not in dp: dp[new_s] = subset + [idx] # まだ登録されていなければ記録 # half に最も近い和を探す best_sum = min(dp.keys(), key=lambda s: abs(half - s)) groupA_idx = dp[best_sum] … 複数の画像からカメラ姿勢を推定 するコンペ。 今年はデータセット内のシーンを 分離するタスクが追加
10. Step2: 実⾏速度を知る 10 • 推測するな、計測せよ ◦ ボトルネックとなっている処理‧ 原因を特定する ◦ perf/py-spy/cProfile/GPU Nsight • ⾔語‧HWの特性を意識する ◦ C/C++/Python ... ◦ CPU/GPUのスペック どのくらい/何が遅いのかを調べる github.com/jlfwong/speedscope
11. Kaggleの実⾏環境 11 Accelerator CPU vCPU Clock RAM GPU Memory CUDA/ Tensor Core TFLOPS (FP16) None (CPU only) Broadwell-EP (Xeon E5-2698) 4 2.2GHz 32GB - - - - P100 Skylake-SP 4 2.0GHz 32GB Pascal (2016) 16GB (HBM2) 3,584 / - 18.7 T4 x2 Skylake-SP 4 2.0GHz 32GB Turing (2018) 16GB (GDDR6) x2 2,560 / 320 130 (total) L4 x4 Skylake-SP 48 2.2GHz 196GB Ada Lovelace (2023) 24GB (GDDR6) x4 7,424 / 232 485 (total) TPU v3-8 Skylake-SP 96 2.0GHz ~330GB TPU v3 16GB (HBM) x8 - 125 (total)
12. 実践編 12 • 複素数平⾯上の集合判定の⾼速化 ◦ 普通にPythonで実装すると... → 4096 x 4096点 で 351.9 sec (遅い) def mandelbrot_kernel(c): z = c for i in range(max_iter): # max_iter = 500 z = z * z + c # zが閾値を超えたら終了 if abs(z) > 2: return i return max_iter def compute_mandelbrot(image): dx = (xmax - xmin) / width dy = (ymax - ymin) / height # ピクセルごとに複素数を計算 for j in range(height): for i in range(width): y = ymin + j * dy x = xmin + i * dx image[j][i] = mandelbrot_kernel(complex(x, y)) return image Mandelbrot 集合
13. ⾼速化① JITコンパイル 13 from numba import njit @njit(cache=True, fastmatch=True) def mandelbrot_kernel(c): z = c for i in range(max_iter): z = z * z + c # zが閾値を超えたら終了 if abs(z) > 2: return i return max_iter @njit(cache=True, fastmatch=True) def compute_mandelbrot(image): dx = (xmax - xmin) / width dy = (ymax - ymin) / height ⋮ • Just-In-Timeコンパイル ◦ プログラム実⾏中に機械語に変換 してコードを最適化する仕組み ◦ PyPy, CPython 3.13 (実験的) • Numba ◦ Pythonの数値計算を⾼速化するた めのライブラリ ◦ デコレータだけでほとんど実装 コストなく追加できる → 8.216 sec (x42.8倍)
14. ⾼速化② Python → C++ 14 • C/C++はPythonより2桁速い ◦ Kaggle notebookでも使える ◦ ファイルに書き出し→コンパイル ◦ subprocess, pybind11 inline int mandelbrot_kernel(float cx, float cy, int maxIter) { double zr = cx; double zi = cy; for (int i = 0; i < maxIter; ++i) { double zr2 = zr * zr - zi * zi + cx; double zi2 = 2.0 * zr * zi + cy; zr = zr2; zi = zi2; if (zr * zr + zi * zi > 4.0) { return i; } } return maxIter; } void compute_mandelbrot(std::vector<int>& image, double xmin, double xmax, double ymin, double ymax, int w, int h, int maxIter) { const double dx = (xmax - xmin) / static_cast<double>(w); const double dy = (ymax - ymin) / static_cast<double>(h); for (int j = 0; j < h; ++j) { const double cy = ymin + static_cast<double>(j) * dy; for (int i = 0; i < w; ++i) { const double cx = xmin + static_cast<double>(i) * dx; image[j * w + i] = mandelbrot_kernel(cx, cy, maxIter); } } } → 7.351 sec (x47.9倍)
15. 役に⽴った例(2)（C/C++で実装） 15 🥇 • PythonからC/C++のバイナリを呼ぶ ◦ チェスAIは思考が⾼速な⽅が強い（2倍で勝率+10%） ◦ メダル圏はほぼ誰もPythonを書いていなかった リソース制限下で効率的なチェスAI の強さを競う。 提出サイズ <64KB, メモリ <5MB という組み込み並みの制約がある ⇦ C ⇦ C ⇦ C ⇦ C++ ⇦ C++
16. Step3: 並列処理を知る 16 • どのように分割するか ◦ タスク並列: 異なる処理を同時に実⾏ ◦ データ並列: 同じ処理を⼀括適⽤ ▪ SIMD命令 ▪ GPU/TPUスレッド 問題を分割し、複数のコアや演算器を活⽤する thebeardsage.com/cuda-dimensions-mapping-and-indexing/ 分割に失敗した状態 GPU内のGRID/Block/Threadの構成
17. ⾼速化③ マルチスレッド（ループ並列） 17 • OpenMP ◦ C/C++向けに標準化された並列 プログラミングフレームワーク ◦ ループ内の処理を簡単に並列化 ◦ vCPU=4 → 4056 ms（x86.8 倍） ◦ vCPU=96 (TPU v3-8) → 235.2 ms（x1500 倍） #include <omp.h> … void compute_mandelbrot(std::vector<int>& image, double xmin, double xmax, double ymin, double ymax, int w, int h, int maxIter) { const double dx = (xmax - xmin) / static_cast<double>(w); const double dy = (ymax - ymin) / static_cast<double>(h); #pragma omp parallel for collapse(2) for (int j = 0; j < h; ++j) { const double cy = ymin + static_cast<double>(j) * dy; for (int i = 0; i < w; ++i) { const double cx = xmin + static_cast<double>(i) * dx; image[j * w + i] = mandelbrot_kernel(cx, cy, maxIter); } } }
18. ⾼速化④ SIMD命令 18 • SIMD命令 ◦ Single Instruction, Multiple Data ◦ 専⽤レジスタに複数の値を詰め込 み、加算‧乗算などを⼀括実⾏ ◦ AVX2: Haswell (2013)以降 ◦ AVX-512: Skylake-SP (2017) 以降 #include <immintrin.h> … static inline void mandelbrot_kernel_avx2(float cx_base, float cy, float dx, int* out_iters) { const __m256 idx = _mm256_set_ps(7,6,5,4,3,2,1,0); const __m256 vdx = _mm256_set1_ps(dx); const __m256 vcb = _mm256_set1_ps(cx_base); const __m256 cx = _mm256_fmadd_ps(idx, vdx, vcb); // cx_base + idx*dx const __m256 cyv = _mm256_set1_ps(cy); const __m256 four = _mm256_set1_ps(4.0f); __m256 zr = cx; __m256 zi = cyv; __m256i iters = _mm256_setzero_si256(); for (int it = 0; it < max_iter; ++it) { __m256 zr2 = _mm256_mul_ps(zr, zr); __m256 zi2 = _mm256_mul_ps(zi, zi); __m256 r2 = _mm256_add_ps(zr2, zi2); __m256 active = _mm256_cmp_ps(r2, four, _CMP_LE_OQ); int mask = _mm256_movemask_ps(active); if (mask == 0) break; // iters += active ? 1 : 0 const __m256i one = _mm256_set1_epi32(1); iters = _mm256_add_epi32(iters, _mm256_and_si256(one, _mm256_castps_si256(active))); … codingame.com/playgrounds/283/sse-avx-vectorization/what-is-sse-and-avx → 69.18ms （x5090 倍）※vCPU=96
19. ☑ 本⽇のゴール Pythonの処理を 速くする⽅法を知る 5000倍 達 成 の
20. ☑ 本⽇のゴール Pythonの処理を 速くする⽅法を知る 5000倍 の 「ドラゴンボール」第25巻, 集英社
21. ☑ 本⽇のゴール Pythonの処理を 速くする⽅法を知る 180000倍 の 5000倍
22. ⾼速化⑤ GPU/TPUを利⽤する 22 • JAX ◦ numpy⾵のインターフェースで XLAを使って⾼速化 ◦ ⾃動ベクトル化, GPU/TPU並列化 今こそはじめるJAX/Flax⼊⾨ Part 1 # pip install jax[cuda12] # pip install jax[tpu] import jax import jax.numpy as jnp … @jax.jit def compute_mandelbrot(xmin, xmax, ymin, ymax, width, height, max_iter): xs = jnp.linspace(xmin, xmax, width, dtype=jnp.float32) ys = jnp.linspace(ymin, ymax, height, dtype=jnp.float32) X, Y = jnp.meshgrid(xs, ys, indexing="xy") C = X + 1j * Y Z = C iters = jnp.zeros((height, width), dtype=jnp.int32) … ◦ T4 x2 → 865.4 ms（x407 倍） ◦ TPU v3-8 → 115.9 ms（x3040 倍）
23. 役に⽴った例(3)（JAXで⾼速化） 23 • Forward modelを利⽤して精度を上げる ◦ 地震波形データからConvNeXt + Unetなどで学習 → 学習データが多すぎて1週間以上かかる… ◦ Forward modelのコードをJAXに変換し⾼速化 ◦ 結果を直接最適化し、スコアを⼤きく向上 地震波形データから地下構造（⾳響 速度分布）を求める課題。 地下構造→波形データ（Forward） の実装は公開されている 🥈 地震波形データ 地下速度分布 Forward Inversion（課題） 邪道（？）: 勾配を取れば 直接改善できる! 正攻法：画像モデルの学習
24. ⾼速化⑥ CUDAを書く 24 • CUDA ◦ Compute Unified Device Architecture ◦ GPGPUのための並列計算プラット フォーム/プログラミングモデル ◦ C/C++拡張として書く ◦ GPUを数千~数万スレッド単位で 同時実⾏できる ◦ nvccでコンパイルして実⾏ ⼈類の99%はバージョン確認にしか使っていないが、 nvccは本来はCUDAコンパイラ → 1.894 ms (x186000 倍) ※ L4x4で⾮同期&最適化 #include <cuda_runtime.h> … __global__ void mandelbrot_kernel_u16( uint16_t* __restrict__ out, int width, int hSeg, int yOffset, float xmin, float ymin, float dx, float dy, int max_iter ) { const int x = blockDim.x * blockIdx.x + threadIdx.x; const int yLocal = blockDim.y * blockIdx.y + threadIdx.y; const bool inBounds = (x < width) & (yLocal < hSeg); const int y = yOffset + yLocal; // 画素の複素数 c = (cx, cy) float cx = xmin + x * dx; float cy = ymin + y * dy; … // kernel実行とhostへの非同期コピー CK(cudaEventRecord(g[i].evStart, g[i].stream)); mandelbrot_kernel_u16<<<grid, block, 0, g[i].stream>>>( g[i].d_out, width, g[i].hSeg, g[i].yOffset, xmin, ymin, dx, dy, max_iter ); CK(cudaEventRecord(g[i].evStop, g[i].stream)); CK(cudaMemcpyAsync( h_out + static_cast<size_t>(g[i].yOffset) * width, g[i].d_out, sizeof(uint16_t) * width * g[i].hSeg, cudaMemcpyDeviceToHost, g[i].stream ));
25. まとめ 25 • ⾼速化‧並列化をするには… ◦ アルゴリズムを知る ◦ 実⾏速度を知る ◦ 並列化を知る • ⾼速化すると何がうれしい？ ◦ 前処理‧データ⽣成‧学習‧推論に 余裕を⽣むことで、コンペ本来の 試⾏錯誤に集中することができる • JAX, C++, CUDAを書いて ライバルと差をつけましょう！ ⼿法 実⾏時間 速度⽐ [CPU] Baseline 351.9 sec x1 Numba 8.216 sec x 42.8 C++ 7.351 sec x 47.9 OpenMP 235.2 ms x 1496 SIMD 69.18 ms x 5086 [GPU / TPU] JAX (T4x2) 865.4 ms x 406.6 JAX (TPU v3-8) 115.9 ms x 3052 CUDA (T4 x2) 4.437 ms x 79310 CUDA (L4 x4) 1.894 ms x 185800