Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/callbacks/test_weight_averaging.py (Python)

Code before:

    # At and after starting step, but not on update frequency
    assert callback.should_update(step_idx=10)  # First update
    assert not callback.should_update(step_idx=11)
    assert not callback.should_update(step_idx=14)
    assert callback.should_update(step_idx=15)  # Second update

    # Test with epoch-based updates
    callback = EMAWeightAveraging(update_starting_at_epoch=2)

    assert not callback.should_update(epoch_idx=0)
    assert not callback.should_update(epoch_idx=1)
    assert callback.should_update(epoch_idx=2)
    assert callback.should_update(epoch_idx=3)


def test_ema_weight_averaging_checkpoint_save_load(tmp_path):
    """Test that EMAWeightAveraging correctly saves and loads checkpoints."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)

    # Train and create checkpoint
    _train(model, dataset, tmp_path, callback, will_crash=True)

    # Resume from checkpoint
    model2 = TestModel()
    callback2 = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)
    checkpoint_path = str(tmp_path / "lightning_logs" / "version_0" / "checkpoints" / "*.ckpt")

    _train(model2, dataset, tmp_path, callback2, checkpoint_path=checkpoint_path)

    assert callback2._average_model is not None


@pytest.mark.parametrize("decay", [0.9, 0.99, 0.999, 0.9999])
def test_ema_weight_averaging_decay_values(tmp_path, decay):
    """Test EMAWeightAveraging with different decay values."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=decay, update_every_n_steps=1)
    _train(model, dataset, tmp_path, callback)

    assert callback._average_model is not None


Review comment (bug):
Would need to set `crash_on_epoch` for the `TestModel` to raise exception when epoch reaches. Also, this is probably where the tests failed.

```suggestion
    model = TestModel()
    model.crash_on_epoch = 2
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)

    # Train and create checkpoint
    _train(model, dataset, tmp_path, callback, will_crash=True)
```

Code after:

    # At and after starting step, but not on update frequency
    assert callback.should_update(step_idx=10)  # First update
    assert not callback.should_update(step_idx=11)
    assert not callback.should_update(step_idx=14)
    assert callback.should_update(step_idx=15)  # Second update

    # Test with epoch-based updates
    callback = EMAWeightAveraging(update_starting_at_epoch=2)

    assert not callback.should_update(epoch_idx=0)
    assert not callback.should_update(epoch_idx=1)
    assert callback.should_update(epoch_idx=2)
    assert callback.should_update(epoch_idx=3)


def test_ema_weight_averaging_checkpoint_save_load(tmp_path):
    """Test that EMAWeightAveraging correctly saves and loads checkpoints."""
    model = TestModel()
    model.crash_on_epoch = 2
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)

    # Train and create checkpoint
    _train(model, dataset, tmp_path, callback, will_crash=True)

    # Resume from checkpoint
    model2 = TestModel()
    callback2 = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)
    import glob  # should be at the top

    _train(
        model2,
        dataset,
        tmp_path,
        callback2,
        checkpoint_path=glob.glob((tmp_path / "checkpoints" / "*.ckpt").as_posix())[0],
    )

    assert callback2._average_model is not None


@pytest.mark.parametrize("decay", [0.9, 0.99, 0.999, 0.9999])
def test_ema_weight_averaging_decay_values(tmp_path, decay):
    """Test EMAWeightAveraging with different decay values."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=decay, update_every_n_steps=1)
    _train(model, dataset, tmp_path, callback)