Function graph_compression_gw

examples/backends/plot_optim_gromov_pytorch.py:200–239

solve min_a GW(C1,C2,a, a2) by gradient descent

(nb_nodes, C2, a2, nb_iter_max=100, lr=1e-2)

def graph_compression_gw(nb_nodes, C2, a2, nb_iter_max=100, lr=1e-2):
    """solve min_a GW(C1,C2,a, a2) by gradient descent"""

    # use pyTorch for our data

    C2_torch = torch.tensor(C2)
    a2_torch = torch.tensor(a2)

    a0 = rng.rand(nb_nodes)  # random_init
    a0 /= a0.sum()  # on simplex
    a1_torch = torch.tensor(a0).requires_grad_(True)
    C0 = np.eye(nb_nodes)
    C1_torch = torch.tensor(C0).requires_grad_(True)

    loss_iter = []

    for i in range(nb_iter_max):
        loss = gromov_wasserstein2(C1_torch, C2_torch, a1_torch, a2_torch)

        loss_iter.append(loss.clone().detach().cpu().numpy())
        loss.backward()

        # print("{:03d} | {}".format(i, loss_iter[-1]))

        # performs a step of projected gradient descent
        with torch.no_grad():
            grad = a1_torch.grad
            a1_torch -= grad * lr  # step
            a1_torch.grad.zero_()
            a1_torch.data = ot.utils.proj_simplex(a1_torch)

            grad = C1_torch.grad
            C1_torch -= grad * lr  # step
            C1_torch.grad.zero_()
            C1_torch.data = torch.clamp(C1_torch, 0, 1)

    a1 = a1_torch.clone().detach().cpu().numpy()
    C1 = C1_torch.clone().detach().cpu().numpy()

    return a1, C1, loss_iter


nb_nodes = 3