我在一个共享计算资源的环境中工作,也就是说,我们有一些服务器机器配备了几个Nvidia Titan X GPU.
对于小到中等大小的型号,12GB的Titan X通常足以让2-3人在同一GPU上同时进行训练.如果模型足够小以至于单个模型没有充分利用Titan X的所有计算单元,那么与在另一个训练过程之后运行一个训练过程相比,这实际上可以导致加速.即使在并发访问GPU确实减慢了单个培训时间的情况下,仍然可以灵活地让多个用户同时在GPU上运行.
TensorFlow的问题在于,默认情况下,它在启动时会在GPU上分配全部可用内存.即使对于一个小的2层神经网络,我也看到12 GB的Titan X已用完.
有没有办法让TensorFlow只分配4GB的GPU内存,如果有人知道这个数量对于给定的模型来说足够了？
我正在sess.run()循环中使用多个tensorflow推论，碰巧有些推论对我的GPU来说太重了。
我收到如下错误：
2019-05-23 15:37:49.582272: E tensorflow/core/common_runtime/executor.cc:623]
Executor failed to create kernel. Resource exhausted: OOM when allocating tensor of shape [306] and type float
我希望能够捕获这些特定的OutOfMemory错误，但不能捕获其他错误（这可能是由于错误的输入格式或损坏的图形所致。）
显然，其结构类似于：
try:
sess.run(node_output, feed_dict={node_input : value_input})
except:
do_outOfMemory_specific_stuff()
不起作用，因为其他类型的错误将导致对该do_outOfMemory_specific_stuff函数的调用。
任何想法如何捕获这些OutOfMemory错误？