TensorRTでfloat型tensorをResizeやSliceの第二以降の引数に指定するモデルが内部エラーになるのでフォーラムに連絡してみた
とあるディープラーニングモデルをONNXに変換してTensorRTで動かそうとしたらInternal Errorが出るのでNVIDIAさんのフォーラムに凸ってみた。
該当スレッドはここ。
モデルに関する簡単な説明
まずはONNXのResizeオペレータとSliceオペレータについて。
Resizeはoutput_feature = Resize(input_feature, roi, scales, sizes)のような形式で第四引数sizesで拡大・縮小後の画像サイズを指定すると拡大・縮小してくれるオペレータ。
Sliceはoutput_feature = Slice(input_feature, starts, ends, axes, steps)のような形式で第二引数以降で指定したスライス操作を行ってくれるオペレータ。
これらのオペレータに対して第二引数以降の引数にfloat型tensorでサイズ計算などを実行してから整数型にキャストしたものを入れていた。例えばResizeには直前のConvolution結果のshapeを2倍に拡大するためにfloat型でwidth, heightを×２してからFloorで端数切り捨て＋キャストで整数型にしてsizesに渡していた。
フォーラムに投稿
TensorRTはプラグインとモデルのパーサーについてはオープンソースで公開されているが、他の部分はプロプリエタリでソースコードは非公開になっている。内部エラーということで、NVIDIAの中の人しか対応できない、ということでフォーラムに連絡することにした。
Issue報告時のガイドラインによると最小セットの再現モデルを提示した方が早そうだったので、不要なオペレータを削除しまくって再現するモデルを作った。
結論
いくつかディスカッションがあったので、結論をまとめると以下の通り。
- Resizeのサイズ指定やSliceのスライス指定にfloat型のtensorを使えない
- 途中でCastで整数型にしていてもダメ
- TensorRT 8.4 でfloat型tensorに対応される
- JetPack 5.0 DP（TensorRT 8.4.0 EA）で対応されていることを確認済み（2022/05/01 追記）
- ただし、TensorRT 8.4 でもtensorをたどった先がモデル入力だった場合は非対応
- プラグイン出力はResizeのサイズ指定やSliceのスライス指定に使えない（非対応）
補足
ディスカッションおよび途中で教えてもらったドキュメント8.5節の内容をざっくり整理すると以下の通り。
- ReiszeやSliceの第二引数以降の引数はshape tensorとして扱われる
- shape tensorはint32またはboolのみ対応
- shape tensorは0次元または1次元のみ対応
- shape tensorの計算はnetworkのビルド（ディープラーニングモデルを内部のエンジン形式に変換）する際のフェーズ１で計算される
- shape計算はモデルのノードを入力側にたどりながら計算する
- たどる時にshape tensorとして解釈できないノードがあると内部エラーになっているっぽい（想像）
- shape計算はTensorRTがCPUで実行する
- 実際にモデルの推論処理を動かすわけではないため、shape計算には制約がある
