PyTorch统计函数详解与应用实践

1. PyTorch中的统计学函数概述

在深度学习项目中，统计分析是不可或缺的一环。PyTorch作为当前最流行的深度学习框架之一，提供了丰富的统计学函数库，这些函数在数据预处理、模型评估和结果分析中都扮演着关键角色。不同于传统的统计学软件，PyTorch的统计函数能够直接在GPU上高效执行，并且完美支持自动微分，这使得它们特别适合深度学习工作流。

PyTorch的统计函数主要分布在几个核心模块中：

• torch：包含基础统计函数如mean(), std(), var()等
• torch.distributions：提供概率分布相关的操作
• torch.special：包含特殊数学函数如伽马函数、贝塞尔函数等

这些函数不仅支持标量计算，更重要的是能够高效处理高维张量，这对于处理图像、文本等复杂数据尤为重要。例如，在处理一批图像数据时，我们可以轻松计算整个batch的均值、方差等统计量，而无需手动编写循环。

提示：PyTorch统计函数的一个关键优势是它们会自动保持梯度信息，这使得我们可以在模型训练过程中直接使用这些统计量作为损失函数的一部分。

2. 基础统计函数详解

2.1 中心趋势度量

中心趋势度量是最基础的统计函数，PyTorch提供了完整的实现：

python复制import torch

# 创建示例数据
data = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

# 均值计算
mean_all = torch.mean(data)          # 所有元素的均值
mean_dim0 = torch.mean(data, dim=0)  # 沿第0维的均值
mean_dim1 = torch.mean(data, dim=1)  # 沿第1维的均值

# 中位数计算
median_val = torch.median(data)      # 所有元素的中位数

在实际项目中，我们经常需要计算批处理数据的统计量。例如，在图像标准化处理中，通常会计算整个数据集的均值和标准差：

python复制# 假设images是一个4D张量[batch, channel, height, width]
channel_means = torch.mean(images, dim=[0, 2, 3])  # 计算每个通道的均值
channel_stds = torch.std(images, dim=[0, 2, 3])    # 计算每个通道的标准差

2.2 离散程度度量

离散程度度量对于理解数据分布至关重要：

python复制# 方差和标准差
variance = torch.var(data)          # 总体方差
std_dev = torch.std(data)           # 总体标准差
unbiased_var = torch.var(data, unbiased=True)  # 无偏方差估计

# 极差
range_val = torch.max(data) - torch.min(data)

# 百分位数
percentile_25 = torch.quantile(data, 0.25)  # 25百分位数
percentile_75 = torch.quantile(data, 0.75)  # 75百分位数

在模型训练中，我们经常使用这些统计量来监控权重分布或梯度变化。例如，跟踪权重矩阵的标准差可以帮助我们诊断梯度消失或爆炸问题：

python复制for name, param in model.named_parameters():
    if 'weight' in name:
        print(f"{name}: mean={torch.mean(param):.4f}, std={torch.std(param):.4f}")

2.3 高阶统计量

PyTorch还支持计算更复杂的高阶统计量：

python复制# 偏度和峰度
from torch.distributions import moments

skewness = moments.skewness(data)  # 偏度
kurtosis = moments.kurtosis(data)  # 峰度

# 协方差和相关系数
cov_matrix = torch.cov(data.T)     # 协方差矩阵
corr_matrix = torch.corrcoef(data.T)  # 相关系数矩阵

这些高阶统计量在数据分析阶段特别有用。例如，在特征工程中，我们可以使用相关系数矩阵来识别高度相关的特征：

python复制# 计算特征间的相关系数
feature_corr = torch.corrcoef(features.T)

# 找出高度相关的特征对
high_corr = torch.abs(feature_corr) > 0.8

注意：在计算协方差和相关系数时，PyTorch要求输入数据至少包含2个样本（行），否则会抛出错误。对于单样本数据，需要特殊处理。

3. 概率分布相关函数

3.1 常见概率分布

