在 C++ 中运行 TensorFlow 时的显着 RAM 成本

Question

我目前正在使用经过训练的 TensorFlow 图在带有 C++ API 的 GPU 机器上编写一些推理代码。

这是我的设置：

• 平台：CentOS 7
• TensorFlow 版本：TensorFlow 1.5
• CUDA 版本：CUDA 9.0
• C++ 版本：C++11

有几个问题我正在努力解决。

1) 首先，我跟随this tutorial 学习了在C++ 中加载图形的基本模板。本教程中的示例非常简单，但是当我运行该程序时（在 GPU 机器上），该程序占用了几乎 0.9G 的 RAM。

2) 我的图表比那个教程中的图表要复杂得多。大约有 20 层，层中的节点数从 300 到 5000 不等。

我的（伪）代码 sn-p 在这里。为简单起见，我只保留导致（潜在）内存问题的代码：

tensorflow::Tensor input = getDataFromSomewhere(...);
int length = size of the input;
int g_batch_size = 50;

// 1) Create session...
// 2) Load graph...

// 3) Inference
for (int x = 0; x < length; x += g_batch_size) {

    tensorflow::Tensor out;
    auto cur_slice = input.Slice(x, std::min(x + g_batch_size, length));

    inference(cur_slice, out);

    // doSomethingWithOutput(out);
}

// 4) Close session and free session memory


// Inference helper function
tensorflow::Status inference(tensorflow::Tensor& input_tensors, tensorflow::Tensor& out) {

    // This line increases a lot more memory usage
    TensorDict feed_dict = {{"IteratorGetNext:0", input_tensors}};
    std::vector<tensorflow::Tensor> outputs;

    tensorflow::Status status = session->Run(feed_dict, {"final_dense:0"}, {}, &outputs);

    // UpdateOutWithOutputs();

    return tensorflow::Status::OK();
}

创建会话并加载图表后，内存成本约为 1.2G。

然后，正如我在代码中指出的那样，当程序到达 session->Run(...) 时，内存使用量上升到超过 2G。

我不确定这是否是 TensorFlow 的正常行为。我检查了this 和this 线程，但我不太清楚我是否在我的代码中创建了冗余操作。

感谢任何 cmets 或建议！提前致谢！

Answer 1

我发现的问题是Tensorflow 动态库将占用大约 200MB 而CUDA 动态库将占用超过 500MB 内存。所以加载这些库已经占用了大量的内存。