| import pytest |
| import torch.nn as nn |
|
|
| from kornia.metrics import AverageMeter |
| from kornia.x import EarlyStopping, ModelCheckpoint |
| from kornia.x.utils import TrainerState |
|
|
|
|
| @pytest.fixture |
| def model(): |
| return nn.Conv2d(3, 10, kernel_size=1) |
|
|
|
|
| def test_callback_modelcheckpoint(tmp_path, model): |
| cb = ModelCheckpoint(tmp_path, 'test_monitor') |
| assert cb is not None |
|
|
| metric = {'test_monitor': AverageMeter()} |
| metric['test_monitor'].avg = 1.0 |
|
|
| cb(model, epoch=0, valid_metric=metric) |
| assert cb.best_metric == 1.0 |
| assert (tmp_path / "model_0.pt").is_file() |
|
|
|
|
| def test_callback_earlystopping(model): |
| cb = EarlyStopping('test_monitor', patience=2) |
| assert cb is not None |
| assert cb.counter == 0 |
|
|
| metric = {'test_monitor': AverageMeter()} |
| metric['test_monitor'].avg = 1 |
|
|
| state = cb(model, epoch=0, valid_metric=metric) |
| assert state == TrainerState.TRAINING |
| assert cb.best_score == -1 |
| assert cb.counter == 0 |
|
|
| metric['test_monitor'].avg = 2 |
| state = cb(model, epoch=0, valid_metric=metric) |
| assert state == TrainerState.TRAINING |
| assert cb.best_score == -1 |
| assert cb.counter == 1 |
|
|
| state = cb(model, epoch=0, valid_metric=metric) |
| assert state == TrainerState.TERMINATE |
|
|
