我需要在一个数据集上训练一个卷积神经网络。神经网络本身可以正常工作，按预期运行，但现在我想生成一个混淆矩阵，并将其导出为一个json文件以便进一步查看。

我尝试使用来自sklearn和 torchmetrics的 ConfusionMatrix 函数，但我都无法让它们工作，因为它们给出的是不同的错误。

Sklearn错误：

ValueError: Classification metrics can't handle a mix of binary and continuous targets

Torchmetrics错误：

AttributeError: 'list' object has no attribute 'shape'

不幸的是，我还没能找到这两者的可行解决方案。

import logging
import os
import pickle
from pathlib import Path
from typing import Any, Callable

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix

#from IPython.sphinxext.ipython_directive import test
from torch import Tensor, nn
from torch.utils.data import DataLoader, random_split
from torch.utils.data import Dataset as _Dataset

#from torchmetrics.classification import ConfusionMatrix
from .ecg import build_model_prototype

"""
uv run python -m bach_ben.fibrillation
"""

def get_dataset_dir() -> str:
    p = os.getenv("DATASET_DIR")
    if p is None:
        raise ValueError("DATASET_DIR environment variable not set")
    logger = logging.getLogger(__name__)
    logger.log(level=logging.INFO, msg=f"using DATASET_DIR {p}")
    return p

def get_model_dir() -> str:
    p = os.getenv("MODEL_DIR")
    if p is None:
        raise ValueError("MODEL_DIR environment variable not set")
    logger = logging.getLogger(__name__)
    logger.log(level=logging.INFO, msg=f"using MODEL_DIR {p}")
    return p

class MitBihAtrialFibrillationDataset(_Dataset[tuple[torch.Tensor, torch.Tensor]]):
    def __init__(
        self,
        dataset_root: Path,
        sample_preprocessing: Callable[[torch.Tensor], torch.Tensor] | None = None,
        dtype_samples: torch.dtype = torch.float32,
        dtype_labels: torch.dtype = torch.float32,
    ) -> None:
        super().__init__()

        self._samples, self._labels = self._load_samples_labels(dataset_root)

        self._samples = self._samples.to(dtype_samples)
        self._labels = self._labels.to(dtype_labels)

        if sample_preprocessing is not None:
            self._samples = sample_preprocessing(self._samples)

    def __len__(self) -> int:
        return len(self._labels)

    def __getitem__(self, index: Any) -> tuple[torch.Tensor, torch.Tensor]:
        return self._samples[index], self._labels[index]


    @staticmethod
    def _load_samples_labels(dataset_root: Path) -> tuple[torch.Tensor, torch.Tensor]:
        af_dir = dataset_root / "atrial_fibrillation"
        sinus_dir = dataset_root / "sinus_rhythm"
        for d in (af_dir, sinus_dir):
            if not d.exists():
                raise ValueError(f"{d} does not exist!")
        def load_samples_labels_for_class(
            class_name: str,
        ) -> tuple[torch.Tensor, torch.Tensor]:
            sample_buffer = []
            label = ""

            for pickle_file in (dataset_root / class_name).glob("*.pickle"):
                with pickle_file.open("rb") as in_file:
                    sample, label = pickle.load(in_file)
                sample_buffer.append(torch.tensor(sample.T).unsqueeze(dim=0))

            samples = torch.cat(sample_buffer)
            n = len(samples)
            labels = torch.zeros(n) if label == "(N" else torch.ones(n)

            return samples, labels

        sinus_samples, sinus_labels = load_samples_labels_for_class("sinus_rhythm")
        af_samples, af_labels = load_samples_labels_for_class("atrial_fibrillation")

        samples = torch.cat([sinus_samples, af_samples]).to(torch.float32)
        labels = torch.cat([sinus_labels, af_labels]).to(torch.int64)

        return samples, labels



