超实用！用tinygrad绘制混淆矩阵，3步搞定模型分类性能分析

2026-02-04 04:09:20作者：郜逊炳

在机器学习模型评估中，混淆矩阵是直观展示分类效果的核心工具。tinygrad作为轻量级深度学习框架，不仅支持高效模型训练，还能通过简洁代码实现专业的混淆矩阵可视化。本文将带你通过三个简单步骤，用tinygrad快速生成混淆矩阵，轻松掌握模型分类性能分析技巧。

准备工作：安装与环境配置

首先确保已安装tinygrad框架。如果尚未安装，可通过以下命令克隆仓库并完成基础配置：

git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad
cd tinygrad
pip install -e .

tinygrad的模块化设计让扩展功能变得简单，我们将使用其内置的张量操作和可视化工具来构建混淆矩阵。核心功能实现可参考框架的examples/目录下的分类任务示例，其中包含了模型训练与评估的完整流程。

步骤1：生成模型预测结果

在绘制混淆矩阵前，需要获取模型对测试集的预测结果。以下是使用tinygrad实现图像分类预测的基础代码框架：

from tinygrad.tensor import Tensor
from tinygrad.nn import Linear, ReLU
import numpy as np

# 加载训练好的模型（示例代码）
model = Linear(28*28, 10)  # 以MNIST手写数字分类为例

# 准备测试数据
test_images = Tensor(np.load("test_images.npy"))
test_labels = np.load("test_labels.npy")

# 获取预测结果
with Tensor.no_grad():
    outputs = model(test_images.reshape(-1, 28*28))
    predictions = outputs.argmax(axis=1).numpy()

这段代码展示了如何使用tinygrad的Tensor类进行推理计算。实际应用中，你可以替换为自己训练的模型，例如examples/beautiful_mnist.py中实现的MNIST分类模型。

步骤2：计算混淆矩阵核心数据

tinygrad虽然没有直接提供混淆矩阵API，但可以通过简单的张量操作实现。以下是计算混淆矩阵的核心函数：

def compute_confusion_matrix(y_true, y_pred, num_classes):
    # 初始化混淆矩阵
    cm = np.zeros((num_classes, num_classes), dtype=int)
    # 填充混淆矩阵
    for true, pred in zip(y_true, y_pred):
        cm[true][pred] += 1
    return cm

# 计算10类分类任务的混淆矩阵
confusion_matrix = compute_confusion_matrix(test_labels, predictions, 10)

该函数通过统计真实标签与预测标签的对应关系，生成混淆矩阵的基础数据。对于更复杂的场景，可以参考tinygrad/nn/目录下的评估工具实现。

步骤3：可视化混淆矩阵

结合matplotlib库，我们可以将混淆矩阵数据转换为直观的热力图。以下是完整的可视化代码：

import matplotlib.pyplot as plt
import seaborn as sns

def plot_confusion_matrix(cm, class_names):
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('预测类别')
    plt.ylabel('真实类别')
    plt.title('模型分类混淆矩阵')
    plt.savefig('confusion_matrix.png')
    plt.close()

# 绘制MNIST数据集的混淆矩阵
class_names = [str(i) for i in range(10)]
plot_confusion_matrix(confusion_matrix, class_names)