如何将报告_tensor_allocations_upon_oom添加到Keras中的RunOptions

Question

我正在尝试使用Keras在GPU上训练神经网络并且我得到"资源耗尽:分配张量时的OOM"错误.它试图分配的特定张量不是很大,所以我假设一些先前的张量消耗了几乎所有的VRAM.错误消息附带一个提示,表明:

提示:如果要在OOM发生时查看已分配的张量列表,请将report_tensor_allocations_upon_oom添加到RunOptions以获取当前分配信息.

这听起来不错,但我该怎么做？RunOptions似乎是一个Tensorflow的东西,我能找到的小文档将它与"会话"联系起来.我正在使用Keras,所以Tensorflow隐藏在一个抽象层下,其会话在另一层下面.

如何挖掘所有内容以设置此选项以使其生效？

Answer 1

它看起来并不难，您需要知道的是，根据文档，传递给的** kwargs参数model.compile将传递给session.run

因此，您可以执行以下操作：

import tensorflow as tf
run_opts = tf.RunOptions(report_tensor_allocations_upon_oom = True)

model.compile(loss = "...", optimizer = "...", metrics = "..", options = run_opts)

并且每次session.run调用时都应直接传递它。