时间戳

首页 > 解决方案 > 在 C++ 中运行 TensorFlow 时的显着 RAM 成本

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

问题描述

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

这是我的设置：

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

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

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

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

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 的正常行为。我已经检查了这个和这个线程，但我不太清楚我是否在我的代码中创建了冗余操作。

任何意见或建议表示赞赏！提前致谢！

标签： c++tensorflow

解决方案

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

推荐阅读