代码实现sinkhorn算法

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\ast 1$和$n\ast 1$的向量。$cost$是声明的代价矩阵，$gamma$是sinkhorn算法中的正则化系数，$maxiters$是最大的迭代次数。

然后sinkhorn算法核心就是进行迭代，具体推导细节在本文中就不进行叙述，迭代的核心就是：设$r$为供应商向量，$c$为需求商向量，一开始将$u$和$v$初始化，然后每次按照$u=r./Kv$, $v=c./Ku$进行迭代即可，其中$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的方法，可以评论分享方法。笔者找到解决方案之后也会过来进行补充！