如何防止tensorflow分配整个GPU内存？
10 浏览
如何防止tensorflow分配整个GPU内存？
我在一个计算资源共享的环境中工作，我们有几台服务器配备了几个Nvidia Titan X GPU。对于小到中等大小的模型来说，Titan X的12 GB通常足够2-3个人同时在同一GPU上进行训练。如果模型足够小，单个模型无法充分利用GPU的所有计算单元，这实际上会比依次运行训练过程更快。即使在并发访问GPU会减慢单个训练时间的情况下，同时让多个用户在GPU上进行训练仍然很方便。
TensorFlow的问题在于，默认情况下，在启动时会分配所有可用的GPU内存。即使是一个小型的两层神经网络，我也看到GPU内存的全部12 GB都被使用掉了。
如果知道对于给定的模型来说4 GB足够，是否有办法让TensorFlow只分配4 GB的GPU内存？