Tensorflow2.0加载图片数据集的两种方式
前言使用tf.keras的ImageDataGenerator生成器完整代码 使用tf.data.Dataset完整代码前言
在tensorflow2.0中加载图片数据集一般有两种方式，第一种是使用tf.keras中的ImageDataGenerator生成器，适合图片分类问题，简单但不灵活；第二种是使用tf.data.Dataset搭配tf.image中的一些图片处理方法，较为灵活。下面分别介绍一下。
我们使用热狗数据集，从这里下载
使用tf.keras的ImageDataGenerator生成器
使用tf.keras的ImageDataGenerator生成器要求文件夹下的图片如图所示的形式放置。hotdog文件夹下分别是用于训练和评估的train和test文件夹，这两个文件夹下面均有hotdog和not-hotdog两个类别文件夹，每个类别文件夹里面是图像文件。
首先创建两个tf.keras.preprocessing.image.ImageDataGenerator实例来分别读取训练数据集和测试数据集中的所有图像文件。rescale参数表示将图片数据归一化到(0,1)范围内，此外还有一些其它参数可用于数据增强。
flow_from_directory函数中，directory参数为每个类别文件夹所在的路径，target_size参数表示将图片调整为224×224尺寸。
import pathlibtrain_dir = "../data/hotdog/train"test_dir = "../data/hotdog/test"train_dir = pathlib.Path(train_dir)train_count = len(list(train_dir.glob('*/*.png')))test_dir = pathlib.Path(test_dir)test_count = len(list(test_dir.glob('*/*.png')))CLASS_NAMES = np.array([item.name for item in train_dir.glob('*') if item.is_dir()])image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255)BATCH_SIZE = 32IMG_HEIGHT = 224IMG_WIDTH = 224train_data_gen = image_generator.flow_from_directory(directory=str(train_dir),batch_size=BATCH_SIZE,target_size=(IMG_HEIGHT, IMG_WIDTH),shuffle=True,classes=list(CLASS_NAMES))test_data_gen = image_generator.flow_from_directory(directory=str(test_dir),batch_size=BATCH_SIZE,target_size=(IMG_HEIGHT, IMG_WIDTH),shuffle=True,classes=list(CLASS_NAMES))
Found 2000 images belonging to 2 classes.Found 800 images belonging to 2 classes.
随机选取9张图片可视化一下：
import matplotlib.pyplot as pltdef show_batch(image_batch, label_batch):plt.figure(figsize=(8, 8))for n in range(9):plt.subplot(3, 3, n + 1)plt.imshow(image_batch[n])plt.title(CLASS_NAMES[label_batch[n] == 1][0].title())plt.axis('off')plt.show()image_batch, label_batch = next(train_data_gen)show_batch(image_batch, label_batch)
然后使用ResNet50预训练模型进行微调：
ResNet50 = tf.keras.applications.resnet_v2.ResNet50V2(include_top=False, weights='imagenet', input_shape=(224, 224, 3))ResNet50.trainable=Falsenet = tf.keras.models.Sequential()net.add(ResNet50)net.add(tf.keras.layers.GlobalAveragePooling2D())net.add(tf.keras.layers.Dense(2, activation='softmax'))net.summary()pile(optimizer=tf.keras.optimizers.Adam(),loss='categorical_crossentropy',metrics=['accuracy'])
使用tf.keras内置的fit_generator方法进行训练：
epoch_steps = train_count // BATCH_SIZEval_steps = test_count // BATCH_SIZEnet.fit_generator(train_data_gen,steps_per_epoch=epoch_steps,epochs=5,validation_data=test_data_gen,validation_steps=val_steps)
训练过程：
Epoch 1/562/62 [==============================] - 31s 505ms/step - loss: 0.3595 - accuracy: 0.8415 - val_loss: 0.2036 - val_accuracy: 0.9212Epoch 2/562/62 [==============================] - 24s 391ms/step - loss: 0.2185 - accuracy: 0.9187 - val_loss: 0.1722 - val_accuracy: 0.9325Epoch 3/562/62 [==============================] - 24s 390ms/step - loss: 0.1748 - accuracy: 0.9339 - val_loss: 0.1339 - val_accuracy: 0.9450Epoch 4/562/62 [==============================] - 24s 394ms/step - loss: 0.1634 - accuracy: 0.9334 - val_loss: 0.1269 - val_accuracy: 0.9500Epoch 5/562/62 [==============================] - 24s 391ms/step - loss: 0.1434 - accuracy: 0.9477 - val_loss: 0.1218 - val_accuracy: 0.9488
完整代码
import tensorflow as tfimport numpy as npimport pathlibfor gpu in tf.config.experimental.list_physical_devices('GPU'):tf.config.experimental.set_memory_growth(gpu, True)train_dir = "../data/hotdog/train"test_dir = "../data/hotdog/test"train_dir = pathlib.Path(train_dir)train_count = len(list(train_dir.glob('*/*.png')))test_dir = pathlib.Path(test_dir)test_count = len(list(test_dir.glob('*/*.png')))CLASS_NAMES = np.array([item.name for item in train_dir.glob('*') if item.is_dir()])image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255)BATCH_SIZE = 32IMG_HEIGHT = 224IMG_WIDTH = 224train_data_gen = image_generator.flow_from_directory(directory=str(train_dir),batch_size=BATCH_SIZE,target_size=(IMG_HEIGHT, IMG_WIDTH),shuffle=True,classes=list(CLASS_NAMES))test_data_gen = image_generator.flow_from_directory(directory=str(test_dir),batch_size=BATCH_SIZE,target_size=(IMG_HEIGHT, IMG_WIDTH),shuffle=True,classes=list(CLASS_NAMES))# 使用ResNet50预训练模型进行微调ResNet50 = tf.keras.applications.resnet_v2.ResNet50V2(weights='imagenet', input_shape=(224, 224, 3), include_top=False)ResNet50.trainable=Falsenet = tf.keras.models.Sequential()net.add(ResNet50)net.add(tf.keras.layers.GlobalAveragePooling2D())net.add(tf.keras.layers.Dense(2, activation='softmax'))net.summary()pile(optimizer=tf.keras.optimizers.Adam(),loss='categorical_crossentropy',metrics=['accuracy'])epoch_steps = train_count // BATCH_SIZEval_steps = test_count // BATCH_SIZEhistory = net.fit_generator(train_data_gen,steps_per_epoch=epoch_steps,epochs=5,validation_data=test_data_gen,validation_steps=val_steps)
使用tf.data.Dataset
定义一个加载图片数据集的函数：
AUTOTUNE = tf.data.experimental.AUTOTUNEdef load_data(path, batch_size, epochs):data_root = pathlib.Path(path)all_image_paths = list(data_root.glob('*/*'))all_image_paths = [str(path) for path in all_image_paths]label_names = sorted(item.name for item in data_root.glob('*/') if item.is_dir())label_to_index = dict((label, index) for index, label in enumerate(label_names))all_image_labels = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]image_ds = tf.data.Dataset.from_tensor_slices(all_image_paths) \.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)label_ds = tf.data.Dataset.from_tensor_slices(all_image_labels)image_label_ds = tf.data.Dataset.zip((image_ds, label_ds))image_count = len(all_image_paths)dataset = image_label_ds.shuffle(buffer_size=image_count) \.batch(batch_size=batch_size) \.repeat(epochs) \.prefetch(buffer_size=AUTOTUNE)return dataset, image_counttrain_path = "../data/hotdog/train"test_path = "../data/hotdog/test"BATCH_SIZE = 32EPOCHS = 5ds_train, train_count = load_data(train_path, BATCH_SIZE, EPOCHS)ds_test, test_count = load_data(test_path, BATCH_SIZE, EPOCHS)
在该函数中，首先获取所有图片样本文件的路径all_image_paths，对每一个图片样本，有对应的标签all_image_labels。接着，使用tf.data.Dataset.from_tensor_slices将字符串列表all_image_paths和all_image_labels转化为张量，对于图片，利用map方法映射到一个图片预处理函数load_and_preprocess_image：
def load_and_preprocess_image(path, size=(224, 224)):image = tf.io.read_file(path)image = tf.io.decode_jpeg(image)image = tf.image.resize(image, size) / 255.return image
该函数将图片调整为224×224尺寸，并归一化到(0,1)范围。
接下来，zip方法将图片数据和标签数据压缩成（图片,标签）对。然后依次使用shuffle、batch、repeat、prefetch方法加载数据集。shuffle方法将数据集进行打乱，batch方法将数据集分批，repeat方法将数据集复制epochs份，使用并行化预处理num_parallel_calls和预存数据prefetch来提升性能。
随机选取9张图片可视化一下：
from matplotlib import pyplot as pltplt.figure(figsize=(8,8))for i,(img,label) in enumerate(ds_train.unbatch().take(9)):ax=plt.subplot(3,3,i+1)ax.imshow(img.numpy())ax.set_title("label = %d"%label)ax.set_xticks([])ax.set_yticks([])plt.show()
然后使用ResNet50预训练模型进行微调：
ResNet50 = tf.keras.applications.resnet_v2.ResNet50V2(weights='imagenet', input_shape=(224, 224, 3), include_top=False)ResNet50.trainable = Falsemodel = tf.keras.models.Sequential([ResNet50,tf.keras.layers.GlobalAveragePooling2D(),tf.keras.layers.Dense(2, activation='softmax')])model.summary()pile(optimizer=tf.keras.optimizers.Adam(),loss='sparse_categorical_crossentropy',metrics=['accuracy'])
使用tf.keras内置的fit方法进行训练：
epoch_steps = train_count // BATCH_SIZEval_steps = test_count // BATCH_SIZEmodel.fit(ds_train, epochs=EPOCHS, steps_per_epoch=epoch_steps,validation_data=ds_test, validation_steps=val_steps)
需要注意，这里传入了参数step_per_epoch和validation_steps，表示每次epoch遍历批量数据的次数，那么每次epoch就会对不同份的训练集和验证集进行遍历，因此之前在加载数据集时需要repeat(epochs)份，从而可以对epochs份数据遍历epochs次；
如果不传入step_per_epoch和validation_steps参数的话，那么之前就无需repeat()，keras会在第一次epoch的遍历中自动计算steps_per_epoch值。
训练过程：
Train for 62 steps, validate for 25 stepsEpoch 1/562/62 [==============================] - 16s 255ms/step - loss: 0.3609 - accuracy: 0.8271 - val_loss: 0.1352 - val_accuracy: 0.9450Epoch 2/562/62 [==============================] - 10s 163ms/step - loss: 0.2186 - accuracy: 0.9116 - val_loss: 0.1122 - val_accuracy: 0.9488Epoch 3/562/62 [==============================] - 10s 169ms/step - loss: 0.1927 - accuracy: 0.9278 - val_loss: 0.1075 - val_accuracy: 0.9563Epoch 4/562/62 [==============================] - 9s 138ms/step - loss: 0.1683 - accuracy: 0.9355 - val_loss: 0.1096 - val_accuracy: 0.9563Epoch 5/562/62 [==============================] - 10s 160ms/step - loss: 0.1443 - accuracy: 0.9497 - val_loss: 0.0943 - val_accuracy: 0.9638
完整代码
import tensorflow as tfimport pathlibfor gpu in tf.config.experimental.list_physical_devices('GPU'):tf.config.experimental.set_memory_growth(gpu, True)AUTOTUNE = tf.data.experimental.AUTOTUNEdef load_and_preprocess_image(path, size=(224, 224)):image = tf.io.read_file(path)image = tf.io.decode_jpeg(image)image = tf.image.resize(image, size) / 255.return imagedef load_data(path, batch_size, epochs):data_root = pathlib.Path(path)all_image_paths = list(data_root.glob('*/*'))all_image_paths = [str(path) for path in all_image_paths]label_names = sorted(item.name for item in data_root.glob('*/') if item.is_dir())label_to_index = dict((label, index) for index, label in enumerate(label_names))all_image_labels = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]image_ds = tf.data.Dataset.from_tensor_slices(all_image_paths) \.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)label_ds = tf.data.Dataset.from_tensor_slices(all_image_labels)image_label_ds = tf.data.Dataset.zip((image_ds, label_ds))image_count = len(all_image_paths)dataset = image_label_ds.shuffle(buffer_size=image_count) \.batch(batch_size=batch_size) \.repeat(epochs) \.prefetch(buffer_size=AUTOTUNE)return dataset, image_counttrain_path = "../data/hotdog/train"test_path = "../data/hotdog/test"BATCH_SIZE = 32EPOCHS = 5ds_train, train_count = load_data(train_path, BATCH_SIZE, EPOCHS)ds_test, test_count = load_data(test_path, BATCH_SIZE, EPOCHS)# 使用ResNet50预训练模型进行微调ResNet50 = tf.keras.applications.resnet_v2.ResNet50V2(weights='imagenet', input_shape=(224, 224, 3), include_top=False)ResNet50.trainable = Falsemodel = tf.keras.models.Sequential([ResNet50,tf.keras.layers.GlobalAveragePooling2D(),tf.keras.layers.Dense(2, activation='softmax')])model.summary()pile(optimizer=tf.keras.optimizers.Adam(),loss='sparse_categorical_crossentropy',metrics=['accuracy'])epoch_steps = train_count // BATCH_SIZEval_steps = test_count // BATCH_SIZEmodel.fit(ds_train, epochs=EPOCHS, steps_per_epoch=epoch_steps,validation_data=ds_test, validation_steps=val_steps)
总结一下，tf.data.Dataset各方法的使用顺序：from_tensor_slices -> map -> shuffle -> batch -> repeat -> prefetch
