動的可変長幅入力とバケットバッチングでOCR文字認識を高速化する

MLism代表の木之下です。今回、YomiTokuの軽量モード（--lite）に、CPUでの文字認識をより高速に実行するための新しい仕組みを導入しました。本記事では、従来の軽量モデルが抱えていた課題と、それを解決するために実装した動的可変長幅入力、バケットバッチングについて紹介します。

従来の軽量モデルとその課題

YomiTokuの文字認識には、Transformerベースの文字認識モデルであるPARSeqを採用しています。PARSeqは自己回帰（AR）デコードを行うモデルです。行画像を入力として文字を1文字ずつ順番に予測し、終端記号（EOS）が出力された時点で生成を終了します。

この構成では、処理時間は主に次の2つの要素によって決まります。ひとつは入力画像の幅によって決まるエンコーダの計算量、もうひとつは予測する文字列長によって決まるARループの反復回数です。従来の軽量モデルでは、この2つを一律に短くすることで処理を高速化していました。具体的には、入力画像の横幅を通常モデルの半分に縮小し、予測可能な最大文字列長を50文字に制限していました。これにより、画像処理に必要な計算量を削減するとともに、ARデコードの反復回数にも上限を設けていました。

この方法はシンプルで効果的でしたが、1行あたり50文字までしか認識できないという別の課題が生じます。日本語の文書でも、1行が50文字を超えることは珍しくありません。英文では、さらに長い行が含まれることがあります。長い行が途中で切れてしまうため、「軽量モードは高速だが、文書によっては利用できない」という状態になっていました。

そこで今回は、文字列長の制限を緩和しながら、計算量を効率的に削減する方法を検討し、実装しました。

そもそも、どこに無駄があるのか

固定幅の入力画像とARデコードを組み合わせた従来の構成を詳しく見ていくと、主に2つの無駄があることが分かります。

固定幅パディングの無駄

認識モデルへの入力には、文書から切り出した行画像を使用します。従来は、画像の高さを32pxに揃えたうえで、固定サイズのキャンバスにリサイズとパディングを行っていました。通常モデルの場合、入力サイズは32×800pxです。

実際の文書では、すべての行が800px相当の長さを持っているわけではありません。むしろ、幅の狭い行が大半を占めます。例えば「〒277-8520」のような短い行も、一律に800px幅までパディングされます。その結果、キャンバスの大部分が背景となります。エンコーダは、この背景部分に対応するパッチについても計算を行います。文字が存在しない余白に対しても同じように計算が発生するため、多くの計算量が無駄になっていました。

課題1：固定長幅入力による余分なパディング。短い画像ほど余白が大きくなる

ARデコードがバッチ内の最長行に律速される

もうひとつの問題は、ARデコードの反復回数です。ARデコードはバッチ単位で実行され、バッチ内のすべての行がEOSに到達するまでループが続きます。そのため、バッチ内にひとつでも長い行が含まれていると、ほかの短い行がすでに生成を終えていても、最長の行に合わせてループが回り続けます。

課題2：ARの再帰予測回数がバッチ内の最大文字数に律速される

つまり、短い行は幅方向では固定幅のパディングによって、文字列長方向ではARループによって、長い行に計算量を合わせられていました。従来の軽量モデルで採用していた「画像幅と最大文字列長を一律に短くする」という設計は、こうした無駄を入力制限によって抑える方法だったとも言えます。

新しいアプローチ

今回の実装では、一律の制限によって計算量を抑えるのではなく、各行の実際の長さに応じて必要な計算量を割り当てる方針に変更しました。中心となる仕組みは、動的可変長幅入力とバケットバッチングの2つです。

動的可変長幅入力（dynamic-width）

まず、すべての行画像を固定の800px幅に合わせる処理をやめ、行画像の内容に応じてキャンバス幅を変えられるようにしました。具体的には、テキストが含まれる実際の幅に終端検知用の末尾マージンを加え、エンコーダのパッチ幅の倍数になるように切り上げます。短い行は狭いキャンバスのまま入力できるため、背景部分に対する不要なパッチ計算を削減できます。

解決策1：動的可変長幅入力によるパディング幅の縮小

エンコーダでは、学習済みの位置埋め込みを、実際のパッチグリッドの大きさに合わせて切り出します。これにより、入力幅が変化しても、そのまま推論できるようにしました。

