画像コンペで苦労したであろう点をまとめる
はじめに
2022年に参加した画像分類のコンペティションでは、PyTorch Distributed を活用した分散学習 を駆使し、大規模なモデルをトレーニングしました。特に、超高解像度画像を扱う Mayo Clinic STRIP AI や、大量のクジラ・イルカ個体識別を行う Happywhale コンペでは、単一GPUでは学習が不可能なレベルのデータ量とモデルサイズ に直面しました。
本記事では、
- どのような課題に苦労したのか
- どれくらいのモデルサイズになったのか
- どのようなハードウェア構成で分散学習を行ったのか
- どれくらいの予算で構築したのか
を中心に、体験記としてまとめます。
苦労した点
1. 画像サイズが大きすぎてメモリ不足に
Mayo Clinic STRIP AI のデータは ギガピクセル級の病理スライド画像 であり、1枚あたり数GB のサイズがありました。通常の単一GPUではバッチサイズ1すら難しい 状況で、以下の問題が発生しました。
- メモリオーバーフロー（CUDA out of memory エラー）
- データロードが遅すぎる（I/O ボトルネック）
- 高解像度のままではGPUメモリに収まらない
→ 対策
- タイル分割（Multiple Instance Learning, MIL） を導入
- FP16（混合精度学習） でメモリ消費を抑える
- 分散データローディング（num_workers を増やし、前処理を高速化）
このような工夫をすることで、4GPU構成でバッチサイズ16を確保 できるようになり、学習の安定性が向上しました。
2. Happywhale のデータがクラス数多すぎ問題
Happywhale では 15,000以上の個体（クラス） を分類する必要がありました。通常のソフトマックス分類では計算負荷が高すぎるため、ArcFace（角度マージン分類） を活用しましたが、以下の課題が発生。
- ラベルの不均衡問題（ほとんどのクラスが数枚しかない）
- 学習データが多すぎて 1 エポックが終わらない
- メモリ消費が多く、1GPU では学習できない
→ 対策
- クラスごとのマージン調整（Dynamic Margin ArcFace）
- 分散学習（4GPUでデータを並列処理）
- kNNによる補助識別を追加し、新規個体の識別精度を向上
結果として、最終モデルでは 約0.874のMAP@5（上位2%） を達成。
モデルサイズと学習環境
最終的なモデルサイズ
コンペ終盤では、ConvNeXt Large や EfficientNetV2-L を組み合わせたアンサンブルを作成し、
- 1モデルあたり約250Mパラメータ
- アンサンブル全体で約800Mパラメータ（4モデル）
という規模になりました。これにより、単一GPUではメモリが足りず、分散学習が必須 となりました。
GPU構成とコスト
最適なコストで最大の計算力を確保するため、
- NVIDIA RTX 3090 × 4枚（合計48GB VRAM）
- Intel Xeon W-2245（8コア/16スレッド）
- メモリ128GB（I/O負荷軽減のため）
- NVMe SSD 2TB（データロード高速化）
を使用しました。価格は以下の通り：
| 機材 | 価格（円） |
|---|---|
| RTX 3090 × 4 | 約180万円 |
| Xeon W-2245 | 約15万円 |
| メモリ 128GB | 約12万円 |
| NVMe SSD 2TB | 約5万円 |
| 電源・冷却系統 | 約8万円 |
| 合計 | 約220万円 |
当初はA100を検討しましたが、1枚あたり100万円超え でコストオーバーとなるため、RTX 3090 4枚構成でコストパフォーマンスを最適化 しました。
分散学習の実装
PyTorch の DistributedDataParallel (DDP) を活用し、
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl")
device = torch.device("cuda")
model = ConvNeXtLarge().to(device)
model = DDP(model, device_ids=[dist.get_rank()])
このようにコードを記述し、各GPUにデータを分散。
- 1GPUあたりのバッチサイズを増やす（学習を安定化）
- 学習時間を単体GPUの約3.5倍に短縮（1エポック10時間 → 3時間）
- GPU間の同期を最適化し、分散学習の効率を向上
を実現しました。
結果と感想
- 分散学習なしでは計算が間に合わなかった（特に Happywhale のデータ量が大きすぎる）
- RTX 3090 4枚で十分な性能を発揮（A100ほどの性能は不要だった）
- GPUコストは約220万円で収まり、500万円以下で最適な環境を構築
結果として、
- Mayo Clinic STRIP AI：0.83 AUC（上位5%）
- Happywhale：0.874 MAP@5（上位2%）
という成果を達成。分散学習環境の構築と最適化が、競争力のあるスコア獲得に直結した ことを実感しました。
これから画像分類コンペに挑戦する方は、モデル選定だけでなく、計算環境の設計も重要 であることを意識すると良いでしょう！
Discussion