|
|
|
|
|
import torch |
|
|
from torch.amp import autocast |
|
|
from torch import Tensor |
|
|
from typing import Union, Tuple, Dict |
|
|
|
|
|
M_EPS = 1e-16 |
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
@autocast(device_type="cuda", enabled=True, dtype=torch.float32) |
|
|
def sinkhorn( |
|
|
a: Tensor, |
|
|
b: Tensor, |
|
|
C: Tensor, |
|
|
reg: float = 1e-1, |
|
|
maxIter: int = 1000, |
|
|
stopThr: float = 1e-9, |
|
|
verbose: bool = False, |
|
|
log: bool = True, |
|
|
eval_freq: int = 10, |
|
|
print_freq: int = 200, |
|
|
) -> Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]: |
|
|
device = a.device |
|
|
na, nb = C.shape |
|
|
assert na == a.shape[0] and nb == b.shape[0], f"Shapes of a ({a.shape}) or b ({b.shape}) do not match that of C ({C.shape})" |
|
|
assert reg > 0, f"reg should be greater than 0. Found reg = {reg}" |
|
|
assert a.min() >= 0. and b.min() >= 0., f"Elements in a and b should be nonnegative. Found a.min() = {a.min()}, b.min() = {b.min()}" |
|
|
|
|
|
if log: |
|
|
log = {"err": []} |
|
|
|
|
|
u = torch.ones(na, dtype=a.dtype, device=device) / na |
|
|
v = torch.ones(nb, dtype=b.dtype, device=device) / nb |
|
|
K = torch.exp(-C / reg) |
|
|
|
|
|
it, err = 1, 1 |
|
|
while (err > stopThr and it <= maxIter): |
|
|
u_pre, v_pre = u.clone(), v.clone() |
|
|
KTu = torch.matmul(K.T, u) |
|
|
v = b / (KTu + M_EPS) |
|
|
Kv = torch.matmul(K, v) |
|
|
u = a / (Kv + M_EPS) |
|
|
|
|
|
if torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or torch.any(torch.isinf(u)) or torch.any(torch.isinf(v)): |
|
|
print("Warning: numerical errors at iteration", it) |
|
|
u, v = u_pre, v_pre |
|
|
break |
|
|
|
|
|
if log and it % eval_freq == 0: |
|
|
b_hat = torch.matmul(u, K) * v |
|
|
err = (b - b_hat).pow(2).sum().item() |
|
|
log["err"].append(err) |
|
|
|
|
|
if verbose and it % print_freq == 0: |
|
|
print(f"Iteration {it}, constraint error {err}") |
|
|
|
|
|
it += 1 |
|
|
|
|
|
if log: |
|
|
log["u"] = u |
|
|
log["v"] = v |
|
|
log["alpha"] = reg * torch.log(u + M_EPS) |
|
|
log["beta"] = reg * torch.log(v + M_EPS) |
|
|
|
|
|
P = u.view(-1, 1) * K * v.view(1, -1) |
|
|
if log: |
|
|
return P, log |
|
|
else: |
|
|
return P |
|
|
|
