オッサンはDesktopが好き

自作PCや機械学習、自転車のことを脈絡無く書きます

tensorflowのSessionを予め立てて,推論を高速化する

 こんにちは.changです.
今回は,ディープ・ラーニングの推論処理を高速化します.

 先日作った異常検知のビューワーですが*1*2,一度の推論に6秒かかっていました. 何故こんなにかかるのか疑問に思って調べてみると,tensorflowのSessionを立てるのに時間が掛かっていることが判りました. これを改善するテクニックを紹介します.

0. Sessionとは?

 tensorflowには,graphとSessionという独特の概念があります.

 こちらの方*3によると:

  • graph: 計算の定義.これ自身は計算をしないし,値も保持しない.
  • session: グラフやグラフの一部を実行する.グラフを実行する為のリソースを確保し,計算値を保持する.

 恥ずかしながら特に意識することなく,何となくでソースを書いていました. 実際,kerasを使っていると不要です. kerasのインターフェースが標準となるtensorflow 2系では,表面上は排除されるのだと思います.

 tensorflowはPythonのライブラリですので,所謂,Script言語的に動きます. C言語的に考えると,graphはソースコード上の記述,sessionはコンパイルして実際に計算を実行する部分に相当するのだと思います(個人的な解釈が強いかも知れません).

1. 一度目のRunに時間がかかる???

 推論を実際に行う際には,下記の様にsessionをrunします.

inference.cpp

 /* run session */
    TF_SessionRun(sess,
        nullptr, // Run options.
        &input_op, &input_tensor, 1, // Input tensors, input tensor values, number of inputs.
        &out_op, &output_tensor, 1, // Output tensors, output tensor values, number of outputs.
        nullptr, 0, // Target operations, number of targets.
        nullptr, // Run metadata.
        status // Output status.
    );

 このrunに時間が掛かるので,推論ごとに時間が掛かる様に見えていました. 何か変だなぁ...と思い,runをfor分で回してみました. 結果,1回目のRunにだけ6秒かかり,2回目以降は500 ms程度で済むことが判りました. おそらく,1回目のrunでコンパイル済みのexeをメモリ上に展開するような処理をした後は,その領域を使いまわすのだと思います.

 そこで,予め空のデータで一度sessionをrunする関数を書きました.

inference.cpp

int prepareSession(int model_num, int channel)
{
    //  /* prepare input tensor */
    TF_Output input_op = { TF_GraphOperationByName(graph[model_num], "input_1"), 0 };
    if (input_op.oper == nullptr) {
        return -2;
    }

    TF_Tensor* output_tensor = nullptr;

    /* prepare session */
    status = TF_NewStatus();
    options = TF_NewSessionOptions();
    sess[model_num] = TF_NewSession(graph[model_num], options, status);
    TF_DeleteSessionOptions(options);

    if (TF_GetCode(status) != TF_OK) {
        TF_DeleteStatus(status);
        return -3;
    }

    const std::vector<std::int64_t> input_dims = { 1, IMG_SIZE, IMG_SIZE, channel };
    std::vector<float> input_vals(IMG_SIZE*IMG_SIZE*channel); 
    TF_Tensor* input_tensor = tf_utils::CreateTensor(TF_FLOAT,
                                                    input_dims.data(), input_dims.size(),
                                                    input_vals.data(), input_vals.size() * sizeof(float));
    
    /* prepare output tensor */
    TF_Output out_op = { TF_GraphOperationByName(graph[model_num], "conv2d_23/Sigmoid"), 0 };
    if (out_op.oper == nullptr) {
        return -4;
    }
    
    /* run session */
    TF_SessionRun(sess[model_num],
        nullptr, // Run options.
        &input_op, &input_tensor, 1, // Input tensors, input tensor values, number of inputs.
        &out_op, &output_tensor, 1, // Output tensors, output tensor values, number of outputs.
        nullptr, 0, // Target operations, number of targets.
        nullptr, // Run metadata.
        status // Output status.
    );
    return 0;
}

Viewer上からは,Class毎のPrepare sessionボタンを押すことで実行します.

f:id:changlikesdesktop:20200901054543p:plain:w400

2. 結果

 予めSessionをrunしておくことで,Inferenceボタンで実行する推論は500 ms程度で済む様になりました.

f:id:changlikesdesktop:20200901054523p:plain:w400

 アプリを立ち上げるときや,立ち上がった後にバックグラウンドでこっそりsessionをrunしておくと,サクサクと推論が動いている様に見せる事ができます.

3. 限界

 この方法には限界があります. 今回の様に複数のモデルを同時に使いたい場合,モデル毎にsessionを立てる必要があります. つまり,モデル毎に6秒の準備時間が必要になってしまいます. また,メモリを膨大に消費します.

 複数のモデルを使っていると言っても,今回の場合,モデル間でネットワーク構造が完全に同じです. 下記の様に,推論の入り口と出口のノード名も一緒です.

inference.cpp

 /* prepare input tensor */
    TF_Output input_op = { TF_GraphOperationByName(graph[model_num], "input_1"), 0 };
    if (input_op.oper == nullptr) {
        return -2;
    }
 TF_Output out_op = { TF_GraphOperationByName(graph[model_num], "conv2d_23/Sigmoid"), 0 };
    if (out_op.oper == nullptr) {
        return -4;
    }

 理論上,メモリに展開したsessionのパラメタ(weightとbias)をモデル毎に切り替えれば,モデル間でもsessionを使いまわせる筈です. それが出来れば,準備にかかる時間も,メモリ消費も抑えられます. 残念ながら,tensorflowにはそうした機能が無い様です(今回調べた範囲では). 今後の機能向上に期待です.

4. 結び

 Tensorflow 2.0系では,今回紹介した対策が不要になっているのかも知れません. また調べてみようと思います.

 公開済みのソースを更新しています*4