こんにちは、unerry CTOの伊藤です。 2025年9月、データサイエンティスト上野優人が、北海道で開催された「情報科学技術フォーラム（FIT）」において、 「位置情報データと購買データを活用した広告セグメントの開発」 に関する発表を行いました。 今回の発表は8月の「Google Cloud Next Tokyo」での登壇に続くもので、最先端技術の実装に新卒のエンジニアが挑んだ記録でもあります。 講演内容の核心となる技術、そして若きデータサイエンティストとしての挑戦の舞台裏について、上野に話を聞きました。 登場人物 株式会社unerry テクノロジー＆オペレーション部 データサイエンス&AIチーム 上野 優人（うえの ゆうと） 入社日： 2025年4月 最近の推し： 令和ロマン 筑波大学を卒業後、上智大学大学院 応用データサイエンス学位プログラムを修了。大学院では、「価格・需要変動下における、利益最大化のための販売戦略」に関する研究を行った。在学中より、unerryでの長期インターンを経験し、保有するデータと働く人に魅力を感じて新卒入社。現在は、位置情報・購買データを用いたロジック開発および改善に取り組んでいる。 ＜聞き手＞株式会社unerry CTO 伊藤 清香（いとう さやか） 入社日： 2018年2月 最近の推し： ピェンロー鍋 ガラケーからスマホまで20年以上モバイルWebシステムを開発し、高負荷対策をノリと勘で支えた縁の下の力持ち。人生の節目にあたり、これからはIoTで人々の生活を便利にしようと考えて、当時10人位だったunerryへJoin。会社の成長とともに湯水のように湧き出る課題を解決し、働きやすい職場環境を作ることを生きがいとしている。趣味はサッカー観戦と音声制御技術。 第1章：推薦システムを革新する「Two-Tower モデル」の技術的深掘り 伊藤： 今回の講演の核となった技術について、詳しく教えてください。 上野： はい、講演では、一言でいうと 位置情報データと購買データ を掛け合わせた次世代ターゲティングモデルについてお話ししました。このモデルは、ユーザーが過去にどこで行動したかという情報（位置情報データ）を、どの商品を買ったかという情報（購買データ）と組み合わせることで、より高精度な広告セグメントの構築を実現するものです。 この推論モデルは、unerryの梅田と張が共同で発明した特許（番号：特許7641682）を実装したものです。(*1) そして、その技術的な中核を担っているのが 「Two-Towerモデル」 というアーキテクチャです。これは、大規模ユーザーに対して高速に推論できるという利点から、YouTubeなど大手テック企業で採用されている先進的なアルゴリズムです。 伊藤： その「Two-Towerモデル」が従来の推薦システムと比較して画期的なのはどのような点でしょうか？ 上野： 主に、従来のシステムが抱える大きな課題を解決できる2点にあります。 1. 新商品に対する推薦が可能： 一般的に、小売企業が持つPOSデータだけを使った推薦システムでは、新商品を販売する際、購買データが全くないため、誰に推薦したらよいか分かりません。しかし、Two-Tower モデルは、商品の特徴量（価格、カテゴリなど）から生成したベクトルで推薦を行うため、データがない新商品でも適切なユーザーに推薦できます。 2. 購買履歴がないユーザーにも推薦が可能： リテール（小売）の購買データがないユーザー、つまりそのお店で買ったことがないユーザーは、従来のシステムではターゲティングできませんでした。しかし、当社は位置情報データを持っています。位置情報データから抽出・推定したユーザーの行動DNA（unerry独自の指標：普段の行動傾向を示す）や性別・年代といった特徴量があれば、購買履歴がないユーザーに対しても、「この商品を買いそうだ」という可能性を予測できます。 伊藤： その高速な処理を実現するアーキテクチャについて、具体的に解説いただけますか？ 上野： Two-Tower モデルは、名前の通り、 ユーザーの特徴量と商品の特徴量という2つのタワー で構成されています。 ユーザーの性別や年代といった特徴量、そして商品の価格やカテゴリといった特徴量を、それぞれ深層学習（DNN）で処理することで、意味のある 「ベクトル」 （埋め込み表現、エンベディング）を生成します。 推薦のスコアは、この 「ユーザーベクトル」と「商品ベクトル」の内積 で算出されます。内積が大きいほど、ユーザーがその商品に興味を持っていると判断できます。 高速化の肝は、 オフラインとオンラインの処理を分けている点 です。 ●オフライン処理： 商品のベクトルは頻繁に変わらないため、事前に計算し、データベースに保存しておきます。 ●オンライン処理： ユーザーのベクトルだけをリアルタイムで計算し、保存しておいた商品ベクトルと照合（近似最近傍探索）することで、瞬時に推薦結果を出すことができます。 YouTubeなどのテック系企業で採用されているのも、この「大規模ユーザーに対して瞬時に結果を出せる」というスケーラビリティと速度が最大の要因です。ちなみに、今回採用したベクトルの次元数は128次元で、一般的なシステムで使われる700次元や1000次元と比較しても、 軽量でリーズナブルな計算資源 で済むという利点もあります。 第2章：実装を阻む壁と300回超のトライ＆エラー 伊藤： この最先端の技術を実装する過程で、特に大変だったのはどのようなことでしょうか？ 上野： 非常に多岐にわたりましたが、最大の困難は 「実装の難しさ」 でした。Two-Tower モデルは概念はシンプルですが、適切なベクトルを生成するための深層学習レイヤーの学習が非常にデリケートで難しいと言われています。実際に手を動かすと、なかなか期待通りの精度が出ませんでした。 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ＜補足＞ Two-Towerモデルについて： Google の YouTube 推薦アルゴリズムなど、大手テック企業で採用されており、大規模ユーザーに対して高速に推論できるという点で革新的。ただし扱いが難しくまだ広く浸透していない。 参考動画 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 私のログを確認したところ、モデルの試行回数は300回以上に及びました。最初はもちろん、コードの書き間違い（コーディングミス）も多くありましたが、その後は主に「ベクトルの精度をどう上げるか」という試行錯誤の連続でした。 伊藤： ベクトルの精度向上は、具体的にどのように進めたのでしょうか？ 上野： 精度を上げるためには、モデルに「正解」を教えて学習させる必要があります。私たちは、ユーザーIDに位置情報データの行動パターンから推定した属性を特徴量（性別、年代など）とし、実際の「購買データ」と紐づけました。「このユーザーがこの商品を買った」というデータには「1」（正解）を、「買ってない」というデータには「0」（不正解）を与えます。 そして、モデルが算出した内積スコアが、この正解（1か0）に近づくように、深層学習レイヤーを学習させていくんです。適当なベクトルだと意味のないスコアが出てしまうので、「ここは1ですよ」という正解を与えることで、ベクトルの精度を上げていきました。 伊藤： 講演の登壇準備と、このモデル開発を同時並行で進めるのは、相当な負荷だったと想像します。 上野： おっしゃる通りです。登壇の締め切りに追われる中で、コードを大量に書き、試行錯誤を繰り返す日々でした。しかし、その結果として、 YouTube や他のビッグテック企業が採用しているのと「同じレベルの技術」を、当社のビジネスに組み込むことができたのは、大きな達成感 でした。まさに「困難を乗り越えたからこそ、価値がある」と実感しています。 第3章：学会の独特な雰囲気と、2度の国際的な登壇経験 伊藤： 会場の雰囲気はいかがでしたか？ 上野： 学会の雰囲気は、一般の技術カンファレンスとは異なり、独特の緊張感がありました。リアル会場には20名程度の参加者がいたかと思います。 伊藤： 質問はありましたか？ 上野： はい、お一人の方から質問をいただきました。登壇内容というよりは、当社の事業領域である「人事領域のAI活用」に関する相談でした。これは、技術広報と採用という今回の登壇目的にも合致しており、意義のある交流となりました。 伊藤： 実は、このFITを含めて、上野さんは短期間で連続して登壇されていると聞きました。 上野： はい、プライベートも含めると5ヶ月で4回となります。 ① 5月：日本経営工学会（国内学会） 卒業後に参加。大学院での研究テーマ（中古スマートフォンの販売先最適化）を発表。 ② 7月：ICPR（国際会議、コロンビア） 指導教員の計らいで、単身コロンビアへ渡航。経営工学に関する研究を発表しました。治安や言語の面で非常にタフな環境でしたが、貴重な経験でした。 ③ 8月：Google Cloud Next Tokyo（国内）クラウド技術大規模カンファレンス ④ 9月：FIT（今回の登壇） 伊藤： コロンビアでの単身登壇は驚きです。短い準備期間での挑戦も大変だったと思いますが、何かエピソードはありますか？ 上野： FIT登壇の準備期間は1週間ほどしかありませんでした。特に大変だったエピソードとして、飛行機の機内で発表練習をしていたことがあります。 飛行機が遅延し、時間ができたため、PDF資料を読み込みながら、頭の中でプレゼンを再生し、タイマーで時間を計るというスタイルで練習を続けていました。ブツブツと声に出すことはしませんでしたが、頭の中ではひたすら時間を調整していました。 また、登壇全体を通して、先輩から非常に手厚いフィードバックをいただきました。 ●「短い言葉で言い切ること」 ●「初見の専門用語をいきなり使ってしまうと、聴衆がついていけなくなる」 といった、スライド作成術から話し方まで、実戦を通じて学ぶことができました。特にGoogle Cloud Nextの際は、他の登壇者との兼ね合いで持ち時間が短くなるという裏事情もありましたが、学んだ技術を活かし、説明の核を外さずにコンパクトにまとめることができたと思います。 第4章：未来の仲間へ。「交流」の場としての学会の価値 伊藤： 学会全体を通して、上野さんが最も重要だと感じたことは何でしょうか。 上野： それはやはり 「交流」 です。 発表者側としては、質問を1人からしか得られなかった反省から、いかに相手に興味を持ってもらえる発表をするかという難しさを痛感しました。一方で、聴衆側として、自社のビジネスに関連のあるセッションには積極的に質問しに行きました。例えば、 自然災害時に避難場所を教えるチャットボット に関する研究は、当社のビジネスとも関連しそうで、非常に興味深く、質問を通して発表者の方と有益な関わりを持つことができました。 学会は、最新の技術動向を知るだけでなく、普段関わることのない研究者や学生とコネクションを作り、自分では気づかなかった新しい観点での気づきを得られる場です。 伊藤： 最後に、同じようにデータサイエンスを深く突き詰めたい学生、そして未来の仲間たちにメッセージをお願いします。 上野： 私は大学院で数理最適化を学び、その専門性が現在のデータサイエンスの仕事にダイレクトに活きています。入社後わずか数ヶ月で、世界的にも先進的な技術であるTwo-Tower モデルの実装に挑戦し、それをビジネスに組み込むという経験ができました。 「学んできたことを、社会の現場で直線的に活かしたい」、「困難な技術に果敢に挑戦し、その成果を世の中に羽ばたかせたい」という熱意 を持った方にとって、unerryは非常に恵まれた環境です。 私たちと共に、最先端のデータサイエンスを深掘りし、世の中を動かす技術を生み出していく仲間になりませんか？ *1 Google Cloud Next Tokyo ‘25の登壇記事もありますので参照ください。 Vertex AIで実現：購買データ x 約1億IDの人流データによる次世代広告ターゲティング / 「 Google Cloud Next Tokyo 」登壇レポート https://www.unerry.co.jp/blog/google-cloud-next... 「Google Cloud Next Tokyo」はGoogle Cloudが年に1回開催するイベントの日本版で、クラウド技術の最新情報や事例の紹介に加え多彩なワークショップなどを含み、今年は2025年8月5日(火)と6日(水)の2日間、東京ビッグサイトで開催されました。 本記事は8月5... unerryでは、行動データの可能性を共に切り拓くデータサイエンティストやエンジニアを募集しています。挑戦できる環境で価値創造に取り組みたい方は、ぜひお問い合わせください。 株式会社unerry 採用ページへ The post 300回超の試行錯誤を経て新卒データサイエンティストが開発に挑む「人流×購買データによる広告ターゲティング手法」 first appeared on 株式会社unerry .