def trainCNN(model: torch.nn.Module, epochs, batch_size):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(0)
    np.random.seed(0)

    dataset_base_dir = Path(get_dataset_dir()) / "mit-bih-atrial-fibrillation"
    model_save_dir = Path(get_model_dir())
    def select_first_channel(x: torch.Tensor) -> torch.Tensor:
        return x[:, 0:1, :]

    ds = MitBihAtrialFibrillationDataset(dataset_base_dir, sample_preprocessing=select_first_channel)
    train_ds, val_ds, test_ds = random_split(ds, [0.6, 0.2, 0.2])
    mean = 0
    std = 0
    for x, _ in train_ds.dataset:
        mean += x[0].mean()
        std += x[0].std()
    mean /= len(train_ds)
    std /= len(train_ds)

    def normalize(x: torch.Tensor) -> torch.Tensor:
        return (x - mean) / std

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=True)
    optim = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.BCEWithLogitsLoss()

    train_hist, val_hist = [], []
    prev_va_loss = 0.0
    for epoch in range(1, epochs + 1):
        model.train()
        tr_loss = 0.0
        for x, y in train_loader:
            x, y = normalize(x.to(device)), y.to(device)
            optim.zero_grad()
            pred = model(x)
            loss = loss_fn(pred, y)
            loss.backward()
            optim.step()
            tr_loss += loss.item() * x.size(0)

        tr_loss /= len(train_loader)
        if tr_loss < 0.036:
            break

        model.eval()
        va_loss = 0.0
        with torch.no_grad():
            for x, y in val_loader:
                x, y = normalize(x.to(device)), y.to(device)
                pred = model(x)
                va_loss += loss_fn(pred, y).item() * x.size(0)
        va_loss /= len(val_loader)
        if va_loss < prev_va_loss:
            torch.save(model.state_dict(), model_save_dir / "test.pt")
        prev_va_loss = va_loss


        train_hist.append(tr_loss)
        val_hist.append(va_loss)
        print(f"Epoch {epoch:02d} | train MSE={tr_loss:.6f} | val MSE={va_loss:.6f}")

    #model.load_state_dict(torch.load(model_save_dir / "test.pt"))
    test_loss = 0.0
    y_true, y_pred = [], []
    model.eval()
    with torch.no_grad():
        for x, y in test_loader:
            x, y = normalize(x.to(device)), y.to(device)
            output = model(x)
            test_loss += loss_fn(output, y).item() * x.size(0)
            y_true.append(y)
            y_pred.append(output)
            print(y)
            print("test")
            print(output)
    #cm = ConfusionMatrix(task="binary", num_classes=2)
    #cm(y_pred, y_true)
    cm = confusion_matrix(y_true, y_pred)
    print(cm)



    fig, axes = plt.subplots(1, 1, figsize=(10, 8), constrained_layout=True)

    #axes.plot(val_hist, label="validation")
    #axes.set_title("Validation MSE")
    #axes.set_xlabel("Epoch")
    #axes.set_ylabel("MSE")
    #axes.legend()
    #axes.grid(True, alpha=0.3)

    #plt.show()




if __name__ == "__main__":
    model = build_model_prototype()
    trainCNN(model, 5, 64)

.ecg只是把我的模型搭建起来，并没有做出什么特别的事情。若有需要，我也可以把那部分的代码给出。

我对PyTorch还挺新手的，所以不知道接下来该怎么做。

解决方案

主要问题在于你的混淆矩阵代码接收到的是：

1. y_true作为一个批次张量的列表
2. y_pred作为原始的logits
3. 不是像0 和1 那样扁平化的二进制标签

对于二分类的混淆矩阵，你需要：

1. 收集跨所有批次的目标
2. 对模型输出应用sigmoid()
3. 将它们阈值化为类别标签
4. 将所有数据扁平化为一个列表 / 数组
5. 将这些离散标签传入confusion_matrix()
6. 在写JSON之前把矩阵转换为.tolist()

下面是你函数的一个清理版本。

import json
import logging
import os
import pickle
from pathlib import Path
from typing import Any, Callable

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
from torch import nn
from torch.utils.data import DataLoader, random_split
from torch.utils.data import Dataset as _Dataset

from .ecg import build_model_prototype


