Google Colab在使用Tensorflow 2.0批处理数据集时未加载图像文件

时间:2020-03-27 13:00:27

标签: python google-colaboratory tensorflow2.0

有一点背景知识,我正在加载约60,000张图像以协作训练GAN。我已经将它们上传到云端硬盘,并且目录结构在root中包含用于不同类(大约7-8)的文件夹。我将它们加载到colab中,如下所示:

root = "drive/My Drive/data/images"
root = pathlib.Path(root)

list_ds = tf.data.Dataset.list_files(str(root/'*/*'))

for f in list_ds.take(3):
  print(f.numpy())

给出输出:

b'drive/My Drive/data/images/folder_1/2994.jpg'
b'drive/My Drive/data/images/folder_1/6628.jpg'
b'drive/My Drive/data/images/folder_2/37872.jpg'

我正在如下进一步处理它们:

def process_path(file_path):
  label = tf.strings.split(file_path, '/')[-2]
  image = tf.io.read_file(file_path)
  image = tf.image.decode_jpeg(image)
  image = tf.image.convert_image_dtype(image, tf.float32)
  return image#, label

ds = list_ds.map(process_path)

BUFFER_SIZE = 60000
BATCH_SIZE = 128

train_dataset = ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

每个图像的大小为128x128。现在遇到问题了,当我尝试在colab中查看批处理时,执行将永远进行,并且永远不会停止,例如,使用以下代码:

for batch in train_dataset.take(4):
  print([arr.numpy() for arr in batch])

之前,我认为batch_size可能是一个问题,因此尝试对其进行更改,但仍然是同样的问题。加载大量文件时,由于colab会造成问题吗?

还是由于使用MNIST(28x28)时正在工作的图像大小?如果是这样,可能的解决方案是什么?

谢谢。

编辑: 删除shuffle语句后,最后几行会在几秒钟内执行。因此,我认为这可能是由于洗牌的BUFFER_SIZE引起的问题,但是即使BUFFER_SIZE减少了,执行时间仍然很长。有解决方法吗?

1 个答案:

答案 0 :(得分:3)

这是我从个人Google云端硬盘加载1.12GB压缩FLICKR图像数据集的方式。首先,我在colab环境中解压缩数据集。 prefetchautotune是可以提高性能的一些功能。此外,我使用本地colab缓存来存储处理后的图像。第一次执行大约需要20秒(假设您已解压缩数据集)。然后,缓存允许后续调用非常快速地加载。

假设您已授权google drive API,我首先将文件夹解压缩

!unzip /content/drive/My\ Drive/Flickr8k
!unzip Flickr8k_Dataset
!ls

然后,我将您的代码与prefetch()autotunecache file一起使用。

import pathlib
import tensorflow as tf

def prepare_for_training(ds, cache, BUFFER_SIZE, BATCH_SIZE):
  if cache:
    if isinstance(cache, str):
      ds = ds.cache(cache)
    else:
      ds = ds.cache()
  ds = ds.shuffle(buffer_size=BUFFER_SIZE)
  ds = ds.batch(BATCH_SIZE)
  ds = ds.prefetch(buffer_size=AUTOTUNE)
  return ds

AUTOTUNE = tf.data.experimental.AUTOTUNE

root = "Flicker8k_Dataset"
root = pathlib.Path(root)

list_ds = tf.data.Dataset.list_files(str(root/'**'))

for f in list_ds.take(3):
  print(f.numpy())

def process_path(file_path):
  label = tf.strings.split(file_path, '/')[-2]
  img = tf.io.read_file(file_path)
  img = tf.image.decode_jpeg(img)
  img = tf.image.convert_image_dtype(img, tf.float32)
  # resize the image to the desired size.
  img =  tf.image.resize(img, [128, 128])
  return img#, label

ds = list_ds.map(process_path, num_parallel_calls=AUTOTUNE)
train_dataset = prepare_for_training(ds, cache="./custom_ds.tfcache", BUFFER_SIZE=600000, BATCH_SIZE=128)
for batch in train_dataset.take(4):
  print([arr.numpy() for arr in batch])

这是使用keras flow_from_directory()的一种方法。这种方法的好处是避免了tensorflow shuffle(),它取决于缓冲区大小,可能需要处理整个数据集。 Keras为您提供了一个迭代器,您可以调用它来获取数据批,并且内置了随机改组。

import pathlib
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

root = "Flicker8k_Dataset"
BATCH_SIZE=128

train_datagen = ImageDataGenerator(
    rescale=1./255 )

train_generator = train_datagen.flow_from_directory(
        directory = root,  # This is the source directory for training images
        target_size=(128, 128),  # All images will be resized
        batch_size=BATCH_SIZE,
        shuffle=True,
        seed=42, #for the shuffle
        classes=[''])

i = 4
for batch in range(i):
  [print(x[0]) for x in next(train_generator)]