PyTorch测试集MSE、MAE、R2指标计算正确性及内置工具咨询

我正在使用人工数据集构建基础神经网络以学习PyTorch，目前需要获取测试集的预测结果并计算包括MSE、MAE、R2在内的统计指标。现咨询：我编写的测试指标计算代码是否正确？是否存在错误？另外，PyTorch中是否有类似scikit-learn的内置函数可一次性输出这些统计结果？

完整代码 #

数据准备 #

import random
import torch
import pandas as pd
import numpy as np
from torch import nn
from torch.utils.data import Dataset,DataLoader,TensorDataset
from torchvision import datasets, transforms
import math


n_input, n_hidden, n_out= 5, 64, 1

# 创建训练和测试数据集
X_train =  pd.DataFrame([[random.random() for i in range(n_input)] for j in range(1000)])
y_train =  pd.DataFrame([[random.random() for i in range(n_out)] for j in range(1000)])

X_test =  pd.DataFrame([[random.random() for i in range(n_input)] for j in range(50)])
y_test  =  pd.DataFrame([[random.random() for i in range(n_out)] for j in range(50)])


test_dataset = TensorDataset(torch.Tensor(X_test.to_numpy().astype(np.float32)), torch.Tensor((y_test).to_numpy().astype(np.float32)))
testloader = DataLoader(test_dataset, batch_size= 32)

# 训练集使用32作为batch size
training_dataset = TensorDataset(torch.Tensor(X_train.to_numpy().astype(np.float32)), torch.Tensor((y_train).to_numpy().astype(np.float32)))
dataloader = DataLoader(training_dataset, batch_size=32, shuffle=True)

模型定义 #

model = nn.Sequential(nn.Linear(n_input, n_hidden),
                      nn.ReLU(),
                      nn.Linear(n_hidden, n_out),
                      nn.ReLU())
loss_function = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

训练循环 #

losses = []
epochs = 1000

for epoch in range(epochs+1): 
    for times,(x_train,y_train) in enumerate(dataloader):
        y_pred  = model(x_train) 
        loss = loss_function(y_pred, y_train) 
   
        model.zero_grad() 
        loss.backward() 

        optimizer.step() 

指标计算（待验证代码） #

from torchmetrics import R2Score

r2score = R2Score()
running_mae = running_mse = running_r2 = 0

with torch.no_grad():
    model.eval()
    for times,(x_test,y_test) in enumerate(testloader):
        y_pred = model(x_test)
        error = torch.abs(y_pred - y_test).sum().data
        squared_error=((y_pred - y_test)*(y_pred - y_test)).sum().data
        running_mae+=error
        running_mse+=squared_error
        running_r2+=r2score(y_pred, y_test)
mse =  math.sqrt(squared_error/ len(testloader))
mae = error / len(testloader)
r2 = running_r2 / len(testloader)
print("MSE:",mse, "MAE:", mae, "R2:", r2) 

代码错误分析 #

你的指标计算代码存在多处逻辑错误，具体如下：

1. MSE/MAE计算错误

   • 最终计算时误用了最后一个batch的squared_error和error，而非累加得到的running_mse和running_mae，导致结果仅反映最后一批数据的情况。
   • 分母选择错误：应该用测试集总样本数（如len(X_test)），而非len(testloader)（loader长度是批次数，和样本数无关，且最后一批可能不满batch_size）。
   • MSE定义错误：你当前对squared_error开根号，得到的是RMSE（均方根误差），而非标准的MSE（均方误差）。

2. R2指标计算错误

   • R2是全局统计指标，不能通过对各batch的R2取平均得到正确结果。R2Score需要通过update方法累积全量数据，最后用compute生成全局结果，而非直接累加每个batch的R2值。

修正后的指标计算代码 #

使用你已导入的TorchMetrics库可以规范计算这些指标，它支持分批累积数据，结果更准确：

from torchmetrics import R2Score, MeanSquaredError, MeanAbsoluteError

# 初始化指标对象
mse_metric = MeanSquaredError()
mae_metric = MeanAbsoluteError()
r2_metric = R2Score()

with torch.no_grad():
    model.eval()
    for x_test, y_test in testloader:
        y_pred = model(x_test)
        # 分批更新指标数据
        mse_metric.update(y_pred, y_test)
        mae_metric.update(y_pred, y_test)
        r2_metric.update(y_pred, y_test)

# 计算最终全局指标
mse = mse_metric.compute()
mae = mae_metric.compute()
r2 = r2_metric.compute()

print(f"MSE: {mse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}")

# 重置指标（若后续需要重复计算）
mse_metric.reset()
mae_metric.reset()
r2_metric.reset()

PyTorch中的替代工具 #

PyTorch生态中的TorchMetrics库就是对标scikit-learn的指标工具，它提供了绝大多数常见统计指标的实现，支持自动累积、分布式训练兼容、指标重置等功能，无需手动编写循环累加逻辑。

如果不想额外依赖库，也可以用torch.nn.functional中的基础函数（如mse_loss、l1_loss）手动计算，但需要自行拼接全量测试集的预测和真实值，不如TorchMetrics便捷。

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