WebGPUでMiDaS深度推定を10倍高速化した
CPUで画像の深度推定をしていた。遅すぎた。
midori261で実装したMiDaSモデルは、CPUのみで処理していた。512×512pxの画像で10秒以上。実用的じゃない。
WebGPUに書き直したら、同じ画像が1秒以下になった。10倍の速度。GPUの並列計算の威力。
なぜ作ったか
midori261のCPU版を使ってみて、処理時間が気になった。
画像1枚で10秒以上待たされる。複数枚処理したら、分単位で待つことになる。これは使えない。
写真展示でこのアプリを使おうと思っていたので、レスポンスは重要でした。来場者が10秒も待つのは、現実的じゃない。
GPUを使えば速くなるのは分かっていた。でも、WebGPUがonnxruntime-webで使えるのか、確信がなかった。
ドキュメントを読んで、executionProvidersに"webgpu"を指定すれば動くと知った。試すことにした。
CPU vs WebGPU
最初にベンチマークを取った。
CPU処理(midori261):
- 512×512px: 約12秒
- 384×384px: 約8秒
- 1024×1024px: 約45秒
これだと、大きな画像は待ちきれない。
WebGPUに変更したら、劇的に改善した。
WebGPU処理(midori262):
- 512×512px: 約1.2秒
- 384×384px: 約0.8秒
- 1024×1024px: 約4.5秒
10倍速い。これなら実用的だ。
CPUは12秒、WebGPUは1.2秒。同じ処理で10倍の差。
WebGPUとは
WebGPUは、ブラウザでGPUを直接使える新しいAPI。WebGLの後継。
WebGLよりも低レベルで、GPUの性能を引き出しやすい。Metal、Vulkan、DirectX 12のような最新のGPU APIに近い設計。
onnxruntime-webは、1.15.0からWebGPUに対応。executionProvidersに"webgpu"を指定するだけで、GPUで推論できる。
const session = await ort.InferenceSession.create(
modelPath,
{ executionProviders: ["webgpu"] }
);
これだけ。CPUの時と比べて、コードはほぼ同じ。
実装の詳細
WebGPUに移行する際、コードの変更は最小限でした。
executionProvidersの指定
async function loadModel() {
if (!session) {
session = await ort.InferenceSession.create(
modelPath,
{ executionProviders: ["webgpu"] }
);
}
return session;
}
ここだけ変更。"cpu"から"webgpu"に変えた。
でも、最初はエラーが出た。WebGPUが使えない環境では、フォールバックする必要がある。
try {
session = await ort.InferenceSession.create(
modelPath,
{ executionProviders: ["webgpu"] }
);
} catch (error) {
console.warn('WebGPU使用不可、CPUにフォールバック:', error);
session = await ort.InferenceSession.create(
modelPath,
{ executionProviders: ["cpu"] }
);
}
これで、WebGPU非対応ブラウザでもCPUで動く。
Tensorの作成
const tensor = new ort.Tensor(
"float32",
floatData,
[1, 3, inferenceHeight, inferenceWidth]
);
Tensorの作成は、CPUの時と同じ。
onnxruntime-webが内部でGPUメモリに転送してくれる。手動でGPUバッファを管理する必要はない。これは楽だった。
WebGLでシェーダーを書いていた時は、バッファの転送を全部自分で管理していた。それに比べると、onnxruntime-webはかなり抽象化されている。
推論実行
const startPerf = performance.now();
const feeds = { "input": tensor };
const results = await session.run(feeds);
const inferenceDuration = performance.now() - startPerf;
推論も同じ。session.run()を呼ぶだけ。
内部でGPUが並列計算してくれる。CPUと違って、複数の演算を同時に実行できる。これが速度の秘密。
9種類のモデル対応
MiDaSには複数のモデルがある。精度と速度のトレードオフ。
実装したモデル:
- midas_model_256.onnx(256×256px)
- midas_model_512.onnx(512×512px)
- midas_model_1024.onnx(1024×1024px)
- midas_model_DPT_Hybrid_384.onnx(384×384px)
- midas_model_DPT_Hybrid_512.onnx(512×512px)
- midas_model_DPT_Hybrid_1024.onnx(1024×1024px)
- midas_model_DPT_Large_384.onnx(384×384px)
- midas_model_DPT_Large_512.onnx(512×512px)
- midas_model_DPT_Large_1024.onnx(1024×1024px)
ユーザーがドロップダウンで選択できるようにした。
DPT_Largeが最も高精度だけど、処理時間が長い。1024pxで約4.5秒。
シンプルな256pxモデルは、0.5秒以下。でも、精度は落ちる。
用途に応じて使い分けられる。
深度マップの可視化
深度値をそのままグレースケールで表示すると、見にくい。
Infernoカラーマップを使った。深い部分は暗い青、浅い部分は明るい黄色。
function getInfernoColor(t) {
const stops = [
{ t: 0.0, color: [0, 0, 3] },
{ t: 0.25, color: [87, 15, 109] },
{ t: 0.5, color: [187, 55, 84] },
{ t: 0.75, color: [249, 142, 8] },
{ t: 1.0, color: [252, 255, 164] }
];
// 線形補間で色を計算
for (let i = 0; i < stops.length - 1; i++) {
const s0 = stops[i], s1 = stops[i + 1];
if (t >= s0.t && t <= s1.t) {
const factor = (t - s0.t) / (s1.t - s0.t);
const r = s0.color[0] + factor * (s1.color[0] - s0.color[0]);
const g = s0.color[1] + factor * (s1.color[1] - s0.color[1]);
const b = s0.color[2] + factor * (s1.color[2] - s0.color[2]);
return [Math.round(r), Math.round(g), Math.round(b)];
}
}
return [0, 0, 0];
}
深度値を0~1に正規化してから、Infernoカラーマップで色を取得。
結果、深度の違いが直感的に分かるようになった。遠くの山は青、手前の木は黄色。
オーバーレイ表示
元の画像と深度マップを重ねて表示する機能も追加した。
overlayCtx.drawImage(imageElement, 0, 0, originalWidth, originalHeight);
overlayCtx.globalAlpha = 0.4;
overlayCtx.drawImage(depthCanvas, 0, 0, originalWidth, originalHeight);
overlayCtx.globalAlpha = 1.0;
元画像を描画してから、深度マップを40%の透明度で重ねる。
これで、どの部分が手前で、どの部分が奥なのか、元の画像と照らし合わせて確認できる。
ハマったポイント
WebGPU非対応ブラウザ
最初、WebGPUが使えない環境での動作を考えていなかった。
Chromeの古いバージョンやSafariで試したら、エラーで止まった。
executionProvidersでフォールバックする仕組みを追加。WebGPU→CPUの順で試すようにした。
モデルの正規化方法の違い
MiDaSのモデルによって、前処理の正規化パラメータが違う。
512pxモデルは mean=[0.485, 0.456, 0.406]、DPT系は mean=[0.5, 0.5, 0.5]。
最初、全モデル同じパラメータで処理していた。結果がおかしい。DPTモデルで真っ暗な深度マップが出力された。
モデル名で分岐して、正しいパラメータを使うように修正。
let mean, std;
if (["midas_model_512.onnx", "midas_model_1024.onnx"].includes(currentModel)) {
mean = [0.485, 0.456, 0.406];
std = [0.229, 0.224, 0.225];
} else {
mean = [0.5, 0.5, 0.5];
std = [0.5, 0.5, 0.5];
}
これで、全モデルが正しく動くようになった。
メモリ管理
最初、複数の画像を連続で処理すると、ブラウザが重くなった。
ONNXSessionをキャッシュしていたけど、古いTensorがメモリに残っていた。
明示的にsession.run()の後、古いTensorを破棄する処理は不要だった。JavaScriptのGCに任せれば良い。
でも、大量の画像を処理する場合は、一定数ごとにsessionを再作成した方が安定する。
パフォーマンス測定
処理時間をリアルタイムで表示する機能を追加した。
const startPerf = performance.now();
const results = await session.run(feeds);
const inferenceDuration = performance.now() - startPerf;
timingDiv.textContent = `推論時間: ${inferenceDuration.toFixed(1)} ミリ秒`;
performance.now()で高精度なタイマーを使用。ミリ秒単位で測定できる。
ユーザーが体感速度を確認できるようにした。
使ってみて
実際にWebGPU版を使ってみて、レスポンスの改善を実感できた。
画像をドロップしてから結果が表示されるまで、1秒以下。これなら待たせない。
写真展示でも問題なく使える。来場者が複数の写真を試しても、ストレスなく動く。
ポイントは以下の3つ:
- executionProvidersに"webgpu"を指定するだけでGPU処理(コード変更最小限)
- 処理時間が10倍改善(12秒→1.2秒)
- 9種類のモデル選択で精度と速度のトレードオフを調整可能
WebGPUの威力を実感できた。onnxruntime-webが抽象化してくれているので、実装は簡単だった。
同じような深度推定を実装している方の参考になれば嬉しいです。
まとめ
今回は、MiDaS深度推定をWebGPUで高速化しました。
ポイントは以下の4つ:
- WebGPU対応のonnxruntime-webを使用(executionProvidersで"webgpu"指定)
- 処理時間を10倍改善(CPU: 12秒 → WebGPU: 1.2秒)
- 9種類のMiDaSモデルに対応(精度と速度のトレードオフ)
- Infernoカラーマップで深度を可視化(直感的に理解しやすい)
GPUの並列計算能力を活用することで、ブラウザでも実用的な深度推定が可能になった。
WebGPUはまだ新しい技術だけど、onnxruntime-webのおかげで簡単に使えた。
さらに深く学ぶには
この記事で興味を持った方におすすめのリンク:
自分の関連記事:
最後まで読んでくださり、ありがとうございました。