TensorFlow1.15.5中如何访问tfds.load加载数据集的单张图像

你遇到的几个错误分属不同原因，逐个解决即可：

• f-string语法错误：这是纯Python语法问题，和TensorFlow无关。你在f-string的外层用了单引号，内部取字典key也用了单引号，Python解析时会提前匹配引号导致语法失败。要么把f-string外层改成双引号，要么把内部字典key的引号改成双引号即可。
• 直接打印张量拿不到值：TF1.15默认是图执行模式，你定义的next_val只是计算图里的张量节点，不是实际存储数值的数组，必须通过Session.run()执行计算图才能拿到真实值，直接打印只会输出张量的形状、类型等元信息。
• dataset[image]下标报错：tf.data.Dataset是流式迭代的数据集对象，本身不支持随机下标访问，只能通过迭代器逐批次取数，不能像列表/数组一样直接用下标取元素。
• 拿不到单张图像：你拿到的batch级张量第一维就是batch维度，只要拿到batch对应的numpy数组，沿第一维索引就能拿到单张图。

可用方案 #

方案1：使用TF1.x原生图模式（无需开启eager execution，兼容性最好） #

import tensorflow as tf
import tensorflow_datasets as tfds

# 加载数据集并做batch处理
dataset = tfds.load('cifar10', split='train', shuffle_files=True)
dataset = dataset.batch(128)

# 创建单次迭代器
iterator = dataset.make_one_shot_iterator()
next_batch = iterator.get_next()

# 通过Session执行计算图，拿到真实数值
with tf.Session() as sess:
    batch_dict = sess.run(next_batch)
    # batch_dict['image'] 是形状为(128, 32, 32, 3)的numpy数组，dtype为uint8
    # 沿第0维（batch维）索引即可取到单张图，例如取batch里第0张图
    single_img = batch_dict['image'][0]
    single_label = batch_dict['label'][0]
    print("单张图像形状：", single_img.shape, "对应标签：", single_label)

注意：就算你修正了f-string的引号问题，图模式下直接打印next_val['image']也只能拿到张量节点信息，拿不到实际像素值，必须经过sess.run()才能得到numpy数组。

方案2：开启Eager Execution模式（和TF2.x用法一致） #

如果你习惯动态图的写法，可以在所有TensorFlow操作之前开启eager模式，就不需要手动维护Session：

import tensorflow as tf
import tensorflow_datasets as tfds
# 必须在导入TF后第一时间开启，否则会报错
tf.enable_eager_execution()

dataset = tfds.load('cifar10', split='train', shuffle_files=True)
dataset = dataset.batch(128)

# 直接迭代取数即可
for batch in dataset.take(1):
    # 直接调用.numpy()把eager张量转成numpy数组
    batch_imgs = batch['image'].numpy()
    batch_labels = batch['label'].numpy()
    # 取单张图
    single_img = batch_imgs[0]
    single_label = batch_labels[0]
    print("单张图像形状：", single_img.shape, "对应标签：", single_label)

注意事项 #

• 不要尝试对tf.data.Dataset对象直接做下标访问，它本身是为流式加载大数据集设计的，不支持随机索引，需要随机访问的话可以提前把所有数据存成numpy数组再索引，但CIFAR10数据集很小可以这么做，大数据集不推荐。
• 图模式下所有张量定义都只是计算图的节点，所有取值操作都必须放到Session的run方法里执行。
• f-string嵌套引号时务必保证内外层引号类型不同，避免语法解析错误。

内容的提问来源于stack exchange，提问作者Alessandro