implicit/tests/gpu_test.py at 73221b103c61e4ac40998882a12c2c97faee0089 · benfred/implicit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import numpy as np
import pytest
from numpy.testing import assert_allclose, assert_array_equal
from recommender_base_test import get_checker_board
from scipy.sparse import coo_matrix

import implicit


@pytest.mark.skipif(not implicit.gpu.HAS_CUDA, reason="needs cuda build")
@pytest.mark.parametrize("k", [4, 16, 64, 128, 1000])
@pytest.mark.parametrize("batch", [1, 10, 100])
@pytest.mark.parametrize("temp_memory", [500_000_000, 5_000_000])
def test_topk_ascending(k, batch, temp_memory):
    num_items = 10000
    factors = 10
    items = np.arange(num_items * factors).reshape((num_items, factors)).astype("float32")
    queries = np.arange(batch * factors).reshape((batch, factors)).astype("float32")
    _check_knn_queries(items, queries, k, max_temp_memory=temp_memory)


@pytest.mark.skipif(not implicit.gpu.HAS_CUDA, reason="needs cuda build")
@pytest.mark.parametrize("k", [4, 64])
@pytest.mark.parametrize("batch", [1, 10, 100])
@pytest.mark.parametrize("temp_memory", [500_000_000, 500_000])
def test_topk_random(k, batch, temp_memory):
    num_items = 1000
    factors = 10

    rs = np.random.default_rng(0)
    items = rs.random(size=(num_items, factors), dtype="float32")
    queries = rs.random(size=(batch, factors), dtype="float32")
    _check_knn_queries(items, queries, k, max_temp_memory=temp_memory)


def _check_knn_queries(items, queries, k=5, max_temp_memory=500_000_000):
    # compute distances on the gpu
    knn = implicit.gpu._cuda.KnnQuery(max_temp_memory=max_temp_memory)
    ids, distances = knn.topk(
        implicit.gpu._cuda.Matrix(items), implicit.gpu._cuda.Matrix(queries), k
    )

    # compute on the cpu
    batch = queries.dot(items.T)
    exact_ids = np.flip(np.argsort(batch)[:, -k:], axis=1)
    exact_distances = np.zeros(exact_ids.shape)
    for r in range(batch.shape[0]):
        exact_distances[r] = batch[r][exact_ids[r]]

    # make sure that we match
    assert_allclose(distances, exact_distances, rtol=1e-06)
    assert_array_equal(ids, exact_ids)


@pytest.mark.skipif(not implicit.gpu.HAS_CUDA, reason="needs cuda build")
def test_calculate_norms():
    num_items = 100
    factors = 8
    items = np.arange(num_items * factors).reshape((num_items, factors)).astype("float32")
    norms = (
        implicit.gpu._cuda.calculate_norms(implicit.gpu._cuda.Matrix(items))
        .to_numpy()
        .reshape(num_items)
    )
    np_norms = np.linalg.norm(items, axis=1)
    assert_allclose(norms, np_norms)


@pytest.mark.skipif(not implicit.gpu.HAS_CUDA, reason="needs cuda build")
@pytest.mark.parametrize(
    "model_class", [implicit.als.AlternatingLeastSquares, implicit.bpr.BayesianPersonalizedRanking]
)
@pytest.mark.parametrize("from_gpu", [True, False])
def test_cpu_gpu_conversion(model_class, from_gpu):
    model = model_class(use_gpu=from_gpu, factors=32)
    user_plays = get_checker_board(50)
    model.fit(user_plays)
    converted = model.to_cpu() if from_gpu else model.to_gpu()
    assert_allclose(
        model.recommend(0, user_plays[0]),
        converted.recommend(0, user_plays[0]),
        rtol=1e-3,
        atol=1e-3,
    )


@pytest.mark.skipif(not implicit.gpu.HAS_CUDA, reason="needs cuda build")
def test_coo_matrix_copies_float_data():
    row = np.array([0, 1, 2, 0, 1, 2], dtype=np.int32)
    col = np.array([0, 1, 2, 1, 2, 0], dtype=np.int32)
    data = np.array([1.5, 2.7, 3.14159, 0.1, 1000.5, -1.5], dtype=np.float32)

    matrix = implicit.gpu.COOMatrix(coo_matrix((data, (row, col)), shape=(3, 3)))

    assert matrix is not None