Python实战：用Sinkhorn算法搞定最优传输问题（附完整代码）

Python实战：用Sinkhorn算法搞定最优传输问题（附完整代码）

最优传输问题在机器学习领域的重要性与日俱增，而Sinkhorn算法凭借其高效稳定的特性成为解决这类问题的利器。本文将带你从零开始实现Sinkhorn算法，并应用于实际数据分布匹配场景。

1. 理解最优传输与Sinkhorn算法

最优传输问题的核心是找到将一种概率分布转换为另一种概率分布的最小成本方案。想象你是一家物流公司的经理，需要将多个仓库的货物分配到各个零售店，同时希望运输成本最低——这就是最优传输问题的现实映射。

传统解法如线性规划在大规模问题上计算成本高昂。Sinkhorn算法通过引入熵正则化项：

H(P) = -ΣP_ij(logP_ij - 1)

将问题转化为可通过迭代矩阵缩放解决的优化问题。其优势在于：

• 计算效率：复杂度从O(n³)降至O(n²)
• 数值稳定：熵正则化避免极端解
• 并行友好：适合GPU加速实现

提示：正则化参数ε控制精度与速度的平衡，通常取0.01-0.1

2. 算法实现关键步骤

2.1 核心迭代过程

Sinkhorn算法的精髓在于交替更新缩放因子u和v：

def sinkhorn_iteration(K, a, b, max_iter=1000, tol=1e-9): u = np.ones_like(a) v = np.ones_like(b) for _ in range(max_iter): u_prev, v_prev = u.copy(), v.copy() # 交替更新 u = a / (K @ v) v = b / (K.T @ u) # 收敛判断 if (np.max(np.abs(u - u_prev)) < tol and np.max(np.abs(v - v_prev)) < tol): break return u, v

2.2 成本矩阵构建

成本矩阵C的设计直接影响传输结果。常见选择包括：

# 构建欧式距离成本矩阵 def euclidean_cost(X, Y): return np.sqrt(np.sum((X[:,None] - Y)**2, axis=2))

3. 完整实现与调优技巧

3.1 完整算法封装

class SinkhornTransport: def __init__(self, epsilon=0.1, max_iter=1000): self.epsilon = epsilon self.max_iter = max_iter def fit(self, a, b, C): K = np.exp(-C / self.epsilon) u, v = sinkhorn_iteration(K, a, b, self.max_iter) self.P_ = np.diag(u) @ K @ np.diag(v) return self

3.2 参数调优指南

• ε选择：通过网格搜索验证不同值的效果
• 迭代控制：监控对偶间隙判断收敛
• 数值稳定：添加小常数防止除零错误

注意：实际实现时应添加对数域计算优化，避免数值下溢

4. 实战案例：图像颜色迁移

让我们将算法应用于实际场景——将一幅图像的色彩风格迁移到另一幅图像：

def color_transfer(source_img, target_img, epsilon=0.01): # 将图像转为RGB分布 src_pixels = source_img.reshape(-1, 3) tgt_pixels = target_img.reshape(-1, 3) # 构建颜色分布直方图 a = np.histogramdd(src_pixels, bins=32)[0].ravel() b = np.histogramdd(tgt_pixels, bins=32)[0].ravel() a, b = a/a.sum(), b/b.sum() # 计算颜色距离矩阵 bin_centers = [np.linspace(0, 255, 32) for _ in range(3)] C = euclidean_cost(bin_centers, bin_centers) # 计算传输计划 st = SinkhornTransport(epsilon=epsilon).fit(a, b, C) # 应用传输矩阵 transferred = st.P_.dot(tgt_pixels) return transferred.reshape(target_img.shape)

5. 性能优化与扩展

5.1 GPU加速实现

使用PyTorch可以轻松实现GPU加速：

import torch def sinkhorn_gpu(a, b, C, epsilon=0.1, device='cuda'): K = torch.exp(-C.to(device)/epsilon) u = torch.ones_like(a).to(device) v = torch.ones_like(b).to(device) for _ in range(1000): u = a.to(device) / (K @ v) v = b.to(device) / (K.T @ u) return (u[:,None] * K * v[None,:]).cpu()

5.2 与其他方法的对比

我们在MNIST数据集上对比了不同方法：

实际项目中，Sinkhorn算法在保持精度的同时显著提升了计算效率。特别是在处理高维数据时，差异更为明显。