用形状分配张量时的OOM - 如何获得更多的GPU内存

Question

[在 Jupyter Lab 环境中运行] 在 TensorFlow 上训练我的 CNN 时：

 history = model.fit(
        train_generator,
        steps_per_epoch=3,
        epochs=5,
        verbose = 1,

当我运行我的算法时，我得到一个'OOM when allocating tensor with shape'。

据我了解，这意味着我没有耗尽足够的 GPU 内存。如何连接 Jupyter 上的服务器以访问更多内存来运行我的训练 NN？

我正在使用以下包和代码来加载图像：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Conduct pre-processing on the data to read and feed the images from the directories into the CNN

# Re-scale data as pixels have value of 0-255
train_datagen = ImageDataGenerator(rescale=1/255)
validation_datagen = ImageDataGenerator(rescale=1/255)

# Feed training dataset images in via batches of 250
train_generator = train_datagen.flow_from_directory (
    'Users\cats-or-dogs\PetImages', # Directory with training set images
    target_size=(300, 300), # Re-size target images
    batch_size = 425, #mini-batch of 250 to make CNN more efficient
    class_mode = 'binary'
)

Answer 1

请让我知道它是否有效。 通常我们可以在导入必要的包后启用mixed-precision，如下所示。它允许更快的计算并且消耗更少的 GPU 内存。因此，我们也可以增加批量大小。但是硬件应该支持这样的设施，所以请先检查它们。 Keras 混合精度 (mp) API 在TensorFlow 2.x 中可用。 开个玩笑，如果你想获得更多的 GPU 内存，那就添加更多的 GPU。因此，您将进行 multi-gpu 训练。但是要使用单个 gpu，mp 是技巧之一。否则，减小batch size可能会解决OOM问题。

policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)

引用官方文档。在 GPU 上使用混合精度时的性能提示。

增加批量大小

如果不影响模型质量，请尝试在使用mixed-precision 时以双 batch size 运行。由于float16 张量使用一半的内存，这通常允许您在不耗尽内存的情况下将批量大小增加一倍。增加批量大小通常会增加训练吞吐量，即您的模型每秒可以运行的训练元素数。

---

此外，我们还可以在每个 epoch 之后使用gc.collect() 来收集垃圾，这将释放一些内存空间，见下文。还有del 未使用的大变量，可能会消耗合理的内存空间。

import tensorflow as tf
import gc

class RemoveGarbaseCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    gc.collect()
...
...
model.fit(train_generator, ...
callbacks=[RemoveGarbaseCallback()])

但是，我们可以在使用tf.keras 的同时使用clear_session()，这将清理所有内容。如果我们在循环中创建模型，这是推荐。因此，我们可以在每次迭代时使用以下代码 sn-p。

for _ in range(no_of_iteration):
   # With `clear_session()` called at the beginning,
   # Keras starts with a blank state at each iteration
   # and memory consumption is constant over time.
   tf.keras.backend.clear_session() # Resets all state generated by Keras

   train_generator = ...
   valid_generator = ...
   
   model =  create_model()
   history = model.fit(.., callbacks=[RemoveGarbaseCallback()])

   # free up some memory space
   del model
   del train_set, valid_set

更新

如你所见：

UnidentifiedImageError: 
cannot identify image file <_io.BytesIO object at 0x0000019F9BC1E950> 

当培训目录中可能有一些不受支持的文件时，就会发生这种情况。要检查文件格式，请运行以下函数：

from collections import Counter
import os
def IMG_EXTENTION(img_path):
    extension_type = []
    file_list = os.listdir(img_path)
    
    for file in file_list: extension_type.append(file.rsplit(".", 1)[1].lower())
        
    print(Counter(extension_type).keys())
    print(Counter(extension_type).values())
    
train_dir = './images' # directory that contains training samples 
IMG_EXTENTION(img_path=train_dir)

在这种情况下，作为方面，它应该包含图像文件格式，即：jpg、jpeg、png 等。现在的问题是在 jupyter 环境下工作时，它自动保存.ipynb 检查点。因此，在您的情况下，它可能与其他图像文件一起保存到培训目录中。这不受支持。在这种情况下，您所要做的就是更改项目目录或更改保存位置。一些指针：1，2

如果您使用的是自定义数据生成器，我建议您使用try 和except 绕过不受支持的文件。同样在flow_from_dataframe而不是flow_from_directory中，我们可以专门传递x_col="id"和y_col="label"，在这种情况下我们可能不会遇到这样的问题。