はじめに NTTデータグループでは、デジタルツインのシステムを設計・構築・活用するノウハウや、各種データ処理ツールを強みに、お客さまの課題を解決し、効率や収益性を高めるデジタルツインサービスを提供しています。 デジタルツインは目的に応じて多様なデータ・機能から構築されます。本記事ではデジタルツインを構成する重要な要素である「点群」に着目し、点群を用いた物体認識技術と、その取り組みにおけるNTTの研究成果を紹介します！ デジタルツインにおける点群の役割 そもそも皆さまは「点群」をご存知でしょうか。 点群とは、LiDAR（Light Detection And Ranging）やレー

こんにちは。イノベーションセンターの加藤です。普段はコンピュータビジョンの技術開発やAIシステムの検証に取り組んでいます。 今回は最新版のPyTorchを使って軽量なTransformerベースOCRモデルであるPARSeq(Permuted Autoregressive Sequence)をTensorRTモデルに変換して高速化した取り組みについて紹介します。 PARSeqとは PARSeqのTensorRT化 PyTorch Lightningによるモデル変換 AutoregressiveとIterative refinementがTensorRT化できない問題 Autoregressive modeのTensorRT化 TorchDynamoの機嫌をとる Iterative refinementのTensorRT化 評価 まとめ PARSeqとは PARSeq 1 はVision Transformer(ViT)を特徴抽出器として用いる文字認識モデルであり、以下の画像のような文章生成の形をとっています。 このような文章生成モデルでは、まず画像をトークンに分割したものをTransformer Encoderで特徴抽出し、これをもとにTransformer Decoderで次の文字トークンの予測を繰り返します。PARSeqの場合は文字トークンの予測方法にオプションがあり、以前の予測を参照しながら１文字ずつ予測するもの(Autoregressive)、一度に全部の文字を予測するもの(Non-autoregressive)、一度予測した文字を入力し直して洗練するもの(Iterative refinement)の三通りのデコード戦略があります。 PARSeqの特徴はTransformerベースでありながら非常に軽量である点です。 Encoder部分は一般的なViTと同様に12層のTransformerレイヤーで構成されていますが、Decoder部分はたった1層しかなく、 一般的なVision Language Modelが数十億のパラメータを抱えている一方でPARSeqは数千万パラメータに留まっています。 PARSeqのTensorRT化 このPARSeqモデルをさらに高速化するために、今回はTensorRTモデルに変換します。 TensorRT 2 は、NVIDIAが提供しているディープラーニングモデルの推論を高速化するためのツールで、さまざまなAIフレームワークが対応している共通フォーマットのONNX 3 からの変換や、PyTorchモデルからの直接変換が可能です。 実はNVIDIAが公式ブログでPARSeqをTensorRT化する記事を公開している 4 のですが、 PARSeqやその依存先のPyTorchのバージョンが古くそのままでは動作しないため、本稿では最新版(PyTorch 2.10, PARSeq 2024年2月版)を使ったTensorRT化の流れを紹介します。 PyTorch Lightningによるモデル変換 PARSeqはPyTorchによって実装されたモデルをPyTorch-Lightningで制御しており、ONNXやTensorRTへの変換はPyTorch-Lightningが提供する関数を利用できます。 NVIDIAのブログでも to_onnx() を利用して一度ONNX化したのち、trtexecと呼ばれるツールを使ってONNXからTensorRTへ変換しています。 今回は to_tensorrt() を利用して、モデルを直接TensorRTに変換してみます。 import torch parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() parseq.model.refine_iters = 0 # Iterative refinementを無効化 parseq.model.decode_ar = False # Non-autoregressive mode output_path = "engine.pt2" img = torch.randn( 1 , 3 , 32 , 128 ) parseq.to_tensorrt(output_path, img, ir= "dynamo" ) これで無事TensorRTモデル engine.pt2 に変換できました。このモデルは以下のように呼び出すことができます。 import torch import torch_tensorrt # <- 必須 parseq = torch.export.load( "engine.pt2" ).module() img = torch.randn( 1 , 3 , 32 , 128 ).cuda() parseq(img) # torch.Size([1, 26, 95]) 26は一度に推測可能な文字数、95は対応文字種 AutoregressiveとIterative refinementがTensorRT化できない問題 しかしながら、この方法ではAutoregressive( decode_ar=True )またはIterative refinement( refine_iters>0 )に対応したモデルを作ろうとするとエラーになってしまいます。 論文ではNon-autoregressiveよりAutoregressiveの方が高精度 5 とされており、またIterative refinementも1回適用するだけでそれなりに精度が向上するため、ぜひこれらのモードもTensorRTで活用したいです。 そこでPARSeqの実装を改造しTensorRT化に挑戦しました。 Autoregressive modeのTensorRT化 まず先ほどと同じ方法ではどこで落ちるかをみてみます。 import torch parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() parseq.model.refine_iters = 0 parseq.model.decode_ar = True # AR mode output_path = "engine.pt2" img = torch.randn( 1 , 3 , 32 , 128 ) parseq.to_tensorrt(output_path, img, ir= "dynamo" ) 表示されるエラーは以下のとおりです。 File "/root/.cache/torch/hub/baudm_parseq_main/strhub/models/parseq/model.py", line 144, in forward if testing and (tgt_in == tokenizer.eos_id).any(dim=-1).all(): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ... torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1)). (Size-like symbols: none) これは「文章の終了を示すEOSトークンが出たら生成を停止する」処理の部分であり、どうもif文による分岐はTensorRTと相性が悪いようです。 しかしこれは１文字生成を繰り返すAutoregressive modeでは必須の処理であるため、1文字生成する実装のみをTensorRT化し、繰り返し部分はモデルの外側でやるように変えてみます。 import pytorch_lightning as pl from torch import Tensor from typing import Optional class PARSeqEncoder (pl.LightningModule): def __init__ (self, model): super ().__init__() self.encoder = model.encoder def forward (self, images: Tensor) -> Tensor: memory = self.encoder(images) return memory class PARSeqDecoder (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed self.pos_queries = model.pos_queries self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) null_ctx = self.text_embed(input_ids[:, : 1 ]) tgt_emb = self.pos_queries[:, :S- 1 ] + self.text_embed(input_ids[:, 1 :]) tgt_emb = torch.cat([null_ctx, tgt_emb], dim= 1 ) tgt_query = self.pos_queries[:, S- 1 :S].expand(B, - 1 , - 1 ) tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, content_mask=tgt_mask) return self.head(decoder_outputs) ここでエンコーダとデコーダが切り離されています。これはエンコードを一度実行したのち、デコードをEOSトークンが出るまで繰り返す必要があるためです。 推論は以下のようになります。 img_transform = T.Compose([ T.Resize(( 32 , 128 ), T.InterpolationMode.BICUBIC), T.ToTensor(), T.Normalize( 0.5 , 0.5 ), ]) _parseq = torch.hub.load( 'baudm/parseq' , 'parseq' , pretrained= True ).eval() bos_id = _parseq.tokenizer.bos_id pad_id = _parseq.tokenizer.pad_id eos_id = _parseq.tokenizer.eos_id parseq_encoder = PARSeqEncoder(_parseq.model) parseq_decoder = PARSeqDecoder(_parseq.tokenizer, _parseq.model) img = Image.open( "world.png" ).convert( "RGB" ) img = img_transform(img).unsqueeze( 0 ) with torch.no_grad(): num_steps = _parseq.model.max_label_length + 1 input_ids = torch.full(( 1 , num_steps), pad_id, dtype=torch.long) input_ids[:, 0 ] = bos_id memory = parseq_encoder(img) preds = [] for i in range (num_steps- 1 ): j = i + 1 logit = parseq_decoder(memory, input_ids[:, :j]) preds.append(logit.softmax(- 1 )) input_ids[:, j:j+ 1 ] = logit.argmax(- 1 ) if (input_ids == eos_id).any(dim=- 1 ).all(): break label, confidence = _parseq.tokenizer.decode(torch.cat(preds, dim= 1 )) print (f "AR result: {label[0]}" ) そして変換は次のように行います。 input_ids の長さは伸び縮みするため最短・最長を指定しておく必要があります。 parseq_encoder.to_tensorrt( "encoder.pt2" , img, ir= "dynamo" ) decoder_input_ids = torch_tensorrt.Input( min_shape=[ 1 , 1 ], opt_shape=[ 1 , num_steps], max_shape=[ 1 , num_steps], dtype=torch.int64) encoder_outputs = torch_tensorrt.Input( min_shape=[ 1 , 128 , 384 ], opt_shape=[ 1 , 128 , 384 ], max_shape=[ 1 , 128 , 384 ], dtype=torch.float32) parseq_decoder.to_tensorrt( "decoder.pt2" , (encoder_outputs, decoder_input_ids), ir= "dynamo" ) TorchDynamoの機嫌をとる しかしながら、なぜかこれはデコーダ( PARSeqDecoder )の変換に失敗します。本来入力する input_ids のトークン長は1以上あれば動作するはずですが、以下のように3以上に限定しなさいというエラーが出てきます。 - Not all values of _1 = L['input_ids'].size()[1] in the specified range _1 <= 26 satisfy the generated guard 3 <= L['input_ids'].size()[1] and L['input_ids'].size()[1] <= 26 Suggested fixes: _1 = Dim('_1', min=3, max=26) これはTensorRT化よりも前の、TorchDynamoがソースコードを解析するときに発生しているエラーなのですが、どこが原因なのかをTorchDynamoを使って探ってみます。 from torch_tensorrt.dynamo.utils import get_torch_inputs, to_torch_device from torch_tensorrt.dynamo._tracer import get_dynamic_shapes_args from torch.export import Dim, export, draft_export arg_inputs = (encoder_outputs, decoder_input_ids) parseq_decoder.to( "cuda" ) device = to_torch_device( "cuda" ) torch_arg_inputs = get_torch_inputs(arg_inputs, device) dynamic_shapes = get_dynamic_shapes_args(parseq_decoder, arg_inputs) ep = draft_export( # エラーが起きても最後まで解析させることで全てのエラーを収集する parseq_decoder, tuple (torch_arg_inputs), dynamic_shapes=dynamic_shapes, ) print (ep._report) すると以下のような警告が確認できます。 ################################################################################################### WARNING: 2 issue(s) found during export, and it was not able to soundly produce a graph. Please follow the instructions to fix the errors. ################################################################################################### 1. Guard Added. A guard was added during tracing, which might've resulted in some incorrect tracing or constraint violation error. Specifically, this guard was added: Ne(s70 - 1, 1), where {'s70': "L['input_ids'].size()[1]"}. This occurred at the following stacktrace: File /opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py, lineno 1776, in _wrapped_call_impl File /opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py, lineno 1787, in _call_impl File /workspace/src/ar_deploy_decoder.py, lineno 31, in forward tgt_emb = self.pos_queries[:, :S-1] + self.text_embed(input_ids[:, 1:]): Locals: self: [None] S: ['s70'] input_ids: ['Tensor(shape: torch.Size([1, s70]), stride: (s70, 1), storage_offset: 0)'] Symbols: s70: L['input_ids'].size()[1] And the following framework stacktrace: File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 404, in is_contiguous_for_memory_format File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 317, in is_contiguous File /opt/venv/lib/python3.12/site-packages/torch/_prims_common/__init__.py, lineno 277, in check_contiguous_sizes_strides if maybe_guard_or_false(x == 1): (以下省略) テンソルを S-1 の長さにスライスするところで S-1 != 1 という制約がDynamoによって導入されています。 どうやらスライスをした時長さ1に なりうる 可変長テンソルは問題があるようです。（おそらく0/1-specialization 6 と呼ばれる処理と関係があるのですが、なぜこうなっているのかはよく分かりません...） そこでスライスを行わない形に実装を直しておきます。 class PARSeqDecoder (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed # self.pos_queries = model.pos_queries self.prefixed_pos_queries = torch.nn.Parameter(torch.cat([torch.zeros_like(model.pos_queries)[:,: 1 ], model.pos_queries], dim= 1 )) self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) tgt_emb = self.prefixed_pos_queries[:, :S] + self.text_embed(input_ids) tgt_query = self.prefixed_pos_queries[:, S:S+ 1 ].expand(B, - 1 , - 1 ) tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, content_mask=tgt_mask) return self.head(decoder_outputs) これで無事変換が通るようになりました。 Iterative refinementのTensorRT化 次にIterative refinementを行うデコーダのTensorRT化を行います。 元のPARSeq実装からrefinementを行う箇所を切り出しPyTorch Lightningでラップします。 class PARSeqRefiner (pl.LightningModule): def __init__ (self, tokenizer, model): super ().__init__() self.tokenizer = tokenizer self.max_label_length = model.max_label_length self.text_embed = model.text_embed self.prefixed_pos_queries = torch.nn.Parameter(torch.cat([torch.zeros_like(model.pos_queries)[:,: 1 ], model.pos_queries], dim= 1 )) self.pos_queries = model.pos_queries self.decoder = model.decoder self.head = model.head def forward (self, memory: Tensor, input_ids: Tensor) -> Tensor: B, S = input_ids.size( 0 ), input_ids.size( 1 ) tgt_emb = self.prefixed_pos_queries[:, :S] + self.text_embed(input_ids) tgt_query = self.pos_queries tgt_mask = torch.triu(torch.ones((S, S), dtype=torch.bool), 1 ).to(tgt_emb.device) tgt_mask[torch.triu(torch.ones((S, S), dtype=torch.bool, device=tgt_emb.device), 2 )] = 0 tgt_padding_mask = (input_ids == self.tokenizer.eos_id).int().cumsum(- 1 ) > 0 decoder_outputs = self.decoder(tgt_query, tgt_emb, memory, query_mask=tgt_mask, content_mask=tgt_mask, content_key_padding_mask=tgt_padding_mask) return self.head(decoder_outputs) refiner_input_ids = torch_tensorrt.Input( min_shape=[ 1 , num_steps], opt_shape=[ 1 , num_steps], max_shape=[ 1 , num_steps], dtype=torch.int64) print ( "==== export refiner ====" ) parseq_refiner.to_tensorrt( "refiner.pt2" , (encoder_outputs, refiner_input_ids), ir= "dynamo" ) こちらは入力トークンが伸び縮みしないのもあり素直に変換できました。 評価 最後にTensorRT化によってどれくらい速くなったかをみてみます。 OCRのベンチマークであるIIIT-5Kに対してさまざまな設定で推論し、1枚あたりのレイテンシをH200 GPU 1台で計測しました。 結果は次の図のようになりました。 例えばAutoregressive(AR)モード・iterative refinement無しではTensorRT変換によって2.58倍の高速化、 Non-Autoregressive(NAR)モードでは3.07倍の高速化を達成しました。 グラフの傾きからiterative refinementも軽量になっていることが分かります。 まとめ 今回の実験では軽量で高性能なOCRモデルであるPARSeqを最新の環境でTensorRT化してみました。 その際、文章生成などでよく用いられるデコーダは入力サイズが動的に変化するため変換に一癖あり、ライブラリが処理しやすいようなプログラムに書き換える必要があることを紹介しました。 https://github.com/baudm/parseq ↩ https://developer.nvidia.com/tensorrt ↩ https://onnx.ai ↩ https://developer.nvidia.com/blog/robust-scene-text-detection-and-recognition-inference-optimization/ ↩ https://arxiv.org/abs/2207.06966 Appendix H ↩ https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html#are-always-specialized ↩