PyTorch的torch.distributions模块提供了丰富的概率分布实现：

python复制from torch.distributions import Normal, Bernoulli, Poisson, Exponential

# 正态分布
normal_dist = Normal(loc=0.0, scale=1.0)  # 均值0，标准差1
samples = normal_dist.sample((1000,))     # 采样1000个点

# 伯努利分布
bernoulli_dist = Bernoulli(probs=0.3)
binary_samples = bernoulli_dist.sample((10,))

# 泊松分布
poisson_dist = Poisson(rate=2.0)
count_samples = poisson_dist.sample((20,))

这些分布在模型构建中非常有用。例如，在变分自编码器(VAE)中，我们通常使用正态分布作为潜在变量的先验分布：

python复制# VAE潜在空间分布
latent_dim = 32
prior = Normal(torch.zeros(latent_dim), torch.ones(latent_dim))
latent_z = prior.sample()  # 从先验采样

3.2 分布间的度量

PyTorch提供了计算分布间距离的函数：

python复制from torch.distributions import kl_divergence

# 定义两个正态分布
p = Normal(loc=0.0, scale=1.0)
q = Normal(loc=1.0, scale=1.5)

# 计算KL散度
kl_pq = kl_divergence(p, q)

# 计算交叉熵
cross_entropy = -torch.mean(p.log_prob(q.sample((1000,))))

KL散度在变分推断中特别重要。例如，在训练VAE时，我们需要最小化潜在变量分布与先验分布间的KL散度：

python复制# 计算VAE的KL损失
q_dist = Normal(q_mu, q_std)  # 编码器输出的分布
kl_loss = torch.mean(kl_divergence(q_dist, prior))

3.3 自定义分布

除了内置分布，我们还可以创建自定义分布：

python复制from torch.distributions import Distribution

class TruncatedNormal(Distribution):
    def __init__(self, loc, scale, low, high):
        self.normal = Normal(loc, scale)
        self.low = low
        self.high = high
        
    def sample(self, sample_shape):
        samples = self.normal.sample(sample_shape)
        return torch.clamp(samples, self.low, self.high)

这种灵活性使得PyTorch能够适应各种复杂的建模需求。例如，在强化学习中，我们可能需要限制动作空间的分布范围：

python复制action_dist = TruncatedNormal(action_mean, action_std, -1.0, 1.0)
actions = action_dist.sample((batch_size,))

4. 特殊数学函数

4.1 伽马函数和相关函数

PyTorch提供了一系列特殊数学函数：

python复制import torch.special as special

# 伽马函数
gamma_vals = special.gammaln(torch.linspace(0.1, 5.0, 10))  # ln(gamma(x))

# 贝塞尔函数
bessel_vals = special.i0(torch.arange(0.0, 5.0, 0.5))  # 第一类修正贝塞尔函数

# 误差函数
erf_vals = special.erf(torch.tensor([-1.0, 0.0, 1.0]))

这些函数在某些特定模型中非常有用。例如，在贝叶斯深度学习中使用Gamma分布作为共轭先验时：

python复制# Gamma分布参数估计
alpha = torch.tensor(2.0)
beta = torch.tensor(1.0)
log_normalizer = special.gammaln(alpha) - alpha * torch.log(beta)

4.2 激活函数相关的统计

许多激活函数本质上也是统计函数：

python复制# softmax函数
logits = torch.randn(10)
probs = torch.softmax(logits, dim=0)

# log_softmax (数值稳定版本)
log_probs = torch.log_softmax(logits, dim=0)

# sigmoid函数
sigmoid_vals = torch.sigmoid(torch.linspace(-5, 5, 10))

在分类任务中，我们经常使用log_softmax结合负对数似然损失：

python复制# 分类损失计算
log_probs = torch.log_softmax(model_output, dim=1)
nll_loss = -torch.mean(torch.sum(targets * log_probs, dim=1))

4.3 其他特殊函数

PyTorch还提供了许多其他有用的特殊函数：