def get_dataset_dir() -> str:
    p = os.getenv("DATASET_DIR")
    if p is None:
        raise ValueError("DATASET_DIR environment variable not set")
    logger = logging.getLogger(__name__)
    logger.info(f"using DATASET_DIR {p}")
    return p


def get_model_dir() -> str:
    p = os.getenv("MODEL_DIR")
    if p is None:
        raise ValueError("MODEL_DIR environment variable not set")
    logger = logging.getLogger(__name__)
    logger.info(f"using MODEL_DIR {p}")
    return p


class MitBihAtrialFibrillationDataset(_Dataset[tuple[torch.Tensor, torch.Tensor]]):
    def __init__(
        self,
        dataset_root: Path,
        sample_preprocessing: Callable[[torch.Tensor], torch.Tensor] | None = None,
        dtype_samples: torch.dtype = torch.float32,
        dtype_labels: torch.dtype = torch.float32,
    ) -> None:
        super().__init__()

        self._samples, self._labels = self._load_samples_labels(dataset_root)

        self._samples = self._samples.to(dtype_samples)
        self._labels = self._labels.to(dtype_labels)

        if sample_preprocessing is not None:
            self._samples = sample_preprocessing(self._samples)

    def __len__(self) -> int:
        return len(self._labels)

    def __getitem__(self, index: Any) -> tuple[torch.Tensor, torch.Tensor]:
        return self._samples[index], self._labels[index]

    @staticmethod
    def _load_samples_labels(dataset_root: Path) -> tuple[torch.Tensor, torch.Tensor]:
        af_dir = dataset_root / "atrial_fibrillation"
        sinus_dir = dataset_root / "sinus_rhythm"

        for d in (af_dir, sinus_dir):
            if not d.exists():
                raise ValueError(f"{d} does not exist!")

        def load_samples_labels_for_class(class_name: str) -> tuple[torch.Tensor, torch.Tensor]:
            sample_buffer = []
            label = ""

            for pickle_file in (dataset_root / class_name).glob("*.pickle"):
                with pickle_file.open("rb") as in_file:
                    sample, label = pickle.load(in_file)
                sample_buffer.append(torch.tensor(sample.T).unsqueeze(dim=0))

            samples = torch.cat(sample_buffer)
            n = len(samples)

            # Assuming "(N" means sinus rhythm -> class 0, otherwise class 1
            labels = torch.zeros(n) if label == "(N" else torch.ones(n)
            return samples, labels

        sinus_samples, sinus_labels = load_samples_labels_for_class("sinus_rhythm")
        af_samples, af_labels = load_samples_labels_for_class("atrial_fibrillation")

        samples = torch.cat([sinus_samples, af_samples]).to(torch.float32)
        labels = torch.cat([sinus_labels, af_labels]).to(torch.float32)

        return samples, labels


def evaluate_binary_classifier(
    model: torch.nn.Module,
    data_loader: DataLoader,
    loss_fn: nn.Module,
    normalize_fn: Callable[[torch.Tensor], torch.Tensor],
    device: str,
) -> tuple[float, list[int], list[int]]:
    model.eval()
    total_loss = 0.0
    y_true: list[int] = []
    y_pred: list[int] = []

    with torch.no_grad():
        for x, y in data_loader:
            x = normalize_fn(x.to(device))
            y = y.to(device)

            logits = model(x).squeeze(-1)
            loss = loss_fn(logits, y)

            total_loss += loss.item() * x.size(0)

            probs = torch.sigmoid(logits)
            preds = (probs >= 0.5).long()

            y_true.extend(y.cpu().long().tolist())
            y_pred.extend(preds.cpu().tolist())

    avg_loss = total_loss / len(data_loader.dataset)
    return avg_loss, y_true, y_pred