実装にあたって注意が必要だったのが、画像末尾のマージンです。文字認識モデルは、テキストの後ろに続く背景領域も利用して「文字列がここで終わる」と判断しています。そのため、文字の直後まで画像を切り詰め、末尾のマージンをなくすと、存在しない文字列の続きを生成しやすくなります。そこで、入力画像の末尾には一定のマージンを残すようにしたうえで、可変幅の入力を前提としてモデル自体を再学習しました。これにより、認識精度を維持しながら、可変幅の入力を利用できるようになりました。

バケットバッチング（bucket batching）

動的可変長幅入力だけでは、バッチ内に幅の異なる画像が混在した場合、最も幅の広い画像に合わせてパディングする必要があります。同様に、ARデコードの再帰回数もバッチ内の最大文字数に律速されるため、幅も文字列長もばらついた画像が同じバッチに混在すると、短い画像ほど無駄な計算が増えてしまいます。

バケッティングなしの場合：幅も文字列長も異なる画像が同一バッチに混在し、短い画像ほど計算が無駄になる

そこで、切り出した行画像を内容幅に基づいて並べ替え、幅の近い画像同士を同じバッチにまとめるバケットバッチングを導入しました。バッチ内の画像幅が揃うことで、最大幅に合わせて追加されるパディングを最小限に抑えられます。

さらに、行画像の幅は、おおむね文字列長を反映します。幅の近い行をまとめると、同じバッチに含まれる文字列の長さも自然と近くなり、この性質がARデコードの効率化にもつながります。短い行を集めたバッチでは、各行が早い段階でEOSに到達するため、ARループを早期に終了できます。長い行は長い行同士でまとめられるため、ひとつの長文が多数の短文を巻き込んで、バッチ全体の反復回数を増やすこともありません。

本手法：可変長入力とバケットバッチングで幅・文字列長の近い画像を同じバッチに集約。バケットごとにARの再帰回数を最小化できる

dynamic-widthは主に「画像幅方向の無駄」を削減し、バケットバッチングは「パディングと文字列長方向の無駄」を削減します。2つを組み合わせることで、各行の実際の長さに応じて計算量を割り当てられる構成になります。

結果

今回の仕組みにより、一律の入力制限に頼らず、推論時の無駄な計算を抑えられるようになりました。その結果、軽量モデルの最大文字列長を、従来の50文字から100文字へ拡張しています。長い文字列を含む行が文書内に存在していても、その影響が同じバッチに含まれる短い行へ波及しにくくなり、短い行は短い行として、長い行は長い行として効率的に処理されます。

特に効果が大きいのはCPU推論です。実際の文書では、ページ幅いっぱいに文字が並ぶ行よりも、短い見出しや項目名、数値、住所など、幅の狭い行が多く含まれます。そのため、固定幅パディングとARループの削減が、そのまま処理時間の短縮につながります。また、今回のモデル再学習に合わせて、軽量モデルでも手書き文字を認識できるようになりました。

処理時間の定量評価

今回の変更による効果を確認するため、同一の文書と実行環境を使用して処理時間を測定しました。計測対象は、文字検出・文字認識・レイアウト解析を含むDocumentAnalyzer全体の実行時間で、各文書についてウォームアップを行った後、複数回実行した平均値を使用しています。

評価は2つに分けて行いました。1つ目は、同じ新軽量モデルを使用し、動的可変長幅入力とバケットバッチングの有無だけを切り替える比較で、今回導入した手法そのものの効果を確認します。2つ目は、通常モード・旧軽量モード・新軽量モードの比較です。こちらは使用するモデルの規模や認識可能な最大文字列長が異なるため、単純な高速化率ではなく、各モードが提供する機能と処理時間の関係を見るための比較です。

1.8倍

手法単体の高速化（CPU・同一モデル）

1.6倍

手法単体の高速化（GPU・同一モデル）

2倍

認識可能な最大文字列長（50→100文字）

2.4秒

新軽量モードの1枚平均（CPU）

測定環境