python复制# 组合数学函数
combinations = special.comb(torch.tensor(10), torch.tensor(2))  # C(10,2)

# 正交多项式
x = torch.linspace(-1, 1, 100)
legendre_vals = special.legendre_polynomial_p(x, n=3)  # 3阶勒让德多项式

这些函数在特定领域非常有用。例如，在物理模拟中使用球谐函数时：

python复制# 球谐函数计算
theta = torch.linspace(0, torch.pi, 50)
phi = torch.linspace(0, 2*torch.pi, 50)
Y_lm = special.sph_harm(l=2, m=1, theta=theta, phi=phi)

5. 统计函数的应用场景

5.1 数据预处理

统计函数在数据标准化和归一化中扮演关键角色：

python复制# 标准化数据
def standardize(data):
    mean = torch.mean(data, dim=0)
    std = torch.std(data, dim=0)
    return (data - mean) / (std + 1e-8)

# 归一化到[0,1]
def normalize(data):
    min_val = torch.min(data, dim=0).values
    max_val = torch.max(data, dim=0).values
    return (data - min_val) / (max_val - min_val + 1e-8)

在实际项目中，我们通常会计算训练集的统计量，然后应用到测试集：

python复制# 计算训练集统计量
train_mean = torch.mean(train_data, dim=0)
train_std = torch.std(train_data, dim=0)

# 应用相同的变换到测试集
test_data_normalized = (test_data - train_mean) / train_std

5.2 模型初始化

统计函数常用于合理的权重初始化：

python复制# Xavier/Glorot初始化
def xavier_init(layer):
    fan_in, fan_out = layer.weight.shape
    std = torch.sqrt(torch.tensor(2.0 / (fan_in + fan_out)))
    layer.weight.data.normal_(0, std)
    
# Kaiming初始化
def kaiming_init(layer):
    fan_in = layer.weight.shape[1]
    std = torch.sqrt(torch.tensor(2.0 / fan_in))
    layer.weight.data.normal_(0, std)

这些初始化方法都基于对权重统计特性的分析，确保信号在网络中能够合理传播。

5.3 损失函数设计

许多损失函数本质上是统计度量：

python复制# 均方误差
def mse_loss(pred, target):
    return torch.mean((pred - target)**2)

# 交叉熵损失
def cross_entropy(pred, target):
    log_probs = torch.log_softmax(pred, dim=1)
    return -torch.mean(torch.sum(target * log_probs, dim=1))

# KL散度损失
def kl_loss(p, q):
    return torch.mean(torch.sum(p * (torch.log(p) - torch.log(q)), dim=1))

在自定义损失函数时，理解这些统计量非常重要。例如，在风格迁移任务中，我们可能会比较特征图的统计特性：

python复制# Gram矩阵计算
def gram_matrix(features):
    batch, channels, height, width = features.shape
    features_flat = features.view(batch, channels, -1)
    gram = torch.bmm(features_flat, features_flat.transpose(1, 2))
    return gram / (channels * height * width)

5.4 模型评估

统计函数在模型评估指标计算中无处不在：

python复制# 准确率
def accuracy(pred, target):
    pred_labels = torch.argmax(pred, dim=1)
    return torch.mean((pred_labels == target).float())

# 精确率和召回率
def precision_recall(pred, target, positive_class=1):
    pred_labels = torch.argmax(pred, dim=1)
    true_pos = torch.sum((pred_labels == positive_class) & (target == positive_class))
    false_pos = torch.sum((pred_labels == positive_class) & (target != positive_class))
    false_neg = torch.sum((pred_labels != positive_class) & (target == positive_class))
    
    precision = true_pos / (true_pos + false_pos + 1e-8)
    recall = true_pos / (true_pos + false_neg + 1e-8)
    return precision, recall

在多任务学习中，我们可能需要计算多个指标的加权平均：

