File size: 2,038 Bytes
c81bc7a | 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 | """Unit tests for model architecture and uncertainty estimation."""
import pytest
import torch
from models.classifier import ChestAIClassifier
from models.uncertainty import mc_predict, enable_dropout
@pytest.fixture(scope="module")
def model():
"""Lightweight model for testing — skip backbone download in CI."""
import unittest.mock as mock
# Mock the BioMedCLIP download for fast CI runs
with mock.patch("models.backbone.create_model_from_pretrained") as m:
mock_visual = mock.MagicMock()
mock_visual.output_dim = 512
mock_visual.return_value = torch.zeros(2, 512)
m.return_value = (mock.MagicMock(visual=mock_visual), None)
clf = ChestAIClassifier(num_classes=14, freeze_backbone=False)
return clf
def test_output_shape(model):
x = torch.randn(2, 3, 224, 224)
with torch.no_grad():
out = model(x)
assert out.shape == (2, 14), f"Expected (2,14), got {out.shape}"
def test_output_range_after_sigmoid(model):
x = torch.randn(4, 3, 224, 224)
with torch.no_grad():
logits = model(x)
probs = torch.sigmoid(logits)
assert probs.min() >= 0.0 and probs.max() <= 1.0
def test_mc_dropout_produces_variance(model):
x = torch.randn(1, 3, 224, 224)
result = mc_predict(model, x, n_samples=5)
assert result["std"].shape == (1, 14)
# Uncertainty should be positive (dropout creates variance)
assert result["std"].mean().item() > 0
def test_uncertainty_keys(model):
x = torch.randn(1, 3, 224, 224)
result = mc_predict(model, x, n_samples=5)
assert set(result.keys()) == {"mean", "std", "entropy", "samples"}
assert result["samples"].shape[0] == 5
def test_enable_dropout_sets_train_mode():
import torch.nn as nn
model_small = nn.Sequential(nn.Linear(10, 10), nn.Dropout(0.3), nn.Linear(10, 5))
model_small.eval()
enable_dropout(model_small)
dropout_layers = [m for m in model_small.modules() if isinstance(m, nn.Dropout)]
assert all(d.training for d in dropout_layers)
|