如何基于二维矩阵构建保留空间信息的ML训练数据集？

我有以下两个二维矩阵：

import numpy as np

matrix_1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
matrix_2 = np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18]])

对应的标签为：

labels = np.array([0, 1])

我要基于这些数据构建用于ML模型训练的数据集：

• 试过给每个矩阵单独存成CSV，但没法用多份CSV训练模型；
• 试过扁平化矩阵，但会丢失空间信息。

希望得到的训练数据格式如下（保留矩阵的二维结构）：

x_train = np.array([[[1, 2, 3],[4, 5, 6],[7, 8, 9]],[[10, 11, 12],[13, 14, 15],[16, 17, 18]]])  
y_train = np.array(["cat", "dog"])

求除了扁平化之外的其他可行方法。

可行解决方案 #

1. 直接构建三维数组（最直接） #

直接将多个二维矩阵堆叠成三维数组作为特征集，标签对应匹配即可：

import numpy as np

matrix_1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
matrix_2 = np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18]])

# 构建x_train：三维数组，shape为(样本数, 行数, 列数)
x_train = np.array([matrix_1, matrix_2])
# 构建y_train：可以直接映射标签，或者用原数字标签
y_train = np.array(["cat", "dog"])
# 如果用原数字标签：y_train = labels

这种方式完全保留矩阵的空间结构，适合卷积神经网络（CNN）等需要空间信息的模型。

2. 用Numpy的stack函数批量处理 #

如果矩阵数量较多，用np.stack更高效：

# 假设你有一个存储所有矩阵的列表
matrices = [matrix_1, matrix_2]
x_train = np.stack(matrices, axis=0)  # axis=0表示在样本维度堆叠
y_train = np.array(["cat", "dog"])

3. 保存为Numpy二进制文件（方便后续复用） #

如果需要把数据集保存下来，避免每次重新构建，可存成.npy或.npz文件：

# 保存单个数组
np.save("x_train.npy", x_train)
np.save("y_train.npy", y_train)

# 或者打包成一个.npz文件
np.savez("dataset.npz", x_train=x_train, y_train=y_train)

# 加载时
data = np.load("dataset.npz")
x_train = data["x_train"]
y_train = data["y_train"]

这种格式完全保留数组的维度信息，加载速度远快于CSV，适合存储多维数据。

4. 使用Pandas存储（兼容表格工具） #

如果需要同时兼容表格类工具查看，可把每个矩阵的行作为嵌套列表存入Pandas的列，标签单独一列：

import pandas as pd

df = pd.DataFrame({
    "features": [matrix_1.tolist(), matrix_2.tolist()],
    "label": ["cat", "dog"]
})

# 保存为CSV（加载时需要把嵌套列表转回numpy数组）
df.to_csv("dataset.csv", index=False)

# 加载时恢复
df_loaded = pd.read_csv("dataset.csv")
x_train = np.array(df_loaded["features"].apply(eval).tolist())
y_train = df_loaded["label"].values

注意：这种方式适合小数据集，大数据集用Numpy二进制更高效。

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