CPU：Intel Core i7-14700K（20コア / 28スレッド）
GPU：NVIDIA GeForce RTX 3090（24GB、1枚使用）
メモリ：32GB
PyTorch 2.6 / CUDA 12.4
通常モード：parseq-large-v4_1（PyTorch）
旧軽量モード：従来の--liteで使用していたparseq-tiny（最大50文字、CPUでは文字検出をONNXで実行）
新軽量モード：parseq-tiny-dynw-v4＋動的可変長幅入力＋バケットバッチング（最大100文字、CPUでは文字検出をONNXで実行）
計測回数：各文書のウォームアップ後にCPUで2回、GPUで3回計測した平均値。実行ごとに数％程度の変動があります

測定対象の文書

帳票、記事、表などを含む日本語文書8枚を使用しました。文書ごとの行数と文字列長は次の通りです。

文書	行数（文字列数）	平均文字列長	最大文字列長
文書1	72	22.9	100
文書2	202	15.5	42
文書3	170	5.5	27
文書4	134	13.0	67
文書5	34	10.8	33
文書6	36	5.3	22
文書7	64	9.7	41
文書8	45	8.9	48

多くの行は数文字から十数文字程度ですが、文書1や文書4のように、旧軽量モードの上限である50文字を超える行も含まれています。

手法単体の効果

まず、モデルの違いによる影響を除くため、新軽量モデルparseq-tiny-dynw-v4を共通して使用し、動的可変長幅入力とバケットバッチングだけをON/OFFして比較しました。OFFの場合は、すべての行画像を固定800px幅にパディングします。ONの場合は、各行の実際の幅に合わせて入力サイズを変更し、幅の近い画像を同じバッチにまとめます。

デバイス	OFF（固定幅）	ON（本手法）	高速化
CPU	4.14秒/枚	2.28秒/枚	1.8倍
GPU	0.60秒/枚	0.38秒/枚	1.6倍

文書ごとの高速化率は、CPUで1.4〜3.1倍、GPUで1.2〜2.4倍でした。行数が多く、短い行を多く含む文書ほど、固定幅パディングとARデコードの無駄が大きくなるため、本手法による効果も大きくなります。

通常・旧軽量・新軽量モードの比較

次に、通常モード、旧軽量モード、新軽量モードの処理時間を比較します。

処理時間（CPU）：

文書	通常モード	旧軽量モード	新軽量モード
文書1	23.2秒	2.0秒	2.9秒
文書2	36.7秒	3.7秒	4.5秒
文書3	27.9秒	3.0秒	2.1秒
文書4	31.1秒	3.0秒	3.0秒
文書5	7.7秒	1.3秒	1.5秒
文書6	7.0秒	1.1秒	1.1秒
文書7	13.4秒	2.0秒	2.2秒
文書8	10.4秒	1.8秒	1.9秒
平均	19.7秒	2.2秒	2.4秒

処理時間（GPU）：

文書	通常モード	旧軽量モード	新軽量モード
文書1	1.20秒	0.30秒	0.42秒
文書2	2.01秒	0.54秒	0.76秒
文書3	1.55秒	0.50秒	0.41秒
文書4	1.68秒	0.44秒	0.50秒
文書5	0.48秒	0.25秒	0.31秒
文書6	0.37秒	0.15秒	0.16秒
文書7	0.79秒	0.33秒	0.36秒
文書8	0.59秒	0.24秒	0.26秒
平均	1.08秒	0.35秒	0.40秒

通常モードとの比較では、新軽量モードはCPUで約8倍、GPUで約2.7倍高速です。ただし、この差には動的可変長幅入力やバケットバッチングの効果だけでなく、モデル自体の規模の違いも含まれます。

旧軽量モードと比較すると、新軽量モードの平均処理時間はCPUで2.2秒から2.4秒、GPUで0.35秒から0.40秒と、わずかに増加しています。ただし、この比較では両者が同じ量の情報を処理しているわけではない点に注意が必要です。

旧軽量モードは最大文字列長が50文字に制限されており、長い行でも画像を狭い入力幅へ縮小したうえで、50文字に到達した時点で認識を打ち切ります。そのため、文字数の多い行では各文字が横方向に強く圧縮され、文字の形状や文字間隔に関する情報も失われやすくなっていました。一方、新軽量モードは入力可能な画像幅を拡張し、行の実際の長さに応じた可変幅入力を使用するため、長い文字列でも1文字あたりの解像度や文字間隔を保ったまま扱えます。さらに、最大100文字までARデコードを継続するため、従来は途中で切れていた長い行も最後まで認識します。つまり、新軽量モードは旧軽量モードより多くの文字を処理し、長い行に対しては、より情報量の多い画像を入力しています。

