0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@okai

数式認識VLM「PP-FormulaNet_plus-M」をONNXに変換してonnxruntimeで動かす

0
Posted at

はじめに

論文をRAG化してMCP経由でAIに読ませることができるMacOSネイティブな文献管理ソフトを開発中です（Zoteroでできるけど物足りなかった）。配布を考えた時にPDFから数式をLaTeXに起こすパイプラインを.appバンドル内に全てまとめるため、依存を極力減らしたいという背景がありました。
数式のLaTex形式での抽出は論文PDFを機械可読にする用途では避けて通れない工程です。PaddleOCR に入っている PP-FormulaNet_plus-M は、このタスクでは品質と速度のバランスがよく、しかもApache 2.0ライセンスということで目をつけました。でもPaddle専用モデルなので今回の使用目的ではあまり好ましくありませんでした。

PaddlePaddle のランタイムは重く（モデル本体の 2〜3 倍のメモリを使う）、Python からしか呼べません。アプリに組み込むなら ONNX にして onnxruntime から呼びたいところです。ところが公式の変換ツール paddle2onnx にそのまま通すと、変換は成功するのに実行できない壊れたONNXが出てきます。

本記事では、この壊れた時の中身と、backbone / head の分割エクスポートによる回避策、変換後の実測値をまとめます。

概要

  • PP-FormulaNet_plus-M の推論グラフを paddle2onnx で一発変換すると、Loop 内の If 分岐が rank 不整合を起こし、onnxruntime が実行時に落ちる
  • 回避策は「生成ループを ONNX に入れない」こと。backbone と 1 ステップぶんの head に分割してエクスポートし、生成ループは呼び出し側で書く
  • 分割エクスポートは FahNos/pp_formula_to_onnx の手順が使える。エクスポート後に float64 → float32 の修正が必要
  • 変換後は CPU で 約 1.05 秒/式。PaddleOCR 経由（3.9 秒/式)の約 3.7 倍速でした

背景

PP-FormulaNet_plus-M は「数式画像 → LaTeX」のエンコーダ・デコーダ型モデルです。画像エンコーダが特徴 (1, 144, 2048) を作り、自己回帰デコーダが語彙 50,000 からトークンを 1 つずつ選びます。

問題はこの自己回帰の生成ループがモデルの推論グラフの内側にあることです。Paddle の推論グラフをそのまま ONNX に変換すると、生成ループが ONNX の Loop オペレータに落ち、その内部の分岐グラフで型宣言の不整合が生まれます。逆に言えば、ループさえグラフから追い出せば変換は素直に通るはずです。

なお筆者の動機はすでに述べたことに加えて、docling（IBM の PDF 解析パイプライン）の数式エンリッチメントの置き換えという目的もあります。標準の CodeFormulaV2 は手元の Apple Silicon で 31 秒/式かかり、入力画像に存在しない本文断片を数式の後ろに生成する過剰生成（自己回帰デコードで入力に関係しないトークン列が生成される）も多いため、代替として PP-FormulaNet_plus-M を選定しました。サイズ違いのS/M/L のうち S は % を含む式で出力が停止する問題があり、M はコスパが良さそうだったので採用しました。Lも本稿の方法で動きますが精度はたいして変わらないけれど処理速度は倍くらい遅いという結果でした。

検証環境

項目 値
ハードウェア Apple M4 (メモリ 32 GB)
OS macOS 26.5.2
Python 3.12.12
paddlepaddle 3.1.0
paddle2onnx 2.1.0
onnx 1.17.0
onnxruntime 1.29.0
numpy 1.26.4

Python は 3.12 を使いました。3.13 には paddle2onnx の新しい wheel が無く、pip が Paddle 2.x 時代の 0.8.1 を使って失敗します（手順内で詳述）。

一発変換が失敗するまで

まず素直にやった場合の失敗の記録です。PaddleOCR が ~/.paddlex/official_models/PP-FormulaNet_plus-M/ に落とす推論グラフ（inference.json + inference.pdiparams）を paddle2onnx に渡すと、変換自体は警告を出すものの完了し、720 MB の model.onnx が得られます。

