Ubuntu重启后直接运行TensorFlow C++应用出现CUDA分配错误

问题场景 #

你在Ubuntu上通过Bazel构建TensorFlow C++应用时全程正常，构建命令如下：

bazel build -c opt --local_resources 3072,4.0,1.0 --verbose_failures --config=monolithic //tensorflow/app/...

构建完成后执行bazel-bin/tensorflow/app/run也能顺畅运行，但系统重启后直接启动这个二进制文件，却频繁抛出CUDA内存分配错误：

May 11 12:29:05 tegra-ubuntu start.sh[800]: 2018-05-11 12:29:04.622723: E tensorflow/stream_executor/cuda/cuda_driver.cc:967] failed to alloc 2304 bytes on host: CUDA_ERROR_UNKNOWN 
May 11 12:29:05 tegra-ubuntu start.sh[800]: 2018-05-11 12:29:04.622849: E tensorflow/stream_executor/cuda/cuda_driver.cc:967] failed to alloc 2304 bytes on host: CUDA_ERROR_UNKNOWN

你发现只要重新执行一次Bazel构建再运行就恢复正常，但构建耗时太长，希望系统启动后能直接运行应用，无需重复构建。

解决方案 #

这个问题的核心是TensorFlow默认会在启动时尝试分配全部GPU显存，而系统重启后GPU显存的初始化状态和构建时存在差异，导致分配失败。你可以通过修改代码，给TensorFlow Session设置GPU显存动态增长的配置，让显存按需分配，就能彻底解决这个问题，而且不需要重新构建。

修改你的Session初始化代码，加入以下配置：

SessionOptions opts;
opts.config.mutable_gpu_options()->set_allow_growth(true);
session->reset(tensorflow::NewSession(opts));

设置allow_growth=true后，TensorFlow会根据实际运行需求逐步申请GPU显存，而不是一次性抢占全部可用显存，这样系统重启后直接运行二进制文件也不会再出现内存分配错误了。

内容的提问来源于stack exchange，提问作者marc_