tensorflow - 从 tensorflow 脚本中捕获 CUDA_ERROR_OUT_OF_MEMORY

Question

当你想训练一个神经网络时，你需要设置一个batch size。批量越大，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
...

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

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

但是，cuda 错误 oom 似乎无法像这样捕获。实际上，我认为该错误已经在 中被捕获了。下一个 tf.Dataset 类的功能。我看到的似乎实际上是由 oom 错误捕获生成的日志。我不知道如何检测此日志以对 oom 事件使用react。

Answer 1

下一个 () 的 tf.compat.v1.Dataset 方法，通过应用调用:

iter(my_dataset).next()

已经捕捉到 OOM 错误。然后，它只是在 stderr channel 中记录错误后尝试生成下一批。您无法自己捕获 OOM 错误，因为 tensorflow api 已经做到了。
不过，您可以通过读取 stderr 来跟踪错误。就我而言，我以这种方式在命令行中启动我的学习脚本:

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

所以我只需要把它改成:

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

为了将 stderr 重定向到 stdout，然后解析 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