diff options
| -rw-r--r-- | ot/gpu/bregman.py | 10 | ||||
| -rw-r--r-- | test/test_gpu.py | 89 |
2 files changed, 53 insertions, 46 deletions
diff --git a/ot/gpu/bregman.py b/ot/gpu/bregman.py index 6714098..600ead4 100644 --- a/ot/gpu/bregman.py +++ b/ot/gpu/bregman.py @@ -90,14 +90,14 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9, """ - a = cp.asarray(a, dtype=np.float64) - b = cp.asarray(b, dtype=np.float64) - M = cp.asarray(M, dtype=np.float64) + a = cp.asarray(a) + b = cp.asarray(b) + M = cp.asarray(M) if len(a) == 0: - a = np.ones((M.shape[0],), dtype=np.float64) / M.shape[0] + a = np.ones((M.shape[0],)) / M.shape[0] if len(b) == 0: - b = np.ones((M.shape[1],), dtype=np.float64) / M.shape[1] + b = np.ones((M.shape[1],)) / M.shape[1] # init data Nini = len(a) diff --git a/test/test_gpu.py b/test/test_gpu.py index 1e97c45..51a0cff 100644 --- a/test/test_gpu.py +++ b/test/test_gpu.py @@ -17,63 +17,70 @@ except ImportError: @pytest.mark.skipif(nogpu, reason="No GPU available") -def test_gpu_sinkhorn(): +def test_gpu_dist(): rng = np.random.RandomState(0) - def describe_res(r): - print("min:{:.3E}, max::{:.3E}, mean::{:.3E}, std::{:.3E}".format( - np.min(r), np.max(r), np.mean(r), np.std(r))) - for n_samples in [50, 100, 500, 1000]: print(n_samples) a = rng.rand(n_samples // 4, 100) b = rng.rand(n_samples, 100) - time1 = time.time() - transport = ot.da.OTDA_sinkhorn() - transport.fit(a, b) - G1 = transport.G - time2 = time.time() - transport = ot.gpu.da.OTDA_sinkhorn() - transport.fit(a, b) - G2 = transport.G - time3 = time.time() - print("Normal sinkhorn, time: {:6.2f} sec ".format(time2 - time1)) - describe_res(G1) - print(" GPU sinkhorn, time: {:6.2f} sec ".format(time3 - time2)) - describe_res(G2) - - np.testing.assert_allclose(G1, G2, rtol=1e-5, atol=1e-5) + + M = ot.dist(a.copy(), b.copy()) + M2 = ot.gpu.dist(a.copy(), b.copy()) + + np.testing.assert_allclose(M, M2, rtol=1e-10) + + M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False) @pytest.mark.skipif(nogpu, reason="No GPU available") -def test_gpu_sinkhorn_lpl1(): +def test_gpu_sinkhorn(): rng = np.random.RandomState(0) - def describe_res(r): - print("min:{:.3E}, max:{:.3E}, mean:{:.3E}, std:{:.3E}" - .format(np.min(r), np.max(r), np.mean(r), np.std(r))) + for n_samples in [50, 100, 500, 1000]: + a = rng.rand(n_samples // 4, 100) + b = rng.rand(n_samples, 100) + + wa = ot.unif(n_samples // 4) + wb = ot.unif(n_samples) + + M = ot.dist(a.copy(), b.copy()) + M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False) + + reg = 1 + + G = ot.sinkhorn(wa, wb, M, reg) + G1 = ot.gpu.sinkhorn(wa, wb, M, reg) + + np.testing.assert_allclose(G1, G, rtol=1e-10) + + G2 = ot.gpu.sinkhorn(wa, wb, M2, reg, to_numpy=False) + + +@pytest.mark.skipif(nogpu, reason="No GPU available") +def test_gpu_sinkhorn_lpl1(): + + rng = np.random.RandomState(0) for n_samples in [50, 100, 500]: print(n_samples) a = rng.rand(n_samples // 4, 100) labels_a = np.random.randint(10, size=(n_samples // 4)) b = rng.rand(n_samples, 100) - time1 = time.time() - transport = ot.da.OTDA_lpl1() - transport.fit(a, labels_a, b) - G1 = transport.G - time2 = time.time() - transport = ot.gpu.da.OTDA_lpl1() - transport.fit(a, labels_a, b) - G2 = transport.G - time3 = time.time() - print("Normal sinkhorn lpl1, time: {:6.2f} sec ".format( - time2 - time1)) - describe_res(G1) - print(" GPU sinkhorn lpl1, time: {:6.2f} sec ".format( - time3 - time2)) - describe_res(G2) - - np.testing.assert_allclose(G1, G2, rtol=1e-3, atol=1e-3) + + wa = ot.unif(n_samples // 4) + wb = ot.unif(n_samples) + + M = ot.dist(a.copy(), b.copy()) + M2 = ot.gpu.dist(a.copy(), b.copy(), to_numpy=False) + + reg = 1 + + G = ot.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M, reg) + G1 = ot.gpu.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M, reg) + + np.testing.assert_allclose(G1, G, rtol=1e-10) + + G2 = ot.gpu.da.sinkhorn_lpl1_mm(wa, labels_a, wb, M2, reg, to_numpy=False) |