def trainCNN(model: torch.nn.Module, epochs: int, batch_size: int) -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)

    torch.manual_seed(0)
    np.random.seed(0)

    dataset_base_dir = Path(get_dataset_dir()) / "mit-bih-atrial-fibrillation"
    model_save_dir = Path(get_model_dir())
    model_save_dir.mkdir(parents=True, exist_ok=True)

    def select_first_channel(x: torch.Tensor) -> torch.Tensor:
        return x[:, 0:1, :]

    ds = MitBihAtrialFibrillationDataset(
        dataset_base_dir,
        sample_preprocessing=select_first_channel,
    )

    train_ds, val_ds, test_ds = random_split(ds, [0.6, 0.2, 0.2])

    # Compute normalization stats on the training split only
    mean = 0.0
    std = 0.0
    for x, _ in train_ds:
        mean += x.mean()
        std += x.std()

    mean /= len(train_ds)
    std /= len(train_ds)

    def normalize(x: torch.Tensor) -> torch.Tensor:
        return (x - mean) / (std + 1e-8)

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.BCEWithLogitsLoss()

    train_hist = []
    val_hist = []

    best_val_loss = float("inf")
    best_model_path = model_save_dir / "best_model.pt"

    for epoch in range(1, epochs + 1):
        model.train()
        train_loss = 0.0

        for x, y in train_loader:
            x = normalize(x.to(device))
            y = y.to(device)

            optimizer.zero_grad()

            logits = model(x).squeeze(-1)
            loss = loss_fn(logits, y)

            loss.backward()
            optimizer.step()

            train_loss += loss.item() * x.size(0)

        train_loss /= len(train_loader.dataset)

        val_loss, _, _ = evaluate_binary_classifier(
            model=model,
            data_loader=val_loader,
            loss_fn=loss_fn,
            normalize_fn=normalize,
            device=device,
        )

        train_hist.append(train_loss)
        val_hist.append(val_loss)

        print(f"Epoch {epoch:02d} | train loss={train_loss:.6f} | val loss={val_loss:.6f}")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), best_model_path)

    # Load best model before testing
    if best_model_path.exists():
        model.load_state_dict(torch.load(best_model_path, map_location=device))

    test_loss, y_true, y_pred = evaluate_binary_classifier(
        model=model,
        data_loader=test_loader,
        loss_fn=loss_fn,
        normalize_fn=normalize,
        device=device,
    )

    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])

    print("\nTest loss:", test_loss)
    print("Confusion matrix:")
    print(cm)

    cm_json = {
        "labels": ["sinus_rhythm", "atrial_fibrillation"],
        "matrix": cm.tolist(),
        "test_loss": test_loss,
        "num_samples": len(y_true),
    }

    with open(model_save_dir / "confusion_matrix.json", "w", encoding="utf-8") as f:
        json.dump(cm_json, f, indent=4)

    disp = ConfusionMatrixDisplay(
        confusion_matrix=cm,
        display_labels=["sinus_rhythm", "atrial_fibrillation"],
    )
    disp.plot()
    plt.show()

    # Optional: loss curve
    plt.figure(figsize=(8, 5))
    plt.plot(train_hist, label="train")
    plt.plot(val_hist, label="validation")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training and Validation Loss")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()


if __name__ == "__main__":
    model = build_model_prototype()
    trainCNN(model, epochs=5, batch_size=64)

为什么你原来的版本失败

这部分就是问题所在：

y_true.append(y)
y_pred.append(output)
cm = confusion_matrix(y_true, y_pred)

因为在循环结束后，数据看起来更像是这样的：

y_true = [tensor([...]), tensor([...]), tensor([...])]
y_pred = [tensor([...]), tensor([...]), tensor([...])]

而y_pred仍然是连续的logits，而不是二进制的类别预测。

sklearn.metrics.confusion_matrix() 期望的是类似如下的输入：

y_true = [0, 1, 0, 1, 1]
y_pred = [0, 1, 0, 0, 1]

不是张量的批次，也不是原始模型分数。

导出的文件看起来会是这样的：

{
    "labels": [
        "sinus_rhythm",
        "atrial_fibrillation"
    ],
    "matrix": [
        [TN, FP],
        [FN, TP]
    ],
    "test_loss": 0.12345,
    "num_samples": 240
}