如何将含振动信号数据的Python字典转换为TensorFlow数据集

核心思路 #

你的方向完全正确——给每个样本配对对应类别标签，再转成TF Dataset就行。问题大概率出在数组拼接或标签映射的细节上，下面是可直接运行的实现步骤：

步骤1：整理样本与标签 #

先遍历字典，把每个类别的样本和对应的整数标签逐一配对，保证标签和样本严格对应：

import tensorflow as tf
import numpy as np

# 假设你的振动数据字典是vibration_data，键为类别名称，值是shape=(样本数, n)的二维数组
vibration_data = {
    "class_0": np.random.rand(50, 100),  # 示例：50个样本，每个含100个采样点
    "class_1": np.random.rand(50, 100),
    # ... 剩下8个类别
}

# 把类别名称映射为整数标签
class_names = list(vibration_data.keys())
label_map = {name: idx for idx, name in enumerate(class_names)}

# 收集所有样本和对应标签
all_samples = []
all_labels = []
for class_name, samples in vibration_data.items():
    # 给当前类别的所有样本打上相同标签
    labels = [label_map[class_name]] * len(samples)
    all_samples.extend(samples)
    all_labels.extend(labels)

# 转成numpy数组方便后续处理
all_samples_np = np.array(all_samples)
all_labels_np = np.array(all_labels)

步骤2：转换为TensorFlow Dataset #

1D CNN要求输入格式为(样本数, 采样点数量, 通道数)，如果你的振动信号是单通道，需要给样本增加一个通道维度：

# 新增通道维度（单通道信号设为1）
all_samples_np = np.expand_dims(all_samples_np, axis=-1)

# 转换成TF Dataset
dataset = tf.data.Dataset.from_tensor_slices((all_samples_np, all_labels_np))

步骤3：数据集预处理（推荐） #

根据训练需求添加打乱、批量、预取等操作，提升训练效率：

# 打乱数据集（buffer_size设为总样本数）
dataset = dataset.shuffle(buffer_size=len(all_samples_np))

# 划分训练/验证集（示例按8:2拆分）
train_size = int(0.8 * len(all_samples_np))
train_dataset = dataset.take(train_size)
val_dataset = dataset.skip(train_size)

# 设置批量大小并开启预取
batch_size = 32
train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
val_dataset = val_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

常见问题排查 #

• 如果之前拼接失败，先检查每个类别对应的二维数组形状是否统一，确保所有样本的采样点数量n完全一致；
• 标签必须是整数类型，TF分类任务默认接受整数标签，若需要独热编码可后续用tf.one_hot转换。

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