46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from mmselfsup.utils.clustering import PIC, Kmeans
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_faiss_in_clutering():
|
|
with patch('mmselfsup.utils.clustering.faiss') as faiss:
|
|
yield faiss
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_faiss(mock_faiss_in_clutering):
|
|
mock_PCAmatrix = mock_faiss_in_clutering.PCAMatrix.return_value
|
|
mock_GpuIndexFlatL2 = mock_faiss_in_clutering.GpuIndexFlatL2.return_value
|
|
|
|
mock_PCAmatrix.apply_py.return_value = np.random.rand(10, 8)
|
|
mock_GpuIndexFlatL2.search.return_value = (
|
|
np.random.rand(1000, 6),
|
|
np.random.rand(1000, 6),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize('verbose', [True, False])
|
|
def test_kmeans(mock_faiss, verbose):
|
|
fake_input = np.random.rand(10, 8).astype(np.float32)
|
|
pca_dim = 2
|
|
|
|
kmeans = Kmeans(2, pca_dim)
|
|
loss = kmeans.cluster(fake_input, verbose=verbose)
|
|
assert loss is not None
|
|
|
|
with pytest.raises(AssertionError):
|
|
loss = kmeans.cluster(np.random.rand(10, 8), verbose=verbose)
|
|
|
|
|
|
def test_pic(mock_faiss):
|
|
fake_input = np.random.rand(1000, 16).astype(np.float32)
|
|
pic = PIC(pca_dim=8)
|
|
res = pic.cluster(fake_input)
|
|
assert res == 0
|