diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..b01b8c3 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,75 @@ +import sys + +sys.path.insert(0, sys.path[0] + "/../") + +import os +import pytest +import torch +import torch.nn as nn +import numpy as np + +from abl.models.nn import LeNet5, SymbolNet +from abl.models.basic_model import BasicModel + + +class TestBasicModel(object): + @pytest.mark.parametrize("num_classes", [4, 10]) + @pytest.mark.parametrize("image_size", [(28, 28, 1), (45, 45, 1)]) + @pytest.mark.parametrize("cls", [LeNet5, SymbolNet]) + @pytest.mark.parametrize("criterion", [nn.CrossEntropyLoss]) + @pytest.mark.parametrize("optimizer", [torch.optim.RMSprop]) + @pytest.mark.parametrize("device", [torch.device("cpu"), torch.device("cuda:0")]) + def test_models(self, num_classes, image_size, cls, criterion, optimizer, device): + cls = cls(num_classes=num_classes, image_size=image_size) + criterion = criterion() + optimizer = optimizer(cls.parameters(), lr=0.001) + + self.num_classes = num_classes + self.image_size = image_size + self.model = BasicModel(cls, criterion, optimizer, device) + + self.data_X = [ + np.random.rand(image_size[2], image_size[0], image_size[1]).astype( + np.float32 + ) + for i in range(5) + ] + self.data_y = np.random.randint(0, num_classes, (5,)) + + self._test_fit() + self._test_predict() + self._test_predict_proba() + self._test_val() + self._test_save() + self._test_load() + + def _test_fit(self): + self.model.fit(X=self.data_X, y=self.data_y) + + def _test_predict(self): + predict_result = self.model.predict(X=self.data_X) + assert predict_result.dtype == int + assert predict_result.shape == (5,) + assert (0 <= predict_result).all() and (predict_result < self.num_classes).all() + + def _test_predict_proba(self): + predict_result = self.model.predict_proba(X=self.data_X) + assert predict_result.dtype == np.float32 + assert predict_result.shape == (5, self.num_classes) + assert (0 <= predict_result).all() and (predict_result <= 1).all() + + def _test_val(self): + accuracy = self.model.val(X=self.data_X, y=self.data_y) + assert type(accuracy) == float + assert 0 <= accuracy <= 1 + + def _test_save(self): + self.model.save(1, "results/test_models") + assert os.path.exists("results/test_models/1_net.pth") + assert os.path.exists("results/test_models/1_opt.pth") + os.remove("results/test_models/1_net.pth") + os.remove("results/test_models/1_opt.pth") + + def _test_load(self): + self.model.save(1, "results/test_models") + self.model.load(1, "results/test_models")