|
|
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 |
|
|
|
