TensorFlow堆叠多传感器张量后执行train_test_split出现索引错误如何解决

报错原因 #

1. scikit-learn的train_test_split函数原生适配Numpy数组，执行拆分逻辑时会生成Numpy格式的索引数组，用该数组对输入数据做切片。
2. 你堆叠三个传感器数据时使用tf.stack得到的输出是TensorFlow张量类型，TensorFlow张量不支持直接用普通Numpy数组作为索引进行切片，这是触发报错的直接原因。
3. 单传感器场景可正常运行的原因是你用np.moveaxis处理后得到的是Numpy数组，天然支持Numpy索引切片逻辑。

解决方案 #

你可以根据自身使用场景选择以下任意一种方案解决问题：

方案1：张量转Numpy数组后拆分 #

将堆叠得到的TensorFlow张量转为Numpy数组后再传入拆分函数，后续需要张量格式可再转换：

final = tf.stack([sensor1, sensor2, sensor3], axis=-1).numpy()
X_train, X_test, y_train, y_test = train_test_split(final, y, test_size=0.25, random_state=42, stratify=y)
# 如需转回TensorFlow张量，执行以下操作即可
X_train = tf.convert_to_tensor(X_train)
X_test = tf.convert_to_tensor(X_test)

方案2：先拆分再堆叠 #

先对三个传感器的原始数据执行拆分，再分别对训练集、测试集做张量堆叠：

# 同时拆分三个传感器的原始数据和标签
s1_train, s1_test, s2_train, s2_test, s3_train, s3_test, y_train, y_test = train_test_split(sensor1, sensor2, sensor3, y, test_size=0.25, random_state=42, stratify=y)
# 分别堆叠训练集和测试集
X_train = tf.stack([s1_train, s2_train, s3_train], axis=-1)
X_test = tf.stack([s1_test, s2_test, s3_test], axis=-1)

方案3：索引拆分适配张量 #

先单独拆分样本索引，再用TensorFlow原生的gather方法对张量做切分，无需转换数据格式：

# 仅拆分样本索引
train_idx, test_idx = train_test_split(range(len(final)), test_size=0.25, random_state=42, stratify=y)
# 用tf.gather完成张量切分
X_train = tf.gather(final, train_idx)
X_test = tf.gather(final, test_idx)

内容的提问来源于stack exchange，提问作者Adler Müller