オプティム R&D チームの奥村です。今回は、私がどういう風に技術選定してきたかという振り返りです。2018 年は、動画解析ミドルウェアの開発などに取り組んできましたので、その中で学んだことの一部 (TensorRT/Rust/CUDA/thrust) を振り返る形でお伝えししてみます。
学んだこと
TensorRT
深層学習モデルを高速に推論するために、2017 年末から継続的に取り組んできました。背景は以下を参照してくだい。
TensorRT にビルトインされていないカスタムレイヤーの実装や、前処理・後処理の実装、学習済みモデルのインポートで ONNX 形式に対応しました。ONNX については過去記事をどうぞ。
Caffe 形式、MXNet 形式を取り込むために、これらを ONNX に変換する Python スクリプトを用意しました。私がサポートしたいモデルに必要なレイヤーのみのサポートで、完全ではありません。
学び方
情報が少なく、公式のリファレンスとサンプルが主な情報源でした。全体のアーキテクチャーや、GPU メモリの管理方法、カスタムレイヤーの扱いについては熟知しておく必要性を感じたので、社内の Wiki に日本語に訳し、小さいプログラムを書いて理解を深め、それをモジュールとしてまとめていく、という作業を繰り返しました。エンジニアの作業時間の多くは、何かを理解することに費やす必要があるのです。ここで手を抜くと、後で大きなしっぺ返しを食らいます。小さく作って大きく仕上げる（ボトムアップ・アプローチ）のも重要です。いきなり大きく作ろうとすると、トラブルが起きたときにその原因を把握しづらくなり、「ハマる」ことが多いです。小さく作るとそのリスクを最小化できるのです。小さく仕上げたあとは、大きく仕上げるために構成を練り直します（トップダウン・アプローチ）。
2018/05/01 に onnx-tensorrt が登場し、とてもよい情報源となりました。Facebook、NVIDIA の人がコミッターというのが大きいです。NVIDIA のフォーラムで見かけるような質問の多くは、このレポジトリを見れば解決できると思います。TensorFlow が TensorRT に対応しはじめたのはこれよりも早い 2018/01/20 で、現在では TF-TRT としてドキュメントもまとまっています。
モデルが動くようになってからは、深層学習フレームワークとベンチマークしました。私がベンチマークした限りでは、速度面では明らかに TensorRT が優位であることが分かったので、採用することにしました。
検討できなかったこと
カスタムレイヤーを実装するためのインターフェースが変わりつつあり、古いインターフェースが非推奨になりました。TensorRT 5 で古いインターフェースが動作することもあり、新しいインターフェースに移行できませんでした。
行動認識で登場した 3D ResNets for Action Recognition はまだ動かせていません。TensorRT のビルトインレイヤーの多くは、四次元まで （NCHW など）しかサポートしていないので不安が残りますが、シンプルで優れたアプローチだと思うので、対応したいです。
TVM との比較もちゃんとできませんでした。こちらはまだサポートされているレイヤーが少ない印象で、現時点では TensorRT よりもハードルが高いと思います。
C++ (復習)
TensorRT や thrust に触れる上で避けられないので、復習しました。私の知識は C++11 にも到達していませんでした（深層学習系に突入する前の私の主戦場は C++ ではなく C でした）。今更ですが、標準ライブラリは充実してきたな、と感じました。thread、mutex、condition_variable が多環境で動くのは感動しました。しかし、並行プログラミングには辛さを感じました。また、テンプレートによるコンパイル時間の増大は相変わらずでした。私は Delphi 育ちで、ビルドが数秒で終わらないと、とてもストレスを感じてしまいます。
学び方
cpprefjp の情報が充実していたのであまり困りませんでした。困ったときは、チームメンバー (C++ポケットリファレンス の著者を含む) に聞けば解決しました。
Rust
※ Rust 自体を知りたい方は過去記事を参照してください。
動画解析ミドルウェアを作る上で、Rust にするか、Go にするか、C/C++ にするかは悩みました。Go は経験がありましたが、TensorRT と絡ませる上で C++ とやりとりする必要がありましたが、Go の cgo が辛い、というのもありました。書いて動かしてみないと良いところも悪いところも分からないので、書いたことがない Rust で書いてみることにしました。製品開発ではいきなりこういうことは難しいですが、技術選定段階では新しい言語でも積極的に試すようにしています。駄目だったら、書いたことがある Go や C/C++ にもどればいいだけです。採用する技術は、最低でもある程度使ってみて、他と比較してみないと正しい決定ができないと思っています。
最初に書いたのは ONNX の parser でしたが、これで手応えを感じました。Rust は parser や decoder、交換器系のプログラムを書くのに向いていると感じました。
画像や動画を扱うときは C/C++ で書かれたライブラリを利用することが多いので、sys crate を作るのに慣れました。libyuv、libvpx、libwebm の sys crate を作ったり、mozjpeg-sys のラッパーを書きました。
並行処理がしたかったので、crossbeam-channel を使ってメッセージパッシングできるようになりました。このときに Rust の enum の強力さに感動しました。
学び方
困ったときは、チームメンバーに聞けば解決しました。個人的にはコンパイルの速さとパッケージマネージャーの充実がとても気に入りました。コンパイルが遅いと、ちょこちょこ作って動かすサイクルが遅くなり、生産性が落ちてしまうので、重要視しています。
標準ライブラリにチャネルがあったので当初はそれを使っていましたが、Multi-producer, multi-consumer な FIFO キューとして振る舞えるものが欲しかったので crossbeam-channel に移行しました。送信タイムアウトの指定ができるのも魅力の一つです。本家のリファレンスの基本的なところ を社内の Wiki に日本語に訳し、コードを一つずつ実行して理解を深めました。ちなみに、中身が気になったのでチームメンバーと話していたら、サクッとコードを読み始めて futex に辿り着いていました。
sys crate を作る方法は Using C libraries in Rust: make a sys crate がよくまとまっていました。あとは GitHub にある偉大なコードを参照しながら、自分が書くコードに適用していきました。
慣れるまではコンパイルが中々通らないのにやきもきしましたが、コンパイルが通るということである程度保証される安全性の方が重要ということで気にならなくなりました。
もう、C/C++ にはもどれない・・・。
検討できなかったこと
検討できたことの方が少ないです。future とかよく分かっていません。仕様が現時点ではふわふわしているので、それを言い訳にしておきます。
おまけ：パフォーマンス解析
動画解析ミドルウェアを作る上で、期待したとおりに処理が走っているかを確認する必要がありました。マルチスレッドなプログラムは、単純なテキストログだと人間が追っかけるには辛いですよね。視覚的・直感的に理解できるのは大事です。そこでぱっと思いついたのが Chrome Dev Tools の Timeline ツールです。調査してみると Trace Event Format なるものがあり、これに沿った json ファイルを Chrome に読み込ませると期待通りの結果が得られました。
CUDA/thrust
TensorRT でのカスタムレイヤーの実装、前処理・後処理の高速化のために利用を余儀なくされました。CUDA は 2007 年、thrust は 2009 年に登場したもので、どちらも既に 10 年選手であり枯れた技術です。
学び方
CUDA は、なんだかんだで CUDA C プロフェッショナル プログラミング が一番参考になりました。CUDA についてはズブの素人だったので、新鮮で面白かったです。総和やソートなど、ループ伝搬依存 (loop-carried dependency) があるカーネルをバグなしで書ける気がしなかったので、thrust を使うことにしました。thrust はシンプルな利用性が少なかったのですが、Chainer など深層学習フレームワークの実装が貴重な情報源となりました。
CUDA には明示的な同期と暗黙的な同期があるということなので、意図しないボトルネックを作らないように注意しました。これを把握していないと、「なんか思ったよりも速度が出ないな」という状態に陥る可能性があります。CUDA C/C++ Streams and Concurrency に詳細がありますが、基本方針は「メモリ確保と同期関数使うな」ですね。CUDA カーネルを自分で書くときは、これらに注意すれば事足ります。
注意が必要なのは thrust です。例えばソートは基数ソートで実装されていて、一次領域を必要とします。メソッドには外部から一次領域を指定することはできません。しかし、thrust はよくできていて一次領域の確保・破棄にうまく割り込むことができます。C++ の標準ライブラリのように、カスタムアロケーターを仕込むことができるのです。thrust の各種アルゴリズムは第一引数に実行ポリシー (thrust::execution_policy) を指定するオーバーロードを持っており、ここにカスタムアロケーターを関連付けることができるのです。詳細は Thrust/CUDA tip: reuse temporary buffer across multiple transforms にまとめられています。thrust のサンプルコードにもカスタムアロケーターの例があります。
カスタムアロケーターのメモリ割り当て戦略は、Buddy memory allocation などではなく、シンプルで高速なシングルフレームアロケーターを使いました。
検討できなかったこと
実装した CUDA カーネルの起動パラメータのチューニングまでは手を付けられませんでした。Optunaを使ってFFmpegのエンコードパラメータを最適化してみる という例もありますし、Optuna - A hyperparameter optimization framework でやってみたいと思っています。
さいごに
こうして振り返ってみると、まだまだ学ぶべきことはたくさんあります。これからも尽きることはないでしょう。
私は以下の記事で触れられているような REST/Web、RDBMS/SQL の経験はないので、そういうところもやっていきたいと思っています。
オプティムでは、学ぶことが尽きることのない状況を楽しめるエンジニアを募集しています。興味のある方は、こちらをご覧ください。
