如何在测试集评估阶段正确计算F1分数？

你的代码存在两个核心问题，导致F1分数计算失败：

• 嵌套列表结构错误：每次将一个batch的标签/预测结果以列表形式append到总列表，最终得到[[batch1数据], [batch2数据], ...]的二维嵌套结构，而sklearn.metrics.f1_score要求输入是一维扁平序列。
• CUDA张量转numpy报错：若模型运行在GPU上，labels和predicted是CUDA张量，直接调用.numpy()会抛出设备不匹配错误，必须先转移到CPU再转换格式。

修正后的Scikit-Learn版本代码 #

# Function to test the model 
from sklearn.metrics import f1_score
import time
import torch

since = time.time()
total=0
correct=0
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
y_pred=[]
y_true=[]

# Iterate over data.
with torch.no_grad():
    for inputs, labels in dataloadersTest_dict['Test']:
        inputs = inputs.to(device)
        labels = labels.to(device)
        predicted_outputs = model(inputs)
        _, predicted = torch.max(predicted_outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # 修正1：先转CPU再转numpy，避免CUDA张量报错
        # 修正2：用extend替代append，直接将batch元素扁平添加到总列表
        temp_true = labels.cpu().numpy().tolist()
        temp_pred = predicted.cpu().numpy().tolist()
        y_true.extend(temp_true)
        y_pred.extend(temp_pred)
        
time_elapsed = time.time() - since
test_acc=100 * correct / total

print('Evaluation completed in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Accuracy: %d %%' % (test_acc))
print('F1 Score:')
f1 = f1_score(y_true, y_pred, average='macro')
print(f1)

关键修正说明 #

1. 用.cpu()将CUDA张量转移到CPU后，再调用.numpy()转换格式
2. 用extend()替代append()，把batch中的每个元素直接添加到总列表，生成一维扁平结构，符合f1_score的输入要求

TorchMetrics版本实现（更简洁） #

如果你想用TorchMetrics计算F1分数，可以直接用其内置的指标累加机制，无需手动维护列表：

from torchmetrics import F1Score
import time
import torch

since = time.time()
total=0
correct=0
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 初始化F1指标：根据任务类型指定task，填写你的实际类别数
f1_metric = F1Score(task="multiclass", num_classes=你的类别数, average='macro').to(device)

with torch.no_grad():
    for inputs, labels in dataloadersTest_dict['Test']:
        inputs = inputs.to(device)
        labels = labels.to(device)
        predicted_outputs = model(inputs)
        _, predicted = torch.max(predicted_outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # 累加当前batch的结果到指标
        f1_metric.update(predicted, labels)
        
time_elapsed = time.time() - since
test_acc=100 * correct / total
# 计算最终F1分数
f1 = f1_metric.compute()

print('Evaluation completed in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Accuracy: %d %%' % (test_acc))
print('F1 Score:', f1.item())

TorchMetrics使用注意 #

• 必须根据任务类型设置task参数：二分类用binary，多分类用multiclass，多标签用multilabel
• num_classes需要填写你的数据集实际类别数量
• 指标对象要和模型、数据放在同一设备（GPU/CPU）上

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