TensorFlow中如何将BatchDataset拆分为输入与标签并解决下标报错

错误原因 #

'BatchDataset' object is not subscriptable报错的直接原因是你直接对BatchDataset类型的数据集对象使用了下标切片语法，tf.data.Dataset系列对象本身是数据集迭代器容器，不是张量，不能直接用[:, x, :]这类语法操作，需要用map方法将拆分逻辑应用到数据集内的每一个批次张量上。
另外你的split_window函数的切片逻辑也和当前数据维度不匹配：你生成的每个批次元素是一维张量（shape为(5,)），不需要写三维索引，直接按一维切片取前3位、后2位即可。

修正后的完整可运行代码 #

import tensorflow as tf

input_slice = 3

def split_window(features):
    inputs = features[:input_slice]  # 取前3个元素作为输入
    labels = features[input_slice:]  # 取剩余2个元素作为标签
    return inputs, labels

# 创建批次数据集
dataset = tf.data.Dataset.range(1, 25 + 1).batch(5)
# 对每个批次应用拆分逻辑
dataset = dataset.map(split_window)

# 验证输出
for inputs, labels in dataset:
    print("输入：", inputs.numpy())
    print("标签：", labels.numpy())

运行输出 #

输入： [1 2 3]
标签： [4 5]
输入： [6 7 8]
标签： [ 9 10]
输入： [11 12 13]
标签： [14 15]
输入： [16 17 18]
标签： [19 20]
输入： [21 22 23]
标签： [24 25]

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