Python中绘制分类模型决策边界报错问题求助

问题背景 #

我搭建了一个简单的分类网络，输入是长度为2的向量，输出是整数预测结果。为理解模型的决策边界，使用了含10000个样本的数据集（训练集9000个），特征是[-1,1]^2范围内的二维向量，标签为长度3的独热编码。参考相关逻辑编写了决策边界绘制函数，但调用时报错。

绘制决策边界的代码 #

def plot_decision_boundary(X, y, model, steps=1000, cmap='Paired'):
    """
    绘制模型决策边界和数据点的函数，数据点按真实标签上色
    """
    cmap = plt.get_cmap(cmap)
    
    # 根据数据范围定义感兴趣区域
    xmin, xmax = X[:,0].min() - 1, X[:,0].max() + 1
    ymin, ymax = X[:,1].min() - 1, X[:,1].max() + 1
    steps = 1000
    x_span = np.linspace(xmin, xmax, steps)
    y_span = np.linspace(ymin, ymax, steps)
    xx, yy = np.meshgrid(x_span, y_span)

    # 对感兴趣区域内的所有点做预测
    labels = model.predict(np.c_[xx.ravel(), yy.ravel()])

    # 绘制决策边界
    z = labels.reshape(xx.shape)
    
    fig, ax = plt.subplots()
    ax.contourf(xx, yy, z, cmap=cmap, alpha=0.5)

    # 对训练数据做预测并绘制
    train_labels = model.predict(X)
    ax.scatter(X[:,0], X[:,1], c=y, cmap=cmap, lw=0)
    return(fig,ax)

调用方式与报错信息 #

调用代码：

plot_decision_boundary(train_features, train_labels_onehot, model, cmap='RdBu')

报错内容：

ValueError: 无法将大小为3000000的数组重塑为形状(1000,1000)

问题原因与解决方案 #

原因分析 #

模型是三分类任务，输出为3维独热编码（每个样本对应3个概率值），model.predict()返回结果形状为(1000000, 3)，总元素数3000000，直接尝试重塑为(1000,1000)的二维数组必然维度不匹配。同时，绘制散点图时传入的y是独热编码（形状(n_samples,3)），但plt.scatter的c参数需要一维类别值（如0、1、2），而非多列独热编码。

修复后的代码 #

def plot_decision_boundary(X, y, model, steps=1000, cmap='Paired'):
    """
    绘制模型决策边界和数据点的函数，数据点按真实标签上色
    """
    cmap = plt.get_cmap(cmap)
    
    # 根据数据范围定义感兴趣区域
    xmin, xmax = X[:,0].min() - 1, X[:,0].max() + 1
    ymin, ymax = X[:,1].min() - 1, X[:,1].max() + 1
    x_span = np.linspace(xmin, xmax, steps)
    y_span = np.linspace(ymin, ymax, steps)
    xx, yy = np.meshgrid(x_span, y_span)

    # 对感兴趣区域内的所有点做预测，并取概率最大的类别索引
    predictions = model.predict(np.c_[xx.ravel(), yy.ravel()])
    labels = np.argmax(predictions, axis=1)

    # 重塑为网格形状
    z = labels.reshape(xx.shape)
    
    fig, ax = plt.subplots()
    ax.contourf(xx, yy, z, cmap=cmap, alpha=0.5)

    # 将独热编码标签转为一维类别值
    y_labels = np.argmax(y, axis=1)
    # 绘制训练数据点
    ax.scatter(X[:,0], X[:,1], c=y_labels, cmap=cmap, lw=0)
    return(fig,ax)

关键修改点 #

• 对预测结果用np.argmax(predictions, axis=1)提取每个样本的预测类别，把3维输出转为1维标签，解决维度不匹配问题。
• 把传入的独热编码标签y用np.argmax(y, axis=1)转为一维类别值，适配plt.scatter的c参数要求。
• 移除重复定义的steps = 1000（函数参数已传入，无需重复赋值）。

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