新軽量モデルを固定800px幅のまま処理すると、CPUでは平均4.14秒/枚となり、旧軽量モードの約1.9倍の処理時間が必要です。動的可変長幅入力とバケットバッチングによってこれを2.28秒/枚まで短縮したことで、旧軽量モードに近い処理時間のまま、最大文字列長の拡張（50文字から100文字）、長い行画像の過度な縮小の抑制、手書き文字認識への対応が可能になりました。今回の変更は、旧軽量モードからの高速化そのものではなく、認識できる文字数と入力画像の情報量を増やしながら、処理時間の増加を抑えることを目的としています。

また、行数と総文字数に対する処理時間を図にすると、通常モードと新軽量モードの性質の違いが分かります。通常モードは行数や文字数が増えるにつれて処理時間が大きく増加するのに対し、新軽量モードでは増加が比較的緩やかです。

旧軽量モードと新軽量モードでは、処理している情報量が異なります。旧軽量モードは長い行画像を狭い入力幅へ縮小し、最大50文字で認識を打ち切るのに対し、新軽量モードは行の長さに応じた入力幅を使用し、最大100文字まで認識します。両者の処理時間を比較する際は、この違いを考慮する必要があります。

使い方

軽量モードは、CLIの--liteオプションから利用できます。動的幅推論を前提として学習した軽量モデルparseq-tiny-dynw-v4が使用されます。

yomitoku ${path_data} -f md --lite -d cpu -o results -v

Python APIから利用する場合は、text_recognizerに軽量モデルを指定し、dynamic_widthとbatch_bucketingを有効にします。

import cv2
from yomitoku import DocumentAnalyzer

configs = {
    "ocr": {
        "text_recognizer": {
            "model_name": "parseq-tiny-dynw-v4",  # 動的幅に対応した軽量モデル
            "dynamic_width": True,                 # 切り出し画像を実際の幅で処理
            "batch_bucketing": True,               # 幅の近い画像を同じバッチにまとめる
            "device": "cpu",
        },
        "text_detector": {
            "device": "cpu",
            "infer_onnx": True,                    # CPUでは検出モデルのONNX推論も有効
        },
    },
}

analyzer = DocumentAnalyzer(configs=configs, device="cpu")
img = cv2.imread("sample.jpg")
results, ocr_vis, layout_vis = analyzer(img)
results.to_json("output.json")

parseq-tiny-dynw-v4は、動的幅入力とバケットバッチングを前提として学習したモデルです。そのため、dynamic_width=Trueおよびbatch_bucketing=Trueと組み合わせて使用することを推奨します。なお、ONNX推論では入力サイズが固定されるため、動的幅入力は自動的に無効化されます。

おわりに

従来の軽量モデルでは、入力画像の幅と最大文字列長を一律に制限することで処理を高速化していましたが、その代わりに、50文字を超える長い行を最後まで認識できないという課題がありました。今回、動的可変長幅入力とバケットバッチングを組み合わせることで、短い行には少ない計算量を、長い行には必要な計算量を割り当てられるようになりました。

新軽量モードは、旧軽量モードと比べて平均処理時間がわずかに増えています。その一方で、最大文字列長を50文字から100文字へ拡張し、従来は途中で打ち切られていた長い行も最後まで認識できるほか、軽量モデルでの手書き文字認識にも対応しました。認識範囲を制限して速度を確保するのではなく、無駄な計算を減らすことで、認識範囲を広げながら速度を維持する、というのが今回の変更の方針です。

最新の軽量モデルは、YomiToku v0.13.1およびYomiToku Studioで利用できます。なお、YomiToku StudioはONNXで推論を実行しており、Python版とは実行方式が異なります。動的可変長幅入力やARデコードの早期終了による効果が限定されるため、Python版と比較すると処理速度が遅くなる場合があります。ぜひ、手元の文書で精度と速度をお試しください。

YomiToku Studio : https://yomitoku-studio.mlism.com/
YomiToku : https://github.com/kotaro-kinoshita/yomitoku