使用 python 生成器进行输入时 keras 出现 OOM 错误

OOM error in keras while using a python generator for input

所以我使用了相同的自动编码器模型,批量大小为 10,没有生成器(通过将元素加载到内存中),并且模型运行没有任何问题。

我已经定义了一个 python 生成器,这样我就可以通过以下方式获取更多数据:-

    from sklearn.utils import shuffle
def nifti_gen(samples, batch_size = 5):
  num_samples = len(samples)
  while True:
    for bat in range(0,num_samples,batch_size):
      temp_batch = samples[bat:bat+batch_size]
      batch_data = []
      batch_data = np.asarray(batch_data)

      for i,element in enumerate(temp_batch):
        temp = get_input(element)
        if i == 0:
          batch_data = temp
        else :
          batch_data = np.concatenate((batch_data,temp))
      yield batch_data,batch_data

from sklearn.model_selection import train_test_split
train_samples, validation_samples = train_test_split(IO_paths[:400], test_size=0.1)

train_generator = nifti_gen(train_samples, batch_size=5)
validation_generator = nifti_gen(validation_samples, batch_size=5)

然而,当我尝试训练模型时,我什至在完成一个纪元之前就收到以下错误:-

autoencoder_train = MRA_autoencoder.fit(train_generator, steps_per_epoch= 36 , callbacks= [es,mc] , epochs= 300)
Epoch 1/300
---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-32-65838b7c908e> in <module>()
----> 1 autoencoder_train = MRA_autoencoder.fit(train_generator, steps_per_epoch= 36 , callbacks= [es,mc] , epochs= 300)

8 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

ResourceExhaustedError:  OOM when allocating tensor with shape[500,84,400,400] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[node functional_5/functional_1/conv2d/Conv2D (defined at <ipython-input-32-65838b7c908e>:1) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
 [Op:__inference_train_function_4760]

Function call stack:
train_function

我不知道为什么会发生这种情况,因为我确信我绝对有足够的内存来处理至少 10 个批处理大小。任何帮助将不胜感激! 谢谢

看起来数据量很大。 [500, 84, 400, 400] 是一个非常非常大的数据要处理,在每一层也是如此,最好是恢复到 5 的批量大小或移动到基于多 GPU 云的训练。