如何用mu和sigma标准化PyTorch训练集？函数参数传递疑问

参数传递逻辑 #

你的核心困惑在于mu和sigma的传递方式，正确处理逻辑分两种场景：

• 训练集处理：无需提前计算，让函数内部自动基于训练集统计量计算即可（避免手动计算出错）
• 测试集处理：必须传入训练集计算得到的mu和sigma，绝对不能用测试集自身的统计量，否则会导致数据泄露，破坏模型泛化能力
• 不要传入mu=0、sigma=1，这不属于基于当前数据集的Z-Score标准化，而是强制映射到标准正态分布，不符合需求

修正后的函数实现（适配PyTorch张量特性） #

将mu和sigma设为可选参数，让函数同时支持训练集自动计算、测试集传入参数两种模式，且利用PyTorch的广播机制替代手动遍历，效率更高：

import torch

def norm_set(x: torch.Tensor, mu: torch.Tensor = None, sigma: torch.Tensor = None) -> torch.Tensor:
    ### your code here
    # 自动计算训练集的均值和标准差（按特征维度，dim=0表示对每列计算）
    if mu is None or sigma is None:
        # keepdim保持维度一致，确保后续广播运算正常
        mu = x.mean(dim=0, keepdim=True)
        # 加1e-8避免标准差为0时的除零错误
        sigma = x.std(dim=0, keepdim=True) + 1e-8

    # will add loop to calculate z-score here
    # PyTorch广播机制自动处理所有元素，无需手动遍历
    normalized_x = (x - mu) / sigma

    return normalized_x

调用示例 #

标准化训练集 #

# 函数自动计算训练集的mu和sigma，返回标准化后的张量
x_train_normalized = norm_set(x_train)

# 若需要保存训练集的统计量用于测试集，可单独提取
train_mu = x_train.mean(dim=0, keepdim=True)
train_sigma = x_train.std(dim=0, keepdim=True) + 1e-8

标准化测试集 #

# 传入训练集的mu和sigma，确保测试集与训练集用同一套统计量标准化
x_test_normalized = norm_set(x_test, mu=train_mu, sigma=train_sigma)

关于手动遍历的说明（不推荐） #

如果因特殊需求必须手动遍历张量元素，可参考以下代码（但效率远低于向量化操作）：

def norm_set(x: torch.Tensor, mu: torch.Tensor = None, sigma: torch.Tensor = None) -> torch.Tensor:
    ### your code here
    if mu is None or sigma is None:
        mu = x.mean(dim=0, keepdim=True)
        sigma = x.std(dim=0, keepdim=True) + 1e-8

    # will add loop to calculate z-score here
    normalized_x = torch.zeros_like(x)
    # 遍历每个样本
    for i in range(x.shape[0]):
        # 遍历每个特征
        for j in range(x.shape[1]):
            normalized_x[i][j] = (x[i][j] - mu[0][j]) / sigma[0][j]

    return normalized_x

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