update: pytest markers

Author: kozistr

tests/test_create_optimizer.py

@@ -1,7 +1,10 @@
+import pytest
+
 from pytorch_optimizer import create_optimizer
 from tests.utils import LogisticRegression
 
 
+@pytest.mark.utils
 def test_create_optimizer():
     model = LogisticRegression()
 

tests/test_gradients.py

@@ -7,6 +7,7 @@
 from tests.utils import build_environment, simple_parameter, simple_sparse_parameter
 
 
+@pytest.mark.gradient
 @pytest.mark.parametrize('optimizer_name', [*VALID_OPTIMIZER_NAMES, 'lookahead'])
 def test_no_gradients(optimizer_name):
     p1 = simple_parameter(require_grad=True)
@@ -32,6 +33,7 @@ def test_no_gradients(optimizer_name):
     optimizer.step(lambda: 0.1)  # for AliG optimizer
 
 
+@pytest.mark.gradient
 @pytest.mark.parametrize('no_sparse_optimizer', NO_SPARSE_OPTIMIZERS)
 def test_sparse_not_supported(no_sparse_optimizer):
     param = simple_sparse_parameter()
@@ -45,6 +47,7 @@ def test_sparse_not_supported(no_sparse_optimizer):
         optimizer.step(lambda: 0.1)
 
 
+@pytest.mark.gradient
 @pytest.mark.parametrize('sparse_optimizer', SPARSE_OPTIMIZERS)
 def test_sparse_supported(sparse_optimizer):
     param = simple_sparse_parameter()
@@ -78,6 +81,7 @@ def test_sparse_supported(sparse_optimizer):
         optimizer.step()
 
 
+@pytest.mark.gradient
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_bf16_gradient(optimizer_name):
     # torch.eye does not support bf16
@@ -92,6 +96,7 @@ def test_bf16_gradient(optimizer_name):
     optimizer.step(lambda: 0.1)
 
 
+@pytest.mark.gradient
 def test_sam_no_gradient():
     (x_data, y_data), model, loss_fn = build_environment()
     model.fc1.weight.requires_grad = False

tests/test_load_lr_schedulers.py

@@ -4,16 +4,19 @@
 from tests.constants import INVALID_LR_SCHEDULER_NAMES, VALID_LR_SCHEDULER_NAMES
 
 
+@pytest.mark.utils
 @pytest.mark.parametrize('valid_lr_scheduler_names', VALID_LR_SCHEDULER_NAMES)
 def test_load_optimizers_valid(valid_lr_scheduler_names):
     load_lr_scheduler(valid_lr_scheduler_names)
 
 
+@pytest.mark.utils
 @pytest.mark.parametrize('invalid_lr_scheduler_names', INVALID_LR_SCHEDULER_NAMES)
 def test_load_optimizers_invalid(invalid_lr_scheduler_names):
     with pytest.raises(NotImplementedError):
         load_lr_scheduler(invalid_lr_scheduler_names)
 
 
+@pytest.mark.utils
 def test_get_supported_lr_schedulers():
     assert len(get_supported_lr_schedulers()) == 10

tests/test_load_optimizers.py

@@ -4,16 +4,19 @@
 from tests.constants import INVALID_OPTIMIZER_NAMES, VALID_OPTIMIZER_NAMES
 
 
+@pytest.mark.utils
 @pytest.mark.parametrize('valid_optimizer_names', VALID_OPTIMIZER_NAMES)
 def test_load_optimizers_valid(valid_optimizer_names):
     load_optimizer(valid_optimizer_names)
 
 
+@pytest.mark.utils
 @pytest.mark.parametrize('invalid_optimizer_names', INVALID_OPTIMIZER_NAMES)
 def test_load_optimizers_invalid(invalid_optimizer_names):
     with pytest.raises(NotImplementedError):
         load_optimizer(invalid_optimizer_names)
 
 
+@pytest.mark.utils
 def test_get_supported_optimizers():
     assert len(get_supported_optimizers()) == 29

tests/test_lr_scheduler_parameters.py

@@ -8,6 +8,7 @@
 from tests.utils import Example
 
 
+@pytest.mark.lr_scheduler
 def test_cosine_annealing_warmup_restarts_params():
     optimizer = AdamP(Example().parameters())
 
@@ -36,6 +37,7 @@ def test_cosine_annealing_warmup_restarts_params():
         lr_scheduler.step(epoch=None)
 
 
+@pytest.mark.lr_scheduler
 def test_linear_warmup_lr_scheduler_params():
     optimizer = AdamP(Example().parameters())
 
@@ -60,6 +62,7 @@ def test_linear_warmup_lr_scheduler_params():
         PolyScheduler(optimizer=optimizer, t_max=1, max_lr=1, min_lr=1, init_lr=1, warmup_steps=-1)
 
 
