如何将MNIST的(行,28,28)形状ndarray转为DataFrame适配TensorFlow预测？

方案1：训练时封装自定义PyFunc模型（推荐） #

在训练阶段就把输入格式转换逻辑封装到自定义PyFunc类中，加载模型后可直接用二维DataFrame输入，无需额外处理。

import mlflow
import mlflow.pyfunc
import tensorflow as tf
import pandas as pd
from tensorflow.keras.datasets import mnist

class MNISTModelWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 加载预训练Keras模型
        self.model = tf.keras.models.load_model(context.artifacts["model_path"])
    
    def predict(self, context, model_input):
        # 将二维DataFrame转换为模型要求的(None, 28, 28, 1)形状
        input_array = model_input.values.reshape(-1, 28, 28, 1)
        return self.model.predict(input_array)

# ---------------------- 训练并保存模型示例 ----------------------
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1) / 255.0
x_test = x_test.reshape(-1, 28, 28, 1) / 255.0

# 构建你的CNN模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1)

# 保存原始Keras模型
keras_model_path = "./mnist_keras_model"
model.save(keras_model_path)

# 封装并保存自定义PyFunc模型
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        "mnist_pyfunc_model",
        python_model=MNISTModelWrapper(),
        artifacts={"model_path": keras_model_path}
    )

# ---------------------- 加载模型并预测 ----------------------
# 替换为你的实际run_id
logged_model = "runs:/<你的run_id>/mnist_pyfunc_model"
loaded_model = mlflow.pyfunc.load_model(logged_model)

# 准备二维DataFrame输入
xtest_df = pd.DataFrame(x_test.reshape(-1, 784))
predictions = loaded_model.predict(xtest_df)

方案2：直接加载原始Keras模型 #

如果不需要用DataFrame输入，可跳过PyFunc封装，直接加载MLflow保存的原始Keras模型，传入符合要求的numpy数组即可：

import mlflow.keras
import tensorflow as tf
from tensorflow.keras.datasets import mnist

# 加载原始Keras模型（需训练时用mlflow.keras.log_model保存）
logged_model = "runs:/<你的run_id>/model"
loaded_keras_model = mlflow.keras.load_model(logged_model)

(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 调整输入形状并归一化
x_test_input = x_test.reshape(-1, 28, 28, 1) / 255.0
# 直接执行预测
predictions = loaded_keras_model.predict(x_test_input)

错误原因说明 #

• 直接将x_test转为DataFrame报错：因为x_test是三维数组(10000, 28, 28)，而DataFrame仅支持二维输入。
• 转成784列DataFrame后输入不匹配：模型输入层要求(None, 28, 28, 1)的三维结构，而784列的DataFrame对应二维数组(10000, 784)，形状不兼容。

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