Python实战：用POT库搞定Wasserstein距离计算（附完整代码示例）

Python实战用POT库高效计算Wasserstein距离的完整指南1. 初识Wasserstein距离与POT库在数据科学和机器学习领域衡量两个概率分布之间的差异是一个基础而重要的问题。传统方法如KL散度存在对称性和支撑集限制等问题而Wasserstein距离又称Earth Movers Distance因其优异的几何特性逐渐成为研究热点。这个距离度量直观上可以理解为将一个分布的质量搬运到另一个分布所需的最小工作量。Python Optimal TransportPOT库由Rémi Flamary和Nicolas Courty等人开发是目前最全面的最优传输Python工具包。与scipy.stats中的基础实现相比POT提供了完整的最优传输算法体系从精确线性规划到正则化近似方法GPU加速支持针对大规模计算任务的硬件加速丰富的应用接口包含图像处理、领域适应等前沿应用场景活跃的社区维护持续更新最新研究成果的算法实现安装POT库非常简单只需运行pip install pot对于需要GPU加速的用户建议额外安装cudamatgit clone https://github.com/cudamat/cudamat.git cd cudamat python setup.py install --user2. 核心算法原理与实现对比2.1 传统线性规划方法精确计算Wasserstein距离可以转化为线性规划问题。POT提供了ot.emd()函数实现这一算法import ot import numpy as np # 创建两个一维概率分布 a np.array([0.4, 0.6]) b np.array([0.3, 0.7]) # 计算成本矩阵这里使用欧式距离的平方 M ot.dist(np.arange(2).reshape(-1,1), np.arange(2).reshape(-1,1), metricsqeuclidean) # 计算最优传输矩阵 T ot.emd(a, b, M) # 计算Wasserstein距离 W_dist np.sum(T * M)这种方法虽然精确但时间复杂度高达O(n³)仅适用于小规模问题n1000。实际应用中我们更常使用正则化方法。2.2 Sinkhorn正则化算法Sinkhorn算法通过引入熵正则项将问题转化为可并行计算的迭代形式复杂度降至O(n²)# 设置正则化参数 reg 0.1 # Sinkhorn算法计算传输矩阵 T_reg ot.sinkhorn(a, b, M, reg) # 计算正则化Wasserstein距离 W_dist_reg ot.sinkhorn2(a, b, M, reg)两种方法的性能对比如下指标精确EMDSinkhorn(reg0.1)计算时间(1000点)35.2s0.8s内存占用高中等结果精度精确近似(误差1e-3)可并行性差优秀提示正则化参数reg的选择需要权衡计算速度和精度通常建议在0.01-0.5范围内尝试3. 实战案例图像风格迁移中的应用让我们通过一个完整的图像处理案例展示Wasserstein距离的实际价值。假设我们要将一幅图像的色彩分布迁移到另一幅图像import matplotlib.pyplot as plt from skimage import io # 加载示例图像 src_img io.imread(source.jpg) / 255. target_img io.imread(target.jpg) / 255. # 将图像像素视为三维(RGB)空间中的点 src_pixels src_img.reshape(-1, 3) target_pixels target_img.reshape(-1, 3) # 随机采样以减少计算量 n_samples 1000 idx np.random.randint(0, len(src_pixels), n_samples) src_samples src_pixels[idx] target_samples target_pixels[idx] # 创建均匀分布 a, b np.ones(n_samples)/n_samples, np.ones(n_samples)/n_samples # 计算RGB空间中的成本矩阵 M ot.dist(src_samples, target_samples) # 计算最优传输 reg 0.01 T ot.sinkhorn(a, b, M, reg) # 应用色彩变换 transferred_pixels n_samples * T.T.dot(src_samples) transferred_img transferred_pixels.reshape(target_img.shape) # 可视化结果 fig, axes plt.subplots(1, 3, figsize(15,5)) axes[0].imshow(src_img) axes[1].imshow(target_img) axes[2].imshow(transferred_img) plt.show()这个案例展示了Wasserstein距离在捕捉分布几何特性方面的优势——它不仅能匹配颜色直方图还能保持颜色之间的相对关系。4. 高级技巧与性能优化4.1 一维分布的特化算法对于一维分布Wasserstein距离有闭合解计算复杂度仅为O(n log n)# 生成两个一维分布 n 10000 x np.random.normal(0, 1, n) y np.random.normal(1, 2, n) # 快速计算一维Wasserstein距离 W_dist_1d ot.wasserstein_1d(x, y)4.2 批处理与GPU加速POT支持使用PyTorch进行GPU加速特别适合深度学习场景import torch import ot # 创建PyTorch张量 a torch.tensor([0.5, 0.5], devicecuda) b torch.tensor([0.3, 0.7], devicecuda) M torch.tensor([[0, 1], [1, 0]], devicecuda) # GPU加速的Sinkhorn计算 T ot.sinkhorn(a, b, M, reg0.1)4.3 不平衡最优传输当两个分布的总质量不相等时可以使用不平衡最优传输# 创建质量不相等的分布 a np.array([0.8, 0.8]) # 总质量为1.6 b np.array([0.5, 0.5]) # 总质量为1.0 # 计算不平衡最优传输 reg_m 0.1 # 质量松弛参数 T_unbalanced ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg, reg_m)5. 结果可视化与评估理解Wasserstein距离计算结果的关键是可视化传输方案。以下代码展示了如何可视化两个二维分布之间的最优传输def plot_2D_OT(xs, xt, T): plt.figure(figsize(8,6)) # 绘制样本点 plt.scatter(xs[:,0], xs[:,1], cblue, labelSource) plt.scatter(xt[:,0], xt[:,1], cred, labelTarget) # 绘制传输路径 for i in range(T.shape[0]): for j in range(T.shape[1]): if T[i,j] 1e-4: # 只显示显著传输 plt.plot([xs[i,0], xt[j,0]], [xs[i,1], xt[j,1]], k-, alphaT[i,j]/T.max()*0.8) plt.legend() plt.title(Optimal Transport Plan) plt.show() # 生成二维样本 xs np.random.randn(50, 2) xt np.random.randn(50, 2) np.array([3,0]) # 计算传输矩阵 a, b np.ones(50)/50, np.ones(50)/50 M ot.dist(xs, xt) T ot.emd(a, b, M) # 可视化 plot_2D_OT(xs, xt, T)这种可视化帮助我们直观理解Wasserstein距离的几何意义——它寻找的是样本点之间的最优对应关系而不仅仅是简单的分布矩匹配。