python复制# 多任务评估
def multi_task_metrics(preds, targets, weights):
    metrics = {}
    for task in preds:
        metrics[f"{task}_acc"] = accuracy(preds[task], targets[task])
    
    avg_metric = torch.mean(torch.stack([metrics[name] * weights[name] 
                                       for name in metrics]))
    return avg_metric, metrics

6. 性能优化与高级技巧

6.1 批处理计算

利用PyTorch的向量化操作可以显著提高统计计算效率：

python复制# 低效的实现
def naive_mean(data):
    result = torch.zeros(data.shape[1])
    for i in range(data.shape[0]):
        result += data[i]
    return result / data.shape[0]

# 高效的向量化实现
def vectorized_mean(data):
    return torch.mean(data, dim=0)

在实际应用中，批处理计算可以带来数量级的性能提升：

python复制# 计算大型数据集的统计量（分批处理）
def compute_dataset_stats(dataloader):
    mean = 0.0
    std = 0.0
    count = 0
    
    for batch in dataloader:
        batch = batch[0]  # 假设dataloader返回(data, target)
        batch_mean = torch.mean(batch, dim=[0, 2, 3])
        batch_std = torch.std(batch, dim=[0, 2, 3])
        batch_count = batch.shape[0] * batch.shape[2] * batch.shape[3]
        
        delta = batch_mean - mean
        mean += delta * batch_count / (count + batch_count)
        std = (std * count + batch_std**2 * batch_count + delta**2 * count * batch_count / (count + batch_count)) / (count + batch_count)
        count += batch_count
    
    return mean, torch.sqrt(std)

6.2 GPU加速

PyTorch统计函数天然支持GPU加速：

python复制# 将数据和计算转移到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
large_data = large_data.to(device)

# GPU上的统计计算
gpu_mean = torch.mean(large_data)
gpu_std = torch.std(large_data)

对于非常大的数据集，我们可以使用并行计算：

python复制# 使用DataParallel进行并行统计计算
model = nn.DataParallel(statistical_model)
results = model(large_data)

6.3 自动微分兼容性

PyTorch统计函数的一个关键优势是它们支持自动微分：

python复制# 创建需要梯度的数据
x = torch.randn(10, requires_grad=True)

# 计算统计量并反向传播
y = torch.mean(x**2)
y.backward()

print(x.grad)  # 梯度为2x/10

这个特性使得我们可以在模型训练中使用复杂的统计量：

python复制# 在损失函数中使用高阶统计量
def custom_loss(output, target):
    error = output - target
    mse = torch.mean(error**2)
    skewness = moments.skewness(error)
    return mse + 0.1 * torch.abs(skewness)  # 惩罚不对称误差

6.4 数值稳定性

在实现统计函数时，数值稳定性至关重要：

python复制# 数值稳定的方差计算
def stable_var(data, dim=None):
    if dim is not None:
        mean = torch.mean(data, dim=dim, keepdim=True)
    else:
        mean = torch.mean(data)
    return torch.mean((data - mean)**2, dim=dim)

# 数值稳定的softmax
def stable_softmax(logits, dim=-1):
    logits = logits - torch.max(logits, dim=dim, keepdim=True).values
    exps = torch.exp(logits)
    return exps / torch.sum(exps, dim=dim, keepdim=True)

在处理极端值时，对数空间的计算往往更稳定：

python复制# 对数空间的计算
def log_space_normalize(log_probs):
    log_max = torch.max(log_probs, dim=-1, keepdim=True).values
    log_normalized = log_probs - log_max
    log_normalized = log_normalized - torch.log(torch.sum(torch.exp(log_normalized), dim=-1, keepdim=True))
    return log_normalized

7. 常见问题与解决方案

7.1 维度处理问题

统计函数最常见的困惑之一是维度处理：

python复制data = torch.randn(4, 3, 2)

# 错误的维度处理
try:
    wrong_mean = torch.mean(data, dim=1, keepdim=True)  # 输出形状 [4,1,2]
    # 后续操作可能因为形状不匹配而失败