これを onnxruntime にロードして推論すると、次のエラーで実行不能になりました。

Non-zero status code returned while running Loop node.
Name:'Loop.0' ... output shape mismatch:
declared {} but computed {1}

原因を探ると、Loop.0 > If.7 > Identity.1702 の出力 p2o.yield.13 が rank 0（スカラー）と宣言されているのに、then 分岐は実行時に rank 1 のテンソルを作っているためでした。onnxruntime は宣言どおりに出力バッファを事前確保するので、ここで整合が取れません。

対策と結果は次のとおりです。

試行内容 結果
opset を 16 / 17 / 19 に変える すべて同じ挙動で失敗
モデルを S に変える 同上（モデル固有ではなく変換系の問題）
ORT のグラフ最適化を無効化 回避できず
分岐に Unsqueeze を挿入して rank 1 に揃えるグラフ書き換え If.7 は直るが、別の箇所で同種の不整合が出てだめ
旧版 paddle2onnx 0.8.1 で変換 Paddle 2.x 用（paddle.fluid を要求）。Python 3.13 には Paddle 2.x の wheel が無く環境が作れない
onnx2torch で PyTorch に戻して再エクスポート Pad（opset 18）が未実装で変換不可

グラフ書き換えは 1 箇所ならなんとかなるかもしれませんが、Loop の中に同種の不整合が複数あるため難しいと判断しました。

手順

同じ問題を解いているリポジトリ FahNos/pp_formula_to_onnx が存在しました。推論グラフではなく学習コード側から、backbone と「1 ステップぶんの head」を別々に静的グラフ化してエクスポートする手順を公開しています。生成ループがグラフに入らないので、Loop も If も生まれず、rank 不整合が起きようがありません。以下は筆者が実際にうまくいった手順です。

環境構築: 変換専用 venv（Python 3.12）

python3.12 -m venv pp2onnx-venv

pp2onnx-venv/bin/pip install --upgrade pip

pp2onnx-venv/bin/pip install packaging onnx

pp2onnx-venv/bin/pip install paddle2onnx paddlepaddle

注意点が 2 つあります。packaging と onnx を先に入れないと、paddle2onnx の import が ModuleNotFoundError: packaging で落ちます。また Python 3.13 では paddle2onnx の新しい wheel が無く、pip が古い 0.8.1（Paddle 2.x 用）を呼んでしまうので、3.12 が必須です。

この時点では paddlepaddle 3.3.1 が入りますが、Step 2 でリポジトリの requirements により 3.1.0 へ下がります（最終構成が検証環境の表）。

Step 1: リポジトリのクローンと事前学習済み重みの取得

git clone --depth 1 https://github.com/FahNos/pp_formula_to_onnx ppf2onnx
cd ppf2onnx

mkdir -p pretrained_model
curl -sL -o pretrained_model/PP-FormulaNet_plus-M_pretrained.pdparams \
  https://paddle-model-ecology.bj.bcebos.com/paddlex/official_pretrained_model/PP-FormulaNet_plus-M_pretrained.pdparams

学習コードからエクスポートするので、必要なのは PaddleOCR が配布する推論グラフではなく事前学習済みの .pdparams（617 Mバイト）です。設定ファイル configs/rec/PP-FormuaNet/PP-FormulaNet_plus-M_ONNX.yaml の Global.pretrained_model がこの場所を指しているので置くだけで済みます。モデル名を引数で渡す仕組みはありませんでした。ディレクトリ名 PP-FormuaNet もともとのタイポだと思いますがそのままでよいです。

Step 2: 依存の追加インストール

pp2onnx-venv/bin/pip install -r requirements_onnx.txt

pp2onnx-venv/bin/pip install scikit-image albumentations "numpy<2.0"

