def forward(pred: ProbabilityMeasure, targets: ProbabilityMeasure):
with torch.no_grad():
P = compute_ot_matrix_par(pred.centered().coord.cpu().numpy(),
targets.centered().coord.cpu().numpy())
P = torch.from_numpy(P).type_as(pred.coord).cuda()
xs = pred.centered().coord
xsT = xs.transpose(1, 2)
xt = targets.centered().coord
a: Tensor = pred.probability + 1e-8
a /= a.sum(dim=1, keepdim=True)
a = a.reshape(a.shape[0], -1, 1)
A = torch.inverse(xsT.bmm(a * xs)).bmm(xsT.bmm(P.bmm(xt)))
T = targets.mean() - pred.mean()
return A.type_as(pred.coord), T.detach()
def forward(pred: ProbabilityMeasure,
targets: ProbabilityMeasure,
iters: int = 200):
lambd = 0.002
with torch.no_grad():
P = SOT(iters, lambd).forward(pred.centered(), targets.centered())
xs = pred.centered().coord
xsT = xs.transpose(1, 2)
xt = targets.centered().coord
a = pred.probability + 1e-8
a /= a.sum(dim=1, keepdim=True)
a = a.reshape(a.shape[0], -1, 1)
A = torch.inverse(xsT.bmm(a * xs)).bmm(xsT.bmm(P.bmm(xt)))
T = targets.mean() - pred.mean()
return A.type_as(pred.coord), T.detach()