我在一个计算资源共享的环境中工作,也就是说,我们有几台服务器机器,每台机器都配备了一些 Nvidia Titan X GPU。
对于中小型模型,Titan X 的 12 GB 通常足以让 2-3 人在同一个 GPU 上同时运行训练。如果模型足够小,以至于单个模型不能充分利用 GPU 的所有计算单元,那么与一个接一个地运行训练过程相比,这实际上会导致加速。即使在并发访问 GPU 确实会减慢个人训练时间的情况下,让多个用户同时在 GPU 上训练的灵活性仍然很好。
TensorFlow 的问题在于,默认情况下,它会在启动时分配全部可用的 GPU 内存。即使对于一个小型的两层神经网络,我看到所有 12 GB 的 GPU 内存都用完了。
有没有办法让 TensorFlow 只分配 4 GB 的 GPU 内存,如果知道这对于给定模型来说已经足够了吗?