时间戳

首页 > 解决方案 > 如何在 TensorFlow 中在大量图像上训练 CNN 而不会耗尽 RAM？

python - 如何在 TensorFlow 中在大量图像上训练 CNN 而不会耗尽 RAM？

问题描述

我正在尝试使用 CNN 对一些图像数据进行分类。我有 38 000 张不同尺寸的图像（比如说大约 400x400 像素）。我最初只使用图像的子集并使用 Open CV 将它们加载到列表中，但现在我尝试使用所有图像，我的 RAM 用完了。在训练过程中处理大量数据的正确方法是什么，我可以批量加载和训练它们吗？如果是这样，怎么做？

我正在使用 Python Jupyter Notebook。

标签： pythonimagetensorflowconv-neural-network

解决方案

对于大型数据集，必须将数据分批读入模型，而不是尝试一次加载所有数据，因为这会导致 OOM（内存不足）错误。由于您正在处理图像，因此我建议使用 ImageDataGenerator().flow_from_directory()。文档在[这里][1]。要使用它，您需要将图像排列到目录和子目录中。例如，假设您有一个包含狗图像和猫图像的数据集，并且您想要构建一个分类器来预测图像是狗还是猫。因此，创建一个名为 train 的目录。在 train 目录中创建一个名为 cat 的子目录和一个名为 dogs 的子目录。将猫的图像放在 cat 目录中，将狗的图像放在 dog 目录中。我通常也会拿一些图像用于测试，所以我还创建了一个名为 test 的目录。在其中创建两个子目录猫和狗，它们的名称与它们在火车目录中的名称相同。将您的测试图像放在 dog 和 cat 目录中。然后使用下面的代码加载数据。

train_dir=r'c:\train'
test_dir=r'c:\test'
img_height=400
imh_width=400
batch_size=32
epochs=20
train_gen=ImageDataGenerator(rescale=1/255, validation_split=.2)
                                     .flow_from_directory( train_dir,
                                      target_size=(img_height, img_width),                                                                              
                                      batch_size=batch_size, seed=123,                                                                               
                                      class_mode='categorical',subset='training' 
                                      shuffle=True)
valid_gen= ImageDataGenerator(rescale=1/255, validation_split=.2)
                                     .flow_from_directory( train_dir,
                                      target_size=(img_height, img_width),                                                                              
                                      batch_size=batch_size, seed=123,                                                                               
                                      class_mode='categorical',subset='validation' 
                                      shuffle=False)
test_gen=ImageDataGenerator(rescale=1/255).flow_from_directory(test_dir,
                                          target_size=(img_height, img_width),                                                                              
                                          batch_size=batch_size,
                                          class_mode='categorical', 
                                          shuffle=False)

然后构建并编译您的模型。使用损失作为 categorical_crossentropy。然后拟合模型

history=model.fit(x=train_gen,  epochs=epochs, verbose=1, validation_data=valid_gen)

这是为了创建验证数据而设置的，因此您可以在训练中监控模型性能。训练完成后，您可以在测试集上测试您的模型

accuracy=model.evaluate( test_gen, batch_size=batch_size, verbose=1, steps=None)[1]*100
print ('Model accuracy on the test set is ' accuracy)

  [1]: https://keras.io/api/preprocessing/image/

推荐阅读