从tensorflow脚本中捕获CUDA_ERROR_OUT_OF_MEMORY错误

Question

当你想要训练神经网络时，你需要设置一个批次大小。批次大小越大，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类的下一个函数中被捕获。我看到的似乎实际上是由OOM错误捕获生成的日志。我不知道如何检测此日志以便对OOM事件作出反应。

- DeepProblems

---

这个回答解决了你的问题吗？https://stackoverflow.com/questions/64900712/how-to-write-a-contextmanager-to-throw-and-catch-errors - rok

@rok 我尝试过类似的方法，但它并不起作用。我认为这是因为错误已经被 tf.Dataset 类捕获并记录了。根据我发现的情况，我修改了我的问题。 - DeepProblems

你可以像这样引发资源错误：except tf.errors.ResourceExhaustedError as e: - Bijay Regmi

Answer 1

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

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

- DeepProblems