+@pytest.mark.lr_scheduler
 def test_chebyshev_params():
     with pytest.raises(IndexError):
         get_chebyshev_schedule(2)

tests/test_lr_schedulers.py

@@ -114,6 +114,7 @@
 PROPORTION_LEARNING_RATES = [(1e-1, 1e-1, 2.0), (1e-1, 1e-3, 1.090909)]
 
 
+@pytest.mark.lr_scheduler
 @pytest.mark.parametrize('cosine_annealing_warmup_restart_param', CAWR_RECIPES)
 def test_cosine_annealing_warmup_restarts(cosine_annealing_warmup_restart_param):
     model = Example()
@@ -150,11 +151,13 @@ def test_cosine_annealing_warmup_restarts(cosine_annealing_warmup_restart_param)
         np.testing.assert_almost_equal(expected_lrs[epoch], lr)
 
 
+@pytest.mark.lr_scheduler
 def test_get_chebyshev_scheduler():
     np.testing.assert_almost_equal(get_chebyshev_schedule(3), 1.81818182, decimal=6)
     np.testing.assert_array_equal(chebyshev_perm(5), np.asarray([0, 7, 3, 4, 1, 6, 2, 5]))
 
 
+@pytest.mark.lr_scheduler
 def test_linear_warmup_linear_scheduler():
     optimizer = AdamP(Example().parameters())
     lr_scheduler = LinearScheduler(optimizer, t_max=10, max_lr=1e-2, min_lr=1e-4, init_lr=1e-3, warmup_steps=5)
@@ -164,6 +167,7 @@ def test_linear_warmup_linear_scheduler():
         np.testing.assert_almost_equal(expected_lr, lr_scheduler.get_lr())
 
 
+@pytest.mark.lr_scheduler
 def test_linear_warmup_cosine_scheduler():
     optimizer = AdamP(Example().parameters())
     lr_scheduler = CosineScheduler(optimizer, t_max=10, max_lr=1e-2, min_lr=1e-4, init_lr=1e-3, warmup_steps=5)
@@ -173,6 +177,7 @@ def test_linear_warmup_cosine_scheduler():
         np.testing.assert_almost_equal(expected_lr, lr_scheduler.get_lr(), 5)
 
 
+@pytest.mark.lr_scheduler
 def test_linear_warmup_poly_scheduler():
     optimizer = AdamP(Example().parameters())
     lr_scheduler = PolyScheduler(optimizer=optimizer, t_max=10, max_lr=1e-2, min_lr=1e-4, init_lr=1e-3, warmup_steps=5)
@@ -182,6 +187,7 @@ def test_linear_warmup_poly_scheduler():
         np.testing.assert_almost_equal(expected_lr, lr_scheduler.get_lr(), 6)
 
 
+@pytest.mark.lr_scheduler
 @pytest.mark.parametrize('proportion_learning_rate', PROPORTION_LEARNING_RATES)
 def test_proportion_scheduler(proportion_learning_rate: Tuple[float, float, float]):
     base_optimizer = AdamP(Example().parameters())
