从 tensorflow 脚本中捕获 CUDA_ERROR_OUT_OF_MEMORY答案

【问题标题】：Catch CUDA_ERROR_OUT_OF_MEMORY from a tensorflow script从 tensorflow 脚本中捕获 CUDA_ERROR_OUT_OF_MEMORY
【发布时间】：2021-06-25 13:29:19
【问题描述】：

当你想训练一个神经网络时，你需要设置一个批量大小。批大小越大，GPU 内存消耗就越高。当你缺乏 GPU 内存时，tensorflow 会引发这种消息：

2021-03-29 15:45:04.185417: E tensorflow/stream_executor/cuda/cuda_driver.cc:825] failed to alloc 8589934592 bytes on host: CUDA_ERROR_OUT_OF_MEMORY: out of memory
2021-03-29 15:45:04.229570: E tensorflow/stream_executor/cuda/cuda_driver.cc:825] failed to alloc 7730940928 bytes on host: CUDA_ERROR_OUT_OF_MEMORY: out of memory
2021-03-29 15:45:10.776120: E tensorflow/stream_executor/cuda/cuda_driver.cc:825] failed to alloc 17179869184 bytes on host: CUDA_ERROR_OUT_OF_MEMORY: out of memory
...

解决方案是减小批量大小。当我收到此消息时，我希望能够捕获此异常，因此我可以向视图发送消息，甚至自动减小批量大小以自动化学习行为。就我而言，内存不足来自数据集的加载：

try:
  features, labels = iter(input_dataset).next()
except:
  print("this is my exception") 
  raise

但是，cuda 错误 oom 似乎无法像这样捕获。实际上，我认为错误已经在 tf.Dataset 类的 next 函数中被捕获。我看到的似乎实际上是由 oom 错误捕获生成的日志。我不知道如何检测此日志以便对 oom 事件做出反应。

【问题讨论】：

这能回答你的问题吗？ stackoverflow.com/questions/64900712/…
@rok 我尝试过类似的方法，但它不起作用。我认为这是因为错误已经被 tf.Dataset 类捕获并记录了。我根据我的发现修改了我的问题。
您可以引发资源错误，例如：except tf.errors.ResourceExhaustedError as e:

标签： tensorflow deep-learning neural-network tensorflow2.0 object-detection

【解决方案1】：

应用调用的tf.compat.v1.Dataset的next()方法：

iter(my_dataset).next()

已经捕获 OOM 错误。然后，它只是在 stderr 通道中记录错误后尝试生成下一批。您无法自己捕获 OOM 错误，因为 tensorflow api 已经做到了。

不过，您可以通过读取 stderr 来跟踪错误。就我而言，我是以这种方式在命令行中启动我的学习脚本：

process = subprocess.Popen('py -u train.py')

所以我只需要把它改成：

process = subprocess.Popen('py -u train.py', stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

为了将标准错误重定向到标准输出，然后解析标准输出：

            while True:
              output = process.stdout.readline()
                if output == '' and process.poll() is not None:
                  break
                if output:
                  log_message = output.strip().decode('utf-8')
                  if "CUDA_ERROR_OUT_OF_MEMORY" in log_message:
                    process.kill()
                    print("please decrease batch_size")
                    break

【讨论】：