TensorFLow基础：使用TensorFlow读取大数据

问题描述

作为一个优秀的菜鸟，内存爆炸这个坑一定会踩一下的：
以前的训练数据都是几千的样本，直接用numpy读取到内存中，不会出什么问题

今天突然读取一个4万的数据集，我算了一下，我的样本是227 * 227 * 3的，所以一个样本大小大约是 (227 * 227 * 3 * 8)byte = 151KB 左右，4万张图片就是5898MB，相当于6个G了（不知道计算的对不对，难怪我8个G的内存一下就跑满了），4万数据集在模型训练里面并不算多，所以以前的方法一定有问题，所以，今天打算用半天的时间去TensorFlow官网查一下资料研究一下解决方案，顺便记录下来，方便其他菜鸟。

解决方案用TensorFlow的DataSet

原理

• 既然一次性读取这么多数据会内存爆炸，那么我就分批次读取，当训练数据是，才将图片读取到内存中
• 分批次读取，数据就不能很好的打乱了，所以先读取全部图片的地址和对应的标签，然后打乱，然后根据图片地址数据集进行分批次读取
• 新问题，每次读取新图片需要很长的时间，完全满足不了GPU的速度，就是每次迭代就会产生很大的延迟，势必会影响训练速度，所以就是预读取（上次还没训练还没结束，就提前先读取下次迭代需要的数据）

代码

• 图片地址一样全部加入内存中，只是图片地址，加入后打乱，为了代码简洁，突出重点，图片地址加入内存中，并打乱的方法就不再描述了

img_path = get_all_file(path) # 获得所有的图片地址
image_labels = get_lable_by_file_path(path,img_path) #获得图片地址对应的标签

• 新建一个处理图片和加载图片的方法，这2个方法是在训练模型时，对每一轮的样本分别处理

# 处理图片
def preprocess_image(image):
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.image.resize(image, [227, 227])
  image /= 255.0  # normalize to [0,1] range
  return image
# 加载图片
def load_and_preprocess_image(path):
  image = tf.io.read_file(path)
  return preprocess_image(image)

关键代码来了

• 1.将图片路径处理为path_ds （path的DataSet）
• 2.图片处理方法和图片路径（path的DataSet）绑定
• 3.将标签处理为label_ds （标签的DataSet）
• 4.将图片的DateSet和标签的DataSet绑定

AUTOTUNE = tf.data.experimental.AUTOTUNE
path_ds = tf.data.Dataset.from_tensor_slices(img_path) #步骤1
image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)#步骤2
label_ds = tf.data.Dataset.from_tensor_slices(tf.cast(y, tf.int64)) #步骤3
image_label_ds = tf.data.Dataset.zip((image_ds, label_ds))

• 这需要处理一下ds，打乱顺序、重启repeat、样本切片

BATCH_SIZE = 200
image_count = len(img_path)
ds = image_label_ds.apply(
  tf.data.experimental.shuffle_and_repeat(buffer_size=int(image_count/200)))
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)

• 下面就是新建模型和训练模型，就不多说了，我将训练代码贴上

model.fit(ds,
          epochs=200,
          steps_per_epoch=200
         )