@@ -201,6 +207,7 @@ def test_proportion_scheduler(proportion_learning_rate: Tuple[float, float, floa
         np.testing.assert_almost_equal(proportion_learning_rate[2], rho_scheduler.get_lr(), 6)
 
 
+@pytest.mark.lr_scheduler
 def test_proportion_no_last_lr_scheduler():
     base_optimizer = AdamP(Example().parameters())
     lr_scheduler = CosineAnnealingWarmupRestarts(
@@ -222,6 +229,7 @@ def test_proportion_no_last_lr_scheduler():
         np.testing.assert_almost_equal(2.0, rho_scheduler.get_lr(), 6)
 
 
+@pytest.mark.lr_scheduler
 def test_deberta_v3_large_lr_scheduler():
     try:
         from transformers import AutoConfig, AutoModel

tests/test_optimizer_parameters.py

@@ -8,6 +8,7 @@
 from tests.utils import Example, simple_parameter
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_learning_rate(optimizer_name):
     if optimizer_name in ('alig',):
@@ -23,6 +24,7 @@ def test_learning_rate(optimizer_name):
         optimizer(None, **config)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_epsilon(optimizer_name):
     if optimizer_name in ('nero', 'shampoo', 'scalableshampoo', 'dadaptsgd', 'adafactor', 'lion'):
@@ -40,6 +42,7 @@ def test_epsilon(optimizer_name):
     assert str(error_info.value) == '[-] epsilon -1e-06 must be non-negative'
 
 
+@pytest.mark.cfg
 def test_shampoo_epsilon():
     shampoo = load_optimizer('Shampoo')
     scalable_shampoo = load_optimizer('ScalableShampoo')
@@ -54,6 +57,7 @@ def test_shampoo_epsilon():
         shampoo(None, matrix_eps=-1e-6)
 
 
+@pytest.mark.cfg
 def test_adafactor_epsilon():
     adafactor = load_optimizer('adafactor')
 
@@ -64,6 +68,7 @@ def test_adafactor_epsilon():
         adafactor(None, eps2=-1e-6)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_weight_decay(optimizer_name):
     if optimizer_name in ('nero', 'alig'):
@@ -81,6 +86,7 @@ def test_weight_decay(optimizer_name):
     assert str(error_info.value) == '[-] weight_decay -0.001 must be non-negative'
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['apollo'])
 def test_weight_decay_type(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
@@ -89,6 +95,7 @@ def test_weight_decay_type(optimizer_name):
         optimizer(None, weight_decay_type='dummy')
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['apollo'])
 def test_rebound(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
@@ -97,48 +104,55 @@ def test_rebound(optimizer_name):
         optimizer(None, rebound='dummy')
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['adamp', 'sgdp'])
 def test_wd_ratio(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, wd_ratio=-1e-3)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['lars'])
 def test_trust_coefficient(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, trust_coefficient=-1e-3)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['madgrad', 'lars'])
 def test_momentum(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, momentum=-1e-3)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['ranger'])
 def test_lookahead_k(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, k=-1)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['ranger21'])
 def test_beta0(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, num_iterations=200, beta0=-0.1)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['nero', 'apollo'])
 def test_beta(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, beta=-0.1)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', BETA_OPTIMIZER_NAMES)
 def test_betas(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
@@ -160,6 +174,7 @@ def test_betas(optimizer_name):
             optimizer(None, betas=(0.1, 0.1, -0.1))
 
 
+@pytest.mark.cfg
 def test_reduction():
     parameters = Example().parameters()
     optimizer = load_optimizer('adamp')(parameters)
@@ -168,6 +183,7 @@ def test_reduction():
         PCGrad(optimizer, reduction='wrong')
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['scalableshampoo', 'shampoo'])
 def test_update_frequency(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
@@ -183,18 +199,21 @@ def test_update_frequency(optimizer_name):
         optimizer(None, preconditioning_compute_steps=-1)
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer_name', ['adan', 'lamb'])
 def test_norm(optimizer_name):
     optimizer = load_optimizer(optimizer_name)
     with pytest.raises(ValueError):
         optimizer(None, max_grad_norm=-0.1)
 
 
+@pytest.mark.cfg
 def test_sam_parameters():
     with pytest.raises(ValueError, match=''):
         SAM(None, load_optimizer('adamp'), rho=-0.1)
 
 
+@pytest.mark.cfg
 def test_lookahead_parameters():
     param = simple_parameter()
     optimizer = load_optimizer('adamp')([param])
@@ -216,6 +235,7 @@ def test_lookahead_parameters():
         Lookahead(optimizer, pullback_momentum='invalid')
 
 
+@pytest.mark.cfg
 def test_sam_methods():
     param = simple_parameter()
 
@@ -224,6 +244,7 @@ def test_sam_methods():
     optimizer.load_state_dict(optimizer.state_dict())
 
 
+@pytest.mark.cfg
 def test_safe_fp16_methods():
     param = simple_parameter()
 
@@ -244,12 +265,14 @@ def test_safe_fp16_methods():
     assert optimizer.loss_scale == 2.0 ** (15 - 1)
 
 
+@pytest.mark.cfg
 def test_ranger21_warm_methods():
     assert Ranger21.build_warm_up_iterations(1000, 0.999) == 220
     assert Ranger21.build_warm_up_iterations(4500, 0.999) == 2000
     assert Ranger21.build_warm_down_iterations(1000) == 280
 
 
+@pytest.mark.cfg
 @pytest.mark.parametrize('optimizer', ['ranger21', 'adai'])
 def test_size_of_parameter(optimizer):
     param = simple_parameter(require_grad=False)
@@ -259,6 +282,7 @@ def test_size_of_parameter(optimizer):
         load_optimizer(optimizer)([param], 1).step()
 
 
+@pytest.mark.cfg
 def test_ranger21_closure():
     model: nn.Module = Example()
     optimizer = load_optimizer('ranger21')(model.parameters(), num_iterations=100, betas=(0.9, 1e-9))
@@ -273,6 +297,7 @@ def closure():
     optimizer.step(closure)
 
 
+@pytest.mark.cfg
 def test_adafactor_reset():
     param = torch.zeros(1).requires_grad_(True)
     param.grad = torch.zeros(1)
@@ -281,6 +306,7 @@ def test_adafactor_reset():
     optimizer.reset()
 
 
+@pytest.mark.cfg
 def test_adafactor_get_lr():
     model: nn.Module = Example()