色々と試行錯誤した記憶が薄れてきているので、忘れないうちに書き留めておく。ここでのpytorch/libtorchのversionは1.8.1。 推論にはlibtorchを使うものとする。
1. TorchScriptを使わない
pytorchで書いたモデルをc++に持っていくには一旦TorchScriptに吐いてc++で読み込むのが一般的だと思われる。( 参考: Loading a TorchScript Model in C++ — PyTorch Tutorials 1.10.0+cu102 documentation ) しかし、これを行うと時々推論が遅くなるという問題に遭遇する。 具体的には普通20ms程度で推論できるはずが、なぜかたまに1secかかったりして困る。 そこでpytorchで書いたモデルをc++で書き直すと安定して速い。 もう少し頑張ってくれTorchScript。
2. tensorを作るときはまとめてから
以下のように小さいtensorをたくさん作って後でがっちゃんこする処理はめちゃくちゃ遅い。
std::vector<std::vector<float>> data; std::vector<torch::Tensor> vtensor; for(auto&& d : data) { vtensor.emplace_back(torch::from_blob(d.data(), {1, N}, torch::kFloat32)); } auto batch = torch::vstack(vtensor);
そこで、以下のようにstd::vector<std::vector<float>>を一旦std::vector<float>に一列にして、一度にtorch::from_blobするとちゃんと速くなる。
std::vector<float> data;
torch::Tensor batch = torch::from_blob(d.data(), {B, N}, torch::kFloat32);
3. boolean indexingを使わない
boolean indexingというのは[True, False, False, True, True, ...]のような配列でindexを保持するやつ。
これを使って
x = torch.randn(3, 4) mask = x.ge(0.5) x[mask] # とか torch.masked_select(x, mask) # とか
のように要素を抜き出すことはよく行われる処理である*1。 同じようにc++で
x.index({"...", mask})
などとすると、この処理を何度も繰り返すような場合には遅くなる。なぜかというと以下の箇所で
boolean indexingから普通のindexingに変換されていて、その処理の中で、
毎回aten::emptyでtensorを作り直しているからである。 memory poolからあてがってくれているとは思うが、手元の処理ではこれがbottleneckとなったため、boolean indexingは一度普通のindexingに変換してから使い回すようにした。
以上。参考になれば幸いです。
