90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from pytorch_lightning import seed_everything
|
|
from pytorch_lightning.metrics.functional.classification import dice_score
|
|
from pytorch_lightning.metrics.functional.precision_recall_curve import _binary_clf_curve
|
|
from pytorch_lightning.metrics.utils import get_num_classes, to_categorical, to_onehot
|
|
|
|
|
|
def test_onehot():
|
|
test_tensor = torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
|
|
expected = torch.stack([
|
|
torch.cat([torch.eye(5, dtype=int), torch.zeros((5, 5), dtype=int)]),
|
|
torch.cat([torch.zeros((5, 5), dtype=int), torch.eye(5, dtype=int)])
|
|
])
|
|
|
|
assert test_tensor.shape == (2, 5)
|
|
assert expected.shape == (2, 10, 5)
|
|
|
|
onehot_classes = to_onehot(test_tensor, num_classes=10)
|
|
onehot_no_classes = to_onehot(test_tensor)
|
|
|
|
assert torch.allclose(onehot_classes, onehot_no_classes)
|
|
|
|
assert onehot_classes.shape == expected.shape
|
|
assert onehot_no_classes.shape == expected.shape
|
|
|
|
assert torch.allclose(expected.to(onehot_no_classes), onehot_no_classes)
|
|
assert torch.allclose(expected.to(onehot_classes), onehot_classes)
|
|
|
|
|
|
def test_to_categorical():
|
|
test_tensor = torch.stack([
|
|
torch.cat([torch.eye(5, dtype=int), torch.zeros((5, 5), dtype=int)]),
|
|
torch.cat([torch.zeros((5, 5), dtype=int), torch.eye(5, dtype=int)])
|
|
]).to(torch.float)
|
|
|
|
expected = torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
|
|
assert expected.shape == (2, 5)
|
|
assert test_tensor.shape == (2, 10, 5)
|
|
|
|
result = to_categorical(test_tensor)
|
|
|
|
assert result.shape == expected.shape
|
|
assert torch.allclose(result, expected.to(result.dtype))
|
|
|
|
|
|
@pytest.mark.parametrize(['pred', 'target', 'num_classes', 'expected_num_classes'], [
|
|
pytest.param(torch.rand(32, 10, 28, 28), torch.randint(10, (32, 28, 28)), 10, 10),
|
|
pytest.param(torch.rand(32, 10, 28, 28), torch.randint(10, (32, 28, 28)), None, 10),
|
|
pytest.param(torch.rand(32, 28, 28), torch.randint(10, (32, 28, 28)), None, 10),
|
|
])
|
|
def test_get_num_classes(pred, target, num_classes, expected_num_classes):
|
|
assert get_num_classes(pred, target, num_classes) == expected_num_classes
|
|
|
|
|
|
@pytest.mark.parametrize(['sample_weight', 'pos_label', "exp_shape"], [
|
|
pytest.param(1, 1., 42),
|
|
pytest.param(None, 1., 42),
|
|
])
|
|
def test_binary_clf_curve(sample_weight, pos_label, exp_shape):
|
|
# TODO: move back the pred and target to test func arguments
|
|
# if you fix the array inside the function, you'd also have fix the shape,
|
|
# because when the array changes, you also have to fix the shape
|
|
seed_everything(0)
|
|
pred = torch.randint(low=51, high=99, size=(100, ), dtype=torch.float) / 100
|
|
target = torch.tensor([0, 1] * 50, dtype=torch.int)
|
|
if sample_weight is not None:
|
|
sample_weight = torch.ones_like(pred) * sample_weight
|
|
|
|
fps, tps, thresh = _binary_clf_curve(preds=pred, target=target, sample_weights=sample_weight, pos_label=pos_label)
|
|
|
|
assert isinstance(tps, torch.Tensor)
|
|
assert isinstance(fps, torch.Tensor)
|
|
assert isinstance(thresh, torch.Tensor)
|
|
assert tps.shape == (exp_shape, )
|
|
assert fps.shape == (exp_shape, )
|
|
assert thresh.shape == (exp_shape, )
|
|
|
|
|
|
@pytest.mark.parametrize(['pred', 'target', 'expected'], [
|
|
pytest.param([[0, 0], [1, 1]], [[0, 0], [1, 1]], 1.),
|
|
pytest.param([[1, 1], [0, 0]], [[0, 0], [1, 1]], 0.),
|
|
pytest.param([[1, 1], [1, 1]], [[1, 1], [0, 0]], 2 / 3),
|
|
pytest.param([[1, 1], [0, 0]], [[1, 1], [0, 0]], 1.),
|
|
])
|
|
def test_dice_score(pred, target, expected):
|
|
score = dice_score(torch.tensor(pred), torch.tensor(target))
|
|
assert score == expected
|