TensorFlow 2.0 意外OOM

Question

我正在尝试使用 tensorflow.keras 训练我的模型，但由于 OOM，它在经过一些时期后失败。 Tensorflow 2.0 已将许多内容标记为已弃用，我不知道应该如何诊断问题。

该网络是一系列 Conv1D 层和一些从一个序列转换为另一个序列的自注意力层。序列是可变长度的，但序列长度与失败时间之间没有相关性。 IE：它可以处理 6 分钟的序列，但在 4 分钟的序列中失败。

with tensorflow.device('/device:gpu:0'):
    m2t = BuildGenerator() #builds and returns model
    m2t.compile(optimizer='adam', loss='mse')
    for epoch in range(1):
        for inout in InputGenerator(params):
            m2t.train_on_batch(inout[0], inout[1])

我尝试过的事情：

1. 删除自注意力层。它仍然失败
2. 移除除少数层之外的所有层。它仍然失败
3. 将所有序列填充到恒定长度。它仍然失败
4. 使用 m2t.predict(inout[0]) 代替 train_on_batch。它失败了，但需要更长的时间。
5. 使用 tensorflow.summary.trace_export。它记录了一些东西，但它不会像页面HERE 所建议的那样在 chrome 中加载。
6. 我查看了THIS 的答案，但是随着 TF-2.0 的变化，我不确定如何正确地做到这一点。

没有其他对 tensorflow 或 keras 的调用。

编辑：根据要求，示例错误日志。每次都略有不同的错误。

其中一些，中间有一些成功的运行。

W tensorflow/core/common_runtime/bfc_allocator.cc:239] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.06GiB with freed_by_count=0. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory were available.

然后从这个开始，还有一个巨大的列表“# chunks of size ...”和“InUse...”

W tensorflow/core/common_runtime/bfc_allocator.cc:419] Allocator (GPU_0_bfc) ran out of memory trying to allocate 43.26MiB (rounded to 45360128).  Current allocation summary follows.
I tensorflow/core/common_runtime/bfc_allocator.cc:869] Bin (256):   Total Chunks: 79, Chunks in use: 79. 19.8KiB allocated for chunks. 19.8KiB in use in bin. 2.2KiB client-requested in use in bin.
...
I tensorflow/core/common_runtime/bfc_allocator.cc:921] Sum Total of in-use chunks: 8.40GiB
I tensorflow/core/common_runtime/bfc_allocator.cc:923] total_region_allocated_bytes_: 9109728768 memory_limit_: 9109728789 available bytes: 21 curr_region_allocation_bytes_: 17179869184
I tensorflow/core/common_runtime/bfc_allocator.cc:929] Stats:
    Limit:                  9109728789
    InUse:                  9024084224
    MaxInUse:               9024084224
    NumAllocs:                   38387
    MaxAllocSize:           1452673536

W tensorflow/core/common_runtime/bfc_allocator.cc:424] 


W tensorflow/core/framework/op_kernel.cc:1622] OP_REQUIRES failed at cwise_ops_common.cc:82 : Resource exhausted: OOM when allocating tensor with shape[1,45000,12,21] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
Traceback (most recent call last):
  File ".\TrainGNet.py", line 380, in <module>
    m2t.train_on_batch(inout[0], inout[1])
  File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\keras\engine\training.py", line 973, in train_on_batch
    class_weight=class_weight, reset_metrics=reset_metrics)
  File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\keras\engine\training_v2_utils.py", line 264, in train_on_batch
  output_loss_metrics=model._output_loss_metrics)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\keras\engine\training_eager.py", line 311, in train_on_batch
  output_loss_metrics=output_loss_metrics))
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\keras\engine\training_eager.py", line 268, in _process_single_batch
  grads = tape.gradient(scaled_total_loss, trainable_weights)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\eager\backprop.py", line 1014, in gradient
  unconnected_gradients=unconnected_gradients)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\eager\imperative_grad.py", line 76, in imperative_grad
  compat.as_str(unconnected_gradients.value))
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\eager\backprop.py", line 138, in _gradient_function
  return grad_fn(mock_op, *out_grads)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\ops\math_grad.py", line 251, in _MeanGrad
  return math_ops.truediv(sum_grad, math_ops.cast(factor, sum_grad.dtype)), None
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\util\dispatch.py", line 180, in wrapper
  return target(*args, **kwargs)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\ops\math_ops.py", line 1066, in truediv
  return _truediv_python3(x, y, name)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\ops\math_ops.py", line 1005, in _truediv_python3
  return gen_math_ops.real_div(x, y, name=name)
File "C:\Program Files (x86)\Microsoft Visual Studio\Shared\Python36_64\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py", line 7950, in real_div
  _six.raise_from(_core._status_to_exception(e.code, message), None)
File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[1,45000,12,21] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:RealDiv] name: truediv/

EDIT2 和 3： 这是一个最小的示例。为我打印“11”后失败。 Edit3 显着减小了尺寸。

from tensorflow.keras.models import Model
from tensorflow.keras.layers import *
import tensorflow.keras.backend as K
import numpy as np
import tensorflow

def BuildGenerator():
    i = Input(shape=(None,2,))

    n_input = 12*21
    to_n = Input(shape=(n_input))
    s_n = Dense(12*21, activation='softmax')(to_n)
    s_n = Reshape((12,21))(s_n)
    n_base = Model(inputs=[to_n], outputs=[s_n])

    b = Conv1D(n_input, 11, dilation_rate=1, padding='same', activation='relu', data_format='channels_last')(i)
    n = TimeDistributed(n_base)(b)

    return Model(inputs=[i], outputs=[n])

def InputGenerator():
    for iter in range(1000):
        print(iter)
        i = np.zeros((1,10*60*1000,2))
        n = np.zeros((1,10*60*1000,12,21))
        yield ([i], [n])

with tensorflow.device('/device:gpu:0'):

    m2t = BuildGenerator()

    m2t.compile(optimizer='adam', loss='mse')

    for epoch in range(1):
        for inout in InputGenerator():
            m2t.train_on_batch(inout[0], inout[1])

Answer 1

这是一个内存分配问题，TensorFlow 试图将带有权重的整个模型图分配给 GPU，但 GPU 的 GDDR RAM 不足以容纳大型模型及其权重。 （TensorFlow 需要更多的内存来分配层，而 CNN 也需要更多的内存）

如果您有足够的 CPU RAM 来加载模型，您可以尝试使用 CPU或缩小模型的形状和大小。内存部分也有帮助。

您还可以对 GPU 进行分割，以便模型被 Tensorflow 切割并一个接一个地循环训练。

使用 watch nvidia-smi 跟踪模型的 Nvidia GPU 内存使用情况以优化模型。

Answer 2

你可以尝试使用：

tf.config.gpu.set_per_process_memory_fraction(0.5)
tf.config.gpu.set_per_process_memory_growth(True)

在 TF-2.0 中。请记住在任何操作之前先删除这些。我的意思是您可以简单地将它们添加到代码的开头。

Answer 3

我的简单推荐：

• 将批量大小减小到最小值，从 1 开始，然后增大 大小。

在大多数情况下，这会有所帮助。