XNNPACKを直接使ってみた

2022/10/27のDeNA/MoT AI技術共有会の発表資料です。

Transcript

1. 2022.10.27 nagiss 株式会社ディー・エヌ・エー ＋ 株式会社 Mobility Technologies XNNPACK を 直接使ってみた
2. 2 項目 01｜背景 02｜XNNPACK とは 03｜直接書いてみる 04｜補足
3. 3 01 背景
4. 4 ▪なんか最近よく見た気がするやつ ▪ ONNX で if を使うと実行時の分岐が選ばれてしまうので辛い ▪ TFLite だと arctan が使えないので辛い ▪ ONNX-Tensorflow で NCHW → NHWC 変換をすると 大量のゴミ Transpose レイヤーが埋め込まれるので辛い 真に受けないでください
5. 5 ▪ じゃあ中の推論 API 直接叩けばいいのでは (?) 真に受けないでください
6. 6 02 XNNPACK とは
7. 7 ▪Google による CPU での NN 推論ライブラリ XNNPACK is a highly optimized library of floating-point neural network inference operators for ARM, WebAssembly, and x86 platforms. XNNPACK is not intended for direct use by deep learning practitioners and researchers; instead it provides low-level performance primitives for accelerating high-level machine learning frameworks, such as TensorFlow Lite, TensorFlow.js, PyTorch, ONNX Runtime, and MediaPipe. XNNPACKは、ARM、WebAssembly、およびx86プラットフォーム用の浮動小数点ニューラル ネットワーク推論演算子の高度に最適化されたライブラリです。XNNPACKは、深層学習の実 践者や研究者が直接使用することを意図していません。代わりに、TensorFlow Lite、 TensorFlow.js、PyTorch、ONNX Runtime、MediaPipeなどの高レベルの機械学習フレームワ ークを高速化するための低レベルのパフォーマンスプリミティブを提供しています。 • (おそらく) 特に ARM に特化している XNNPACK とは
8. 8 ▪どのように最適化しているのか？ ▪ アーキテクチャに最適化された高速な行列積計算 ▪ (これだけで 1 つ資料ができる) ▪ Winograd または FFT を使った高速な畳み込みアルゴリズム ▪ (これだけで 1 つ資料ができる) ▪ 間接的な (?) 畳み込み ▪ https://arxiv.org/abs/1907.02129 ▪ Depthwise Conv を特殊処理 ▪ 良い資料が他にたくさんあるので詳細は割愛 ▪ https://qiita.com/yasudakn/items/f9f906bd21262d067a14 XNNPACK とは
9. 9 ▪どうやって (直接) 使うのか？ ▪ 良い資料が全く無い ▪ じゃあ自分で資料を作ろう XNNPACK とは
10. 10 03 直接書いてみる
11. 11 ▪これだけ ▪依存先含め ./build に静的ライブラリがビルドされる ライブラリのビルド git clone https://github.com/google/XNNPACK.git cd XNNPACK sudo apt install ninja-build scripts/build-local.sh
12. 12 ▪XNNPACK は C/C++ から使うことができる ▪include/xnnpack.h のコメントが詳しい ▪ これとテストコードとかを見て参考にしながら書く ▪環境 ▪ Ubuntu 20.04 LTS、x86-64 ▪目標: 1 層の Conv2d の推論をする ▪ Conv2d の愚直実装と突き合わせて計算結果を確認する ▪(結論を先に書くと) 意外と難しいものではなかった コードを書く
13. 13 ▪main 関数にべた書きしていく コードを書く #include <algorithm> #include <functional> #include <iostream> #include <limits> #include <random> #include <vector> #include <xnnpack.h> int main() {
14. 14 ▪Conv2d のパラメータを用意する コードを書く // パラメータ設定 const size_t batch_size = 1; const size_t input_height = 28; const size_t input_width = 28; const size_t kernel_height = 3; const size_t kernel_width = 3; const size_t padding_left = 1; const size_t padding_top = 1; const size_t padding_right = 1; const size_t padding_bottom = 1; const size_t stride = 1; const size_t dilation = 1; const size_t groups = 1; const size_t group_input_channels = 128; // in_channels / groups const size_t group_output_channels = 128; // out_channels / groups
15. 15 ▪必要な値を計算する コードを書く const size_t out_channels = groups * group_output_channels; const size_t in_channels = groups * group_input_channels; const size_t effective_kernel_height = (kernel_height - 1) * dilation + 1; const size_t effective_kernel_width = (kernel_width - 1) * dilation + 1; const size_t padding_height = padding_top + padding_bottom; const size_t padding_width = padding_left + padding_right; const size_t output_height = (input_height + padding_height - effective_kernel_height) / stride + 1; const size_t output_width = (input_width + padding_width - effective_kernel_width) / stride + 1;
16. 16 ▪入出力と重みのテンソルを用意してランダムな値を入れる コードを書く // 乱数生成器 std::random_device random_device; auto rng = std::mt19937(random_device()); auto f32rng = std::bind(std::uniform_real_distribution<float>(0.0f, 1.0f), std::ref(rng)); // input, kernel, bias のテンソルを用意してランダムな値を入れる std::vector<float> input(batch_size * input_height * input_width * in_channels + XNN_EXTRA_BYTES / sizeof(float)); std::generate(input.begin(), input.end(), std::ref(f32rng)); std::vector<float> kernel(groups * group_output_channels * kernel_height * kernel_width * group_input_channels); std::generate(kernel.begin(), kernel.end(), std::ref(f32rng)); std::vector<float> bias(groups * group_output_channels); std::generate(bias.begin(), bias.end(), std::ref(f32rng)); // 出力テンソルを用意 const size_t output_elements = batch_size * output_height * output_width * out_channels; std::vector<float> output(output_elements); ▪ SIMD 計算の都合で入力テンソルは XNN_EXTRA_BYTES の分だけ余分に確保する必要がある
17. 17 ▪XNNPACK を初期化する ▪allocator の詳細は xnnpack.h に書いてある ▪ allocate, reallocate, deallocate, aligned_allocate, aligned_allocate の 5 つの関数ポインタ持った構造体へのポインタ ▪ nullptr を渡すとデフォルトのアロケータを使ってくれるのでそれで良い コードを書く // XNNPACK の初期化 xnn_initialize(nullptr /* allocator */);
18. 18 ▪operator を作る コードを書く // convolution_op を作る xnn_operator_t convolution_op = nullptr; xnn_create_convolution2d_nhwc_f32( padding_top, // padding_top padding_right, // padding_right padding_bottom, // padding_bottom padding_left, // padding_left kernel_height, // kernel_height kernel_width, // kernel_width stride, // subsampling_height (stride) stride, // subsampling_width (stride) dilation, // dilation_height dilation, // dilation_width groups, // groups group_input_channels, // group_input_channels (in_channels / groups) group_output_channels, // group_output_channels (out_channels / groups) in_channels, // input_channel_stride (in_channels) out_channels, // output_channel_stride (out_channels) kernel.data(), // kernel bias.data(), // bias -std::numeric_limits<float>::infinity(), // output_min (relu -> 0.0f) +std::numeric_limits<float>::infinity(), // output_max 0, // flags NULL, // caches &convolution_op // convolution_op_out );
19. 19 ▪Conv の計算を実行して、後処理をする コードを書く // 入出力のテンソルを指定 xnn_setup_convolution2d_nhwc_f32( convolution_op, batch_size, input_height, input_width, input.data(), output.data(), nullptr /* thread pool */); // conv の実行 xnn_run_operator(convolution_op, nullptr /* thread pool */); // メモリ解放 xnn_delete_operator(convolution_op); convolution_op = nullptr; // メモリ解放 xnn_deinitialize(); ▪thread pool は pthreadpool というライブラリのものを渡す模様 ▪ nullptr を渡せばシングルスレッドで動作するはず
20. 20 ▪(ここまでで XNNPACK のコードは終わり) ▪ わかってしまえば難しいことはなかった コードを書く
21. 21 ▪愚直計算 愚直計算と突き合わせて確認 std::vector<float> naive_output(output_elements); for (auto b = 0; b < batch_size; b++) { for (auto h = 0; h < output_height; h++) { for (auto w = 0; w < output_width; w++) { for (auto g = 0; g < groups; g++) { for (auto o = 0; o < group_output_channels; o++) { for (auto i = 0; i < group_input_channels; i++) { for (auto kh = 0; kh < kernel_height; kh++) { for (auto kw = 0; kw < kernel_width; kw++) { const auto ih = h + kh - padding_top; const auto iw = w + kw - padding_left; if (ih < 0 || ih >= input_width || iw < 0 || iw >= input_width) continue; naive_output[(((b * output_height + h) * output_width + w) * groups + g) * group_output_channels + o] += input[(((b * input_height + ih) * input_width + iw) * groups + g) * group_input_channels + i] * kernel[(((g * group_output_channels + o) * kernel_height + kh) * kernel_width + kw) * group_input_channels + i]; } } } naive_output[(((b * output_height + h) * output_width + w) * groups + g) * group_output_channels + o] += bias[g * group_output_channels + o]; } } } } }
22. 22 ▪突き合わせる 愚直計算と突き合わせて確認 // 計算結果の最初の 5 つを確認 for (auto i = 0; i < 5; i++) { std::cout << "xnnpack:" << output[i] << " naive:" << naive_output[i] << std::endl; } // 差の最大値を確認 auto max_diff = 0.0f; for (auto i = 0; i < output_elements; i++) max_diff = std::max(max_diff, std::abs(output[i] - naive_output[i])); std::cout << "maxdiff=" << max_diff << std::endl; } // end main
23. 23 ▪コンパイル コードを書く g++ -c example.cpp -I../include -I../build/local/pthreadpool-source/include g++ -L../build/local -L../build/local/pthreadpool -L../build/local/cpuinfo -L../build/local/clog example.o ¥ -lXNNPACK -lpthreadpool -lpthread -lcpuinfo -lclog ▪実行結果 xnnpack:127.698 naive:127.698 xnnpack:125.764 naive:125.764 xnnpack:127.382 naive:127.382 xnnpack:126.214 naive:126.214 xnnpack:131.696 naive:131.696 maxdiff=0.000823975 ▪ 一致した！
24. 24 04 補足
25. 25 ▪さっきのコードのそれぞれの推論時間を雑に計測 ▪ batch_size=32 に変更 ▪ -Ofast -march=native を付けてコンパイル ▪ -march=native はあまり効いていなそうだった ▪ 結果: ▪ Naive -> 3.2 sec ▪ XNNPACK -> 0.06 sec ▪計算グラフを作って効率よく NN 全体を計算する機能 ▪ 力尽きたのでそこまで調べられなかった 補足