2 行目は README にありませんが必要です。requirements_onnx.txt に skimage と albumentations が入っていないのに ppocr.data が import するため、入れないと Step 3 が ModuleNotFoundError で 2 回止まります。

Step 3: 静的グラフへのエクスポート

pp2onnx-venv/bin/python3 tools/export_onnx.py \
  --config ./configs/rec/PP-FormuaNet/PP-FormulaNet_plus-M_ONNX.yaml

確認

ls -la output/static_models/

backbone_static.json         backbone_static.pdiparams   (計 276,742,350 B)
head_simplified_static.json  head_simplified_static.pdiparams  (計 315,152,607 B)

このスクリプトは静的グラフ化のあと、内部で paddle2onnx コマンドをサブプロセスとして PATH から呼び出しますがこれがやっかいなところでした。venv を activate せずに venv/bin/python3 を直接叩いていると、静的グラフの出力までは成功するのに ONNX 化だけが [Errno 2] No such file or directory: 'paddle2onnx' で失敗し、ログ末尾が Backbone ONNX: None / Head ONNX: None になります。venv を activate して実行するか、次の Step 4 のように paddle2onnx を手で実行してください（静的グラフは出来ているためやり直しは不要のはず）。

Step 4: paddle2onnx で ONNX へ

mkdir -p output/onnx_models

pp2onnx-venv/bin/paddle2onnx \
  --model_dir ./output/static_models \
  --model_filename backbone_static.json \
  --params_filename backbone_static.pdiparams \
  --save_file ./output/onnx_models/backbone.onnx \
  --opset_version 17

pp2onnx-venv/bin/paddle2onnx \
  --model_dir ./output/static_models \
  --model_filename head_simplified_static.json \
  --params_filename head_simplified_static.pdiparams \
  --save_file ./output/onnx_models/head.onnx \
  --opset_version 17

head 側で Fail to fold onnx model with error: [ShapeInferenceError] (op_type:MatMul, ...) Skip folding. という WARNING が出ますが次の手順で直りますので無視でOKです。

得られる成果物は 2 つ。

成果物 サイズ 入出力
backbone.onnx 277 MB 画像 (N, 3, H, W) float32 → 特徴 (1, 144, 2048)
head.onnx 418 MB トークン列 (N, T) int64 + 特徴 → logits (N, 50000)

Step 5: head の float64 → float32 修正

エクスポート直後の head はグラフ内に float64 が混ざっており、そのままでは onnxruntime でロードできません。同リポジトリの fix_head_onnx.py でこれを修正します。引数はなく、./output/onnx_models/head.onnx → head_fixed.onnx とパスが決め打ちです。

pp2onnx-venv/bin/python3 tools/fix_head_onnx.py

確認

ログはベトナム語で出ます（作者の母語？）。手元では 91 箇所が書き換わりました。

Tổng số lỗi đã sửa trên toàn bộ các đồ thị: 91   ← 全グラフでの修正総数: 91
Xác thực THÀNH CÔNG!                             ← 検証成功

中身は float64 → float32 の一括変換で、initializer・Cast の to 属性・Constant の value・入出力と value_info の elem_type を再帰的に直します。最後に onnx.checker と onnxruntime のロードテストまで走ってくれます。ファイルサイズの差はわずか 52 バイトで、型宣言の書き換えが主であると思われます。

Step 6: onnxruntime でロード確認

確認

pp2onnx-venv/bin/python -c "
import onnxruntime as ort
for f in ['output/onnx_models/backbone.onnx', 'output/onnx_models/head_fixed.onnx']:
    s = ort.InferenceSession(f, providers=['CPUExecutionProvider'])
    print(f)
    for i in s.get_inputs():  print('  in :', i.name, i.shape, i.type)
    for o in s.get_outputs(): print('  out:', o.name, o.shape, o.type)
"

output/onnx_models/backbone.onnx
  in : image ['N', 3, 'H', 'W'] tensor(float)
  out: fetch_name_0 [...] tensor(float)