except Exception as e:
    print(f"错误: {e}")

# 正确的做法是明确指定是否需要保持维度
mean_dim1 = torch.mean(data, dim=1)          # 输出形状 [4,2]
mean_dim1_keep = torch.mean(data, dim=1, keepdim=True)  # 输出形状 [4,1,2]

提示：使用keepdim=True可以保持张量的维度数不变，这在某些广播操作中非常有用。

7.2 NaN和Inf处理

现实数据中经常包含异常值：

python复制# 创建包含NaN和Inf的数据
data = torch.tensor([1.0, 2.0, float('nan'), float('inf'), -float('inf'), 3.0])

# 安全的统计计算
def nanmean(data):
    mask = ~torch.isnan(data)
    return torch.sum(data[mask]) / torch.sum(mask)

def finite_mean(data):
    mask = torch.isfinite(data)
    return torch.mean(data[mask])

在训练过程中，我们可以添加检查来捕获异常：

python复制# 训练循环中的统计检查
for batch in dataloader:
    if not torch.isfinite(batch).all():
        print("发现非有限值!")
        break
        
    # 训练代码...

7.3 随机性与可重复性

统计计算有时涉及随机采样：

python复制# 设置随机种子保证可重复性
torch.manual_seed(42)

# 可重复的采样
samples1 = torch.randn(10)
torch.manual_seed(42)
samples2 = torch.randn(10)
print(torch.allclose(samples1, samples2))  # 应该输出True

在分布式环境中，需要额外注意随机性控制：

