日本語レシートOCRと名前付きエンティティ抽出：AWS SageMaker ServerlessとTriton Inference Serverを使用した複数モデルによる低コスト推論

今回のブログでは、日本語のレシートから情報を抽出するためのサーバーベースの光学式文字認識（OCR）と名前付きエンティティ（NE）のデモをどのように展開したかをお話したいと思います。このデモは、様々な異なるソフトウェア技術の長所を組み合わせて、複数のモデルで低コストに推論を行う方法を示す良いデモだと思います。
現在のOCR/NEデモサービスは、2つの別々の処理パイプラインで構成されており、並行して実行され、基本的に互いに独立しています。一方のパイプラインは「従来の」OCR/NEカスケードで、もう一方はすべての処理を「一度に」行う、いわゆるエンドツーエンドモデルです。
従来のOCR/NEパイプラインでは、4種類のディープラーニングモデルを用意しており、それぞれTensorFlowを用いてゼロから学習、もしくはファインチューニングを行っています。
  1. 画像中の中央のレシートの角を検出し（複数のレシートが見える場合）、まっすぐになるように修正するモデル
  2. テキストの連続する領域の位置を検出するモデル
  3. 検出された各テキスト領域内の文字を認識するためのOCRモデル、および
  4. BERT大規模言語モデル（LLM）に基づくNEモデル。このモデルは画像内で認識されたすべてのテキストを受け取り、上から下、左から右へとソートし、各文字を住所、日付、店名、合計金額など、関心のある数百種類の名前付きエンティティのうちの一つに属するよう分類します。
このパイプラインは、情報抽出タスクを複数のモジュールに分割しているため、「従来型」と呼んでいますが、実際には2つの新規コンポーネントが含まれています。 すなわち、（1）レシート検出・修正モジュール（詳細はこちら）、（4）テキストレベルで動作し、テキストが画像内の元の位置にあったという明示的位置情報を無視するNEモデルです。
Donutは、合成的に生成された膨大な量のテキスト画像データに対して事前に学習させたエンコーダ・デコーダ画像トランスフォーマーモデルであり、従来のパイプラインで使用したものと同じ名前付きエンティティを認識するためにPyTorchを用いてファインチューニングを行ったものです。 このアプローチは、今回紹介したエンドツーエンドの受信処理方法と同様です。
実際のレシートから情報抽出を行うため、バックエンド全体を最も安価なインフラで展開したいと考えています。デモはレスポンシブであることが望ましいのですが、常時稼働させ実際に使用しないときに無駄なランニングコストが発生するのは避けたいです。また、将来的には拡張性の高いサーバーインフラに同じソフトをシームレスに展開し、最小限の変更で済むようにしたいとも考えています。
AWS SageMaker Serverlessは、このようなユースケースに最適です。なぜなら、入ってきたリクエストを処理するために発生する推論時間の合計に対してのみ課金されます。リクエストがなければ、何のコストもかかりません。また、各リクエストをできるだけ早く処理しようというインセンティブが働くので、各コンポーネントモデルをできるだけ高速に最適化することになります。さらに、365日・24時間稼働するシステムをデプロイする準備が整えば、AWS SageMaker Serverless用に構築した同じDockerイメージを、AWS SageMaker Real-time Inference や AWS SageMaker Asynchronousに簡単に移行できます。
AWS SageMakerをデプロイメントターゲットに選んだら、次にその上で動作する推論ソフトを決める必要があります。AWSではGoogleのTensorFlow Serving、IntelのOpenVINO Model Server、NvidiaのTriton Inference Serverなどの選択肢が用意されています。私たちが知る限り、Triton Inference Serverだけが、異なるモデル推論フレームワークで学習または変換された異なるモデルを、1つのサーバーインスタンスで実行できるオプションを提供しています。そこで今回は、OpenVINO、ONNX、PyTorch、TensorFlow、PythonのモデルプラグインでTriton Inference Serverを実行することにしました。
また、Flutterでシンクライアントを開発し、以下のスクリーンショットのように、アプリ内でレシートを撮影したり、PCやモバイル端末のカメラロールに保存済みのレシートを選択したりすることができるようにしました。
そして、その画像をAWSのSageMakerエンドポイントに送り、画像を処理し、次のスクリーンショットに示すように、2つの異なる処理パイプラインから結果を返してアプリに表示します。
これらの異なるパーツがどのように組み合わされているのか、ハイレベルなビューを下図に示します。
このシステム全体をスムーズに稼働させるためには、数多くの課題がありました。その1つが、AWS SageMaker Serverless上でTriton Inference Serverを動作させることでした。SageMaker Serverlessは、デフォルトでセキュリティ上の理由から共有メモリを無効にしています。この問題はソースコードを修正し、Dockerイメージを再コンパイル、再ビルドすることで何とか解決しました。
もう一つの問題は、Triton Inference Server PyTorchプラグインを使って、複雑なPyTorchモデルを実行させることでした。プラグインはモデル学習時によく使われるPython APIではなく、C++ APIを使うため、まずDonutでエンコーダとデコーダのコンポーネントを分離し、それぞれを別々にトレースしてディスクに保存する必要がありました。その後、PyTorchプラグインのC++ APIを通じて、エンコーダを個別に呼び出し、エンコーダの出力でデコーダを個別に呼び出して、読み込むことができます。
Donutのエンコーダとデコーダのモデルはそれぞれ500MB程度と非常に大きく、NEのモデルはさらに大きく800MB程度です。Donutのエンコーダとデコーダの量子化には成功しましたが、この量子化されたモデルを使って推論を行うことはできませんでした。これはまだ解決策を検討しているところです。しかし、NEモデル（BERTエンコーダ型LLM）は、ONNXに変換して量子化した後、サイズを4分の1に縮小することに成功しました（このモデルは、BERTエンコーダ型LLMであることを思い出してください）。そして、この量子化されたONNXモデルをOpenVINOフォーマットに変換し、最終的なシステムで使用することにしました。
最終的なシステムによる推論速度は、従来のパイプラインで約1秒、エンドツーエンドのDonutパイプラインで約10秒です。すべてのモデルは、6Gb RAMを搭載した単一のAWS SageMaker Serverlessインスタンスで同時に実行されています。コールドスタート時間（S3からモデルファイルをコピーし、サーバーレスインスタンスを起動し、リクエストに応答するのにかかる時間）は約45秒です。つまり、最初のリクエストは最大タイムアウトの30秒を超えますが、その後のリクエストはほぼ瞬時に完了します。最初のタイムアウトを回避するために、コールドスタート時間を短縮する努力を続けています。
将来的には、Triton推論サーバをOpenVINO Model Serverで置き換え、Intelハードウェア上で動作するためのソフトウェア最適化を図ることも視野に入れています。特に、モデルの量子化とプルーニング、モデルのロードタイムの短縮、推論速度の高速化など、最新のOpenVINO 2023.0リリースにも大きな期待をしています。そのためには、既存のTensorFlowやPyTorchのモデルをすべてOpenVINOのモデルフォーマットに変換する必要がありますが、これについては今後詳しく書いていく予定ですので、その時はまたよろしくお願いします。
キーワード：TensorFlow, PyTorch, ONNX, OpenVINO, Triton Inference Server, AWS SageMaker, Serverless, Real-time Inference Server, client-server, OCR, NE, 日本語レシート
Tags:
BLOG