output/onnx_models/head_fixed.onnx
  in : decoder_input_ids ['N', 'T'] tensor(int64)
  in : encoder_outputs ['N', 'S', 2048] tensor(float)
  out: fetch_name_0 ['N', 50000] tensor(float)

未使用 initializer の警告が大量に出ますが実行には影響しません。

Step 7: 生成ループを自前で書く

ループをグラフから追い出したので、貪欲法デコードは呼び出し側で行う必要があります。実行サンプル↓

import numpy as np
import onnxruntime as ort

bb   = ort.InferenceSession("backbone.onnx",   providers=["CPUExecutionProvider"])
head = ort.InferenceSession("head_fixed.onnx", providers=["CPUExecutionProvider"])

feat = bb.run(None, {"image": image})[0]          # (1, 144, 2048)  一度だけ
tokens = [BOS]
for _ in range(MAX_STEPS):
    ids = np.array([tokens], dtype=np.int64)      # (1, T)
    logits = head.run(None, {"decoder_input_ids": ids,
                             "encoder_outputs": feat})[0]
    nxt = int(logits[-1].argmax())
    if nxt == EOS:
        break
    tokens.append(nxt)

head は KV キャッシュを持たないため、毎ステップ全プレフィックスを渡し直す形になります。それでも 1 ステップ 6.3 ms（後述）なので、実用上は問題になりませんでした。

実測

Apple M4 の CPU（onnxruntime CPUExecutionProvider）での実測です。

項目 値
backbone 1 回 153.4 ms
head 1 ステップ 6.3 ms
1 式あたり（平均 142 ステップとして） 約 1.05 秒
推論時 RSS 2,232 MB

PaddleOCR 経由の 3.9 秒/式に対して約 3.7 倍速です。おそらく Paddle ランタイムのオーバーヘッドが消えたからでしょう。モデル自体は同じものです。

なお Core ML への変換（ANE / GPU 実行）も試しましたが、変換した mlpackage が Failed to parse the model specification でコンパイルできず断念しました。一部の演算子が Core ML の仕様に適合しないようです。仮に動いたとしても、同系統の小型モデルで測った範囲ではこの種の自己回帰デコーダは 1 ステップの計算が小さすぎてアクセラレータの起動・転送コストが勝り、CPU が一番早かったです。ANE使ったら爆速なんじゃねと思っていましたがCPU 駆動で落ち着きました。GPUはそもそもCUDA必須のようだし。

前処理について

ONNX にできても前処理を移植しないと使えないので、要点だけ記しておきます。PP-FormulaNet_plus-M の前処理は UniMERNet 系で、次の順に行います。

  1. 余白の除去（画像の平均輝度が 128 を超える場合は白黒反転）
  2. 32 の倍数へパディング
  3. 32×32 〜 672×192 の範囲に収める
  4. 正規化（出力テンソルの値域はおおむね −4.6 〜 1.2)
  5. 再び 32 の倍数へパディング

出力は (1, 1, 384, 384) の float32 ですが、backbone の入力は 3 チャンネルなので同じ面を複製して渡します。

結論

PP-FormulaNet_plus-M の ONNX 変換について、要点は次のとおりです。

  • 推論グラフの一発変換は、Loop 内の rank 不整合により onnxruntime で実行できない（opset 16/17/19、S/M の両モデルで再現）
  • 生成ループをグラフに入れない分割エクスポート（backbone + 1 ステップ head）で回避できる。手順は FahNos/pp_formula_to_onnx を参考にしました
  • head は float64 → float32 の修正が必要
  • 変換後は Apple M4 の CPU で約 1.05 秒/式、PaddleOCR 経由の約 3.7 倍速

グラフ内の生成ループは paddle2onnx に限らず変換系の落とし穴になりがちなようなので、ループを呼び出し側で書く方針は他のエンコーダ・デコーダ型モデルにも使えるかもです。

長時間運用での安定性、および他のバージョン構成・他プラットフォームでの再現性については検証していません。また変換後モデルと Paddle 原本の出力一致度についても未検証です。

参考

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?