python复制# 分布式环境中的随机种子设置
def set_seed(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

7.4 内存优化

处理大型数据时的内存管理技巧：

python复制# 内存高效的统计计算
def memory_efficient_mean(dataloader):
    total = 0.0
    count = 0
    for batch in dataloader:
        batch = batch[0]
        total += torch.sum(batch, dim=0)
        count += batch.shape[0]
        del batch  # 显式释放内存
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return total / count

对于特别大的数据集，可以考虑使用在线算法：

python复制# 在线计算均值和方差
def online_stats(iterable):
    n = 0
    mean = 0.0
    M2 = 0.0
    
    for x in iterable:
        n += 1
        delta = x - mean
        mean += delta / n
        delta2 = x - mean
        M2 += delta * delta2
        
    if n < 2:
        return mean, float('nan')
    else:
        return mean, M2 / n

8. 实际案例分析

8.1 图像风格迁移中的统计应用

在风格迁移中，Gram矩阵捕捉了特征的统计相关性：

python复制def gram_matrix(features):
    batch, channels, height, width = features.shape
    features_flat = features.view(batch, channels, -1)
    gram = torch.bmm(features_flat, features_flat.transpose(1, 2))
    return gram / (channels * height * width)

# 风格损失计算
def style_loss(style_features, content_features):
    loss = 0.0
    for style, content in zip(style_features, content_features):
        style_gram = gram_matrix(style)
        content_gram = gram_matrix(content)
        loss += torch.mean((style_gram - content_gram)**2)
    return loss

8.2 变分自编码器中的KL散度

VAE中使用KL散度作为正则项：

python复制def vae_loss(recon_x, x, mu, logvar):
    # 重构损失
    BCE = torch.mean((recon_x - x)**2)
    
    # KL散度
    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    
    return BCE + KLD

8.3 批归一化层的实现

批归一化本质上是统计标准化：

python复制class BatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.eps = eps
        self.momentum = momentum
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
    
    def forward(self, x):
        if self.training:
            mean = torch.mean(x, dim=0)
            var = torch.var(x, dim=0, unbiased=False)
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            mean = self.running_mean
            var = self.running_var
        
        x_normalized = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_normalized + self.beta

8.4 自注意力机制中的统计

自注意力机制中的softmax本质上是统计归一化：

python复制def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = torch.softmax(scores, dim=-1)
    return torch.matmul(p_attn, V), p_attn

9. 性能对比与基准测试

9.1 PyTorch与NumPy统计函数对比

虽然NumPy也有统计函数，但PyTorch的GPU支持使其在大数据上更有优势：

python复制import numpy as np
import time

# 创建大型数据
numpy_data = np.random.randn(10000, 10000)
torch_data = torch.from_numpy(numpy_data)

# NumPy计算
start = time.time()
np_mean = np.mean(numpy_data)
np_time = time.time() - start

# PyTorch CPU计算
start = time.time()
cpu_mean = torch.mean(torch_data)
cpu_time = time.time() - start

# PyTorch GPU计算
torch_data_gpu = torch_data.cuda()
start = time.time()
_ = torch.mean(torch_data_gpu)
torch.cuda.synchronize()  # 确保计时准确
gpu_time = time.time() - start

print(f"NumPy: {np_time:.4f}s, PyTorch CPU: {cpu_time:.4f}s, PyTorch GPU: {gpu_time:.4f}s")

9.2 不同实现的性能差异

有些统计函数有多种实现方式，性能可能不同：

python复制# 三种计算L2范数的方法
def l2_norm_1(x):
    return torch.sqrt(torch.sum(x**2))

def l2_norm_2(x):
    return torch.norm(x, p=2)

def l2_norm_3(x):
    return torch.linalg.norm(x)

# 性能测试
x = torch.randn(1000000)
%timeit l2_norm_1(x)
%timeit l2_norm_2(x)
%timeit l2_norm_3(x)

9.3 批处理大小的影响

批处理大小对统计计算性能有显著影响：

python复制batch_sizes = [16, 32, 64, 128, 256, 512]
times = []

for bs in batch_sizes:
    data = torch.randn(bs, 3, 256, 256).cuda()
    start = time.time()
    _ = torch.mean(data, dim=[0, 2, 3])
    torch.cuda.synchronize()
    times.append(time.time() - start)

10. 扩展与高级主题

10.1 自定义统计函数

我们可以创建自己的统计函数并确保它支持自动微分：

python复制def weighted_quantile(x, weights, q):
    """计算加权分位数"""
    sorted_x, sorted_weights = torch.sort(x), weights[torch.argsort(x)]
    cum_weights = torch.cumsum(sorted_weights, dim=0)
    cutoff = q * cum_weights[-1]
    return sorted_x[torch.searchsorted(cum_weights, cutoff)]

# 测试
x = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
weights = torch.tensor([0.1, 0.2, 0.3, 0.4])
q = weighted_quantile(x, weights, 0.5)
q.backward()

10.2 分布式统计计算

在大规模分布式环境中计算统计量：

python复制import torch.distributed as dist

def distributed_mean(tensor):
    local_sum = torch.sum(tensor)
    local_count = torch.tensor(tensor.numel())
    
    dist.all_reduce(local_sum, op=dist.ReduceOp.SUM)
    dist.all_reduce(local_count, op=dist.ReduceOp.SUM)
    
    return local_sum / local_count

10.3 流式统计计算

对于无法一次性加载到内存的数据，可以使用流式算法：

python复制class StreamingStats:
    def __init__(self):
        self.n = 0
        self.mean = 0.0
        self.M2 = 0.0
    
    def update(self, x):
        self.n += 1
        delta = x - self.mean
        self.mean += delta / self.n
        delta2 = x - self.mean
        self.M2 += delta * delta2
    
    @property
    def variance(self):
        return self.M2 / self.n if self.n > 1 else float('nan')
    
    @property
    def std(self):
        return torch.sqrt(self.variance)

10.4 统计函数的自动微分

理解统计函数的梯度行为很重要：

python复制x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = torch.var(x)
y.backward()
print(x.grad)  # 梯度为2(x - mean)/(n-1)

对于更复杂的统计量，我们可以手动定义梯度：

python复制class MyStatFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return torch.median(x)
    
    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        median_val = torch.median(x)
        mask = (x == median_val).float()
        return grad_output * mask / torch.sum(mask)

11. 调试与验证技巧

11.1 统计函数的单元测试

确保自定义统计函数正确性的方法：

python复制def test_weighted_quantile():
    x = torch.tensor([1.0, 2.0, 3.0, 4.0])
    weights = torch.tensor([0.25, 0.25, 0.25, 0.25])
    assert torch.allclose(weighted_quantile(x, weights, 0.5), torch.tensor(2.5))
    
    weights = torch.tensor([0.1, 0.2, 0.3, 0.4])
    assert torch.allclose(weighted_quantile(x, weights, 0.5), torch.tensor(3.0))

11.2 梯度检验

验证统计函数的梯度计算是否正确：

python复制from torch.autograd import gradcheck

# 创建输入
input = torch.randn(3, dtype=torch.double, requires_grad=True)

# 测试梯度
test = gradcheck(lambda x: torch.var(x, unbiased=True), (input,))
print("梯度检验通过:", test)

11.3 与参考实现对比

将PyTorch实现与已知正确的参考实现对比：

python复制def test_against_scipy():
    from scipy import stats
    data_np = np.random.randn(100)
    data_torch = torch.from_numpy(data_np)
    
    # 偏度对比
    scipy_skew = stats.skew(data_np)
    torch_skew = moments.skewness(data_torch)
    assert np.allclose(scipy_skew, torch_skew.numpy(), atol=1e-5)

11.4 可视化验证

使用可视化验证统计函数的正确性：

python复制import matplotlib.pyplot as plt

def plot_distribution_samples(dist, num_samples=10000):
    samples = dist.sample((num_samples,)).cpu().numpy()
    plt.hist(samples, bins=50, density=True)
    plt.title(f"{dist.__class__.__name__} 采样分布")
    plt.show()

# 测试
normal_dist = Normal(loc=0.0, scale=1.0)
plot_distribution_samples(normal_dist)

12. 最佳实践总结

12.1 统计函数使用原则

1. 明确计算目标：在调用统计函数前，明确你想要计算的是什么统计量，以及需要在哪个维度上计算。

2. 注意无偏估计：torch.var()等函数有unbiased参数控制是否使用无偏估计，在样本量小时差异明显。

3. 保持维度一致性：使用keepdim=True可以避免意外的广播行为，特别是在后续计算中使用统计量时。

4. GPU加速：对于大型数据，将计算转移到GPU可以带来显著加速。

5. 梯度考虑：如果需要在反向传播中使用统计量，确保操作是可微的。

12.2 性能优化建议

1. 向量化操作：尽量使用内置的向量化统计函数，避免手动编写循环。

2. 批处理计算：即使数据不能一次性装入内存，也应尽量使用较大的批处理大小。

3. 内存管理：对于中间结果及时使用del释放，并在CUDA上调用empty_cache()。

4. 选择合适精度：对于统计计算，有时可以使用torch.float32甚至torch.bfloat16来节省内存。

5. 并行计算：对于非常大的数据集，考虑使用多GPU或分布式计算。

12.3 调试与验证策略

1. 小数据测试：先用小规模数据验证统计函数的正确性。

2. 梯度检验：对于自定义统计函数，使用gradcheck验证梯度计算。

3. 参考实现对比：与NumPy、SciPy等成熟库的结果进行交叉验证。

4. 可视化检查：对于分布相关的统计量，通过可视化验证其行为。

5. 单元测试：为关键的统计计算编写单元测试，确保代码变更不会引入错误。

12.4 扩展学习资源

1. 官方文档：PyTorch的torch和torch.distributions模块文档是最权威的参考。

2. 统计学习：《All of Statistics》等教材可以夯实统计基础。

3. 源代码：PyTorch是开源项目，直接阅读统计函数的实现代码可以深入理解其行为。

4. 社区案例：PyTorch论坛和GitHub上有大量实际应用统计函数的案例。

5. 高级应用：研究变分推断、概率编程等主题可以深化统计函数在深度学习中的应用理解。