代码实现sinkhorn算法

INEVGVUP

已于 2023-08-01 22:10:25 修改

阅读量1.5k

点赞数 1

文章标签：算法深度学习

于 2023-08-01 16:45:35 首次发布

本文链接：https://blog.youkuaiyun.com/explorer5568/article/details/132044862

版权

文章介绍了Sinkhorn算法在最优运输问题中的应用，包括手动实现和使用Python库POT求解实例。作者还讨论了手写实现代码和调用专门库的优缺点，并提到未来可能的GPU加速方向。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

sinkhorn算法有很多的应用，可以计算深度学习任务中一些分布之间的距离，也可以计算最优运输问题。本篇文章的侧重点在于利用sinkhorn算法求解最优运输问题，类似于以下这个问题：有 $m$ 个供应商， $n$ 个需求商，并且已知任意两个供应商与需求商之间的运输代价 $M_{ij}$ ，求解运输矩阵 $T$ ，如果读者的需求就是这个，那么看完本文绝对能实现读者的需求，如果是处理深度学习任务中的一些度量问题，看“调用已有库部分”也能有所收获。

本文从两个路径来分享思路，第一个就是靠自己实现代码；第二个就是笔者在探索的过程中，发现了专门解决最优传输问题的库POT(Python Optimal Transport)，那么第二个就是如何使用这个库来解决此问题

手写实现

整体能跑的代码如下：

import numpy as np
import matplotlib.pyplot as plt


source = np.array([0.4, 0.6])
target = np.array([0.3, 0.3, 0.4])
m = source.shape[0]
n =target.shape[0]
source = source.reshape((source.shape[0], 1))
target = target.reshape((target.shape[0], 1))


cost = np.array([[1,1,1],[2,3,4]])
gamma = 0.01
maxiters = 1000
M = np.exp(-cost / gamma)

mu = np.ones((m, 1))
nu = np.ones((n, 1))

thresh = 10 ** (-1)  # stopping criterion
last_mu = np.ones((m, 1)) * -1

for k in range(maxiters):
    mu = np.divide(source, np.dot(M, nu))
    nu = np.divide(target, np.dot(M.T, mu))
    pi = np.dot(
        np.dot(
            np.diag(mu.reshape(-1)),
            M),
        np.diag(nu.reshape(-1))
    )

    sinkhorn_distance = np.sqrt(sum(sum(np.multiply(pi, cost))))
    print(sinkhorn_distance)


print("sinkhorn distance: ", sinkhorn_distance)
print("transport matrix: ", pi)

代码分析：

source = np.array([0.4, 0.6])
target = np.array([0.3, 0.3, 0.4])
m = source.shape[0]
n =target.shape[0]
source = source.reshape((source.shape[0], 1))
target = target.reshape((target.shape[0], 1))


cost = np.array([[1,1,1],[2,3,4]])
gamma = 0.01
maxiters = 1000

这部分代码就是简单创建一下需要解决的问题，供应商向量为 $[0.4, 0.6]$ ，需求商向量为 $[0.3, 0.3, 0.4]$ ，然后分别将其转换成 $m * 1$ 和 $n * 1$ 的向量。 $cos t$ 是声明的代价矩阵， $g amma$ 是sinkhorn算法中的正则化系数， $ma x i t ers$ 是最大的迭代次数。

然后sinkhorn算法核心就是进行迭代，具体推导细节在本文中就不进行叙述，迭代的核心就是：设 $r$ 为供应商向量， $c$ 为需求商向量，一开始将 $u$ 和 $v$ 初始化，然后每次按照 $u = r ./ K v$ , $v = c ./ K u$ 进行迭代即可，其中 $K$ 是指数化处理后的代价矩阵，那核心代码如下所示：

for k in range(maxiters):
    mu = np.divide(source, np.dot(M, nu))
    nu = np.divide(target, np.dot(M.T, mu))
    pi = np.dot(
        np.dot(
            np.diag(mu.reshape(-1)),
            M),
        np.diag(nu.reshape(-1))
    )

    sinkhorn_distance = np.sqrt(sum(sum(np.multiply(pi, cost))))
    print(sinkhorn_distance)

最后运行结果如下所示：

在这里插入图片描述

调用已有库

笔者在GitHub上找各种实现的时候，找到了专门处理OT问题的python库，链接如下：

Github地址

文档地址

调用此python库来解决上面相同的问题，代码如下：

import ot
import numpy as np

a = np.array([0.4, 0.6])
b = np.array([0.3, 0.3, 0.4])

M = np.array([[1,1,1],[2,3,4]])
reg = 0.1

T = ot.emd(a, b, M) # exact linear program
T_reg = ot.sinkhorn(a, b, M, reg) # entropic regularized OT

print("T: ", T)
print("T_reg: ", T_reg)
print(np.sqrt(sum(sum(np.multiply(M, T_reg)))))

运行结果如下所示：
在这里插入图片描述

可以看到调用此库非常简单，读者对着上述代码即可知道各个参数的含义，如果读者想详细了解此库的用法，可以进行查阅。

笔者后续应该有用调用GPU来实现sinkhorn的需求，目前有两种思路，第一种就是将自己写的那部分代码改为由pytorch实现，进而调用GPU；另一种就是将库中的sinkhorn代码复制出来，将其改为由pytorch进而调用GPU实现，因为笔者暂时没看到怎么利用此库来利用GPU进行求解，并且查看库中的代码，发现代码之间的耦合性不强，实现较为简单，复制出来进行改写难度也不大。如果读者知道现成的调用GPU的方法，可以评论分享方法。笔者找到解决方案之后也会过来进行补充！