increased tests

Author: simon-donike

README.md

@@ -134,6 +134,6 @@ If you use this work, please cite:
 
 Developed within the **ESA Φ-lab / OpenSR** initiative. Simon Donike is the main contributor and maintainer of the repository. Cesar Aybar and Julio Contreras contributed the datasets as well as implementation, documentation and publishing support. Prof. Luis Gómez-Chova contributed the remote sensing-specific perspective and signal processing advice.  
 > The development history of this code began in 2020 with the implementation of an SR-GAN for a MSc thesis project. Since then, over several iterations, the codebase has been expanded and many training tweaks implemented, based on the experiences made training SR-GANs for the OpenSR project. The fundamental training outline, training tweaks, normalizations, and inference procedures are built upon that experience.  
-The added complexity that came with (a) the implementation of many different models and blocks, (b) more data sources, (c) according normalizations, and (d) complex testing and documentation structures, was handled to varying degrees with the help of *Codex*. Specifically, the docs, the automated testing workflows, and the normalizer class are in part AI generated. This code and its functionalities have been verified and tested to the best of my ability.
+The added complexity that came with (a) the implementation of many different models and blocks, (b) more data sources, (c) according normalizations, and (d) testing and documentation structures, was handled to varying degrees with the help of *Codex*. Specifically, the docs, the automated testing workflows, and the normalizer class are in (large) parts AI generated. This code and its functionalities have been verified and tested to the best of my ability.
 
 ---

docs/getting-started.md

@@ -2,6 +2,16 @@
 
 This guide walks through installing dependencies, configuring datasets, and launching your first ESA OpenSR experiment. The stack supports Python 3.10-3.12, PyTorch Lightning, and Weights & Biases for experiment tracking.
 
+## Try it in Colab first
+
+For the fastest start, open the interactive notebook in Google Colab and run through the introduction without setting up a local environment.
+
+<p align="center">
+  <a href="https://colab.research.google.com/drive/16W0FWr6py1J8P4po7JbNDMaepHUM97yL?usp=sharing">
+    <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab">
+  </a>
+</p>
+
 > 💡 **Only need inference?** Install the published package instead: `python -m pip install opensr-srgan`. It exposes `load_from_config` and `load_inference_model` so you can instantiate models without cloning the repository. Continue with the rest of this guide when you want to train, fine-tune, or otherwise modify the codebase.
 
 ## 1. Install the environment

tests/test_model_blocks/test_ema.py

@@ -83,3 +83,116 @@ def test_invalid_decay_raises():
     model = TinyNet()
     with pytest.raises(ValueError):
         ExponentialMovingAverage(model, decay=1.5)
+
+
+def test_update_with_num_updates_uses_warmup_decay_and_updates_buffer():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.9, use_num_updates=True)
+    old_shadow = ema.shadow_params["lin.weight"].clone()
+
+    with torch.no_grad():
+        model.lin.weight.add_(1.0)
+        model.scale.fill_(3.0)
+
+    ema.update(model)
+
+    expected_decay = min(0.9, (1 + 1) / (10 + 1))
+    expected = old_shadow.lerp(model.lin.weight.detach(), 1.0 - expected_decay)
+    assert ema.num_updates == 1
+    assert torch.allclose(ema.shadow_params["lin.weight"], expected)
+    assert torch.allclose(ema.shadow_buffers["scale"], torch.tensor([3.0]))
+
+
+def test_register_and_update_skip_frozen_parameters():
+    model = TinyNet()
+    model.lin.bias.requires_grad_(False)
+
+    ema = ExponentialMovingAverage(model, decay=0.5, use_num_updates=False)
+    assert set(ema.shadow_params) == {"lin.weight"}
+
+    with torch.no_grad():
+        model.lin.bias.add_(10.0)
+    ema.update(model)
+
+    assert "lin.bias" not in ema.shadow_params
+
+
+def test_update_registers_new_trainable_parameters_and_buffers():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.5, use_num_updates=False)
+
+    model.extra = nn.Parameter(torch.full((1,), 4.0))
+    model.register_buffer("offset", torch.full((1,), 2.0))
+
+    ema.update(model)
+
+    assert torch.allclose(ema.shadow_params["extra"], torch.tensor([4.0]))
+    assert torch.allclose(ema.shadow_buffers["offset"], torch.tensor([2.0]))
+
+
+def test_apply_to_swaps_buffers_and_rejects_reapply_before_restore():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.9)
+    original_scale = model.scale.clone()
+
+    ema.shadow_buffers["scale"].fill_(5.0)
+    ema.apply_to(model)
+
+    assert torch.allclose(model.scale, torch.tensor([5.0]))
+    with pytest.raises(RuntimeError, match="already applied"):
+        ema.apply_to(model)
+
+    ema.restore(model)
+    assert torch.allclose(model.scale, original_scale)
+    assert not ema.collected_params
+    assert not ema.collected_buffers
+
+
+def test_restore_without_apply_is_noop():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.9)
+    original_weight = model.lin.weight.clone()
+    original_scale = model.scale.clone()
+
+    ema.restore(model)
+
+    assert torch.allclose(model.lin.weight, original_weight)
+    assert torch.allclose(model.scale, original_scale)
+
+
+def test_average_parameters_restores_after_exception():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.9)
+    ema.shadow_params["lin.weight"].add_(2.0)
+    original_weight = model.lin.weight.clone()
+
+    with pytest.raises(RuntimeError, match="boom"):
+        with ema.average_parameters(model):
+            assert torch.allclose(model.lin.weight, ema.shadow_params["lin.weight"])
+            raise RuntimeError("boom")
+
+    assert torch.allclose(model.lin.weight, original_weight)
+    assert not ema.collected_params
+
+
+def test_load_state_dict_restores_device_and_clears_collected_caches():
+    model = TinyNet()
+    ema = ExponentialMovingAverage(model, decay=0.9)
+    state = ema.state_dict()
+    state["decay"] = 0.25
+    state["num_updates"] = None
+    state["device"] = "cpu"
+
+    ema2 = ExponentialMovingAverage(model, decay=0.1)
+    ema2.apply_to(model)
+    assert ema2.collected_params
+
+    ema2.load_state_dict(state)
+
+    assert ema2.decay == pytest.approx(0.25)
+    assert ema2.num_updates is None
+    assert ema2.device == torch.device("cpu")
+    assert not ema2.collected_params
+    assert not ema2.collected_buffers
+    for tensor in ema2.shadow_params.values():
+        assert tensor.device.type == "cpu"

tests/test_models/test_generators.py

@@ -167,6 +167,74 @@ def test_esrgan_generator_warns_about_srresnet_specific_options(capsys):
     )
 
 
+@pytest.mark.parametrize("block_type", ["res", "rcab", "rrdb", "lka"])
+@pytest.mark.parametrize(
+    "scale, expected_shape",
+    [
+        (2, (2, 3, 10, 12)),
+        (4, (2, 3, 20, 24)),
+        (8, (2, 3, 40, 48)),
+    ],
+)
+def test_flexible_generator_forward_scales_output_for_all_blocks(
+    block_type, scale, expected_shape
+):
+    generator = FlexibleGenerator(
+        in_channels=3,
+        n_channels=16,
+        n_blocks=1,
+        small_kernel=3,
+        large_kernel=3,
+        scale=scale,
+        block_type=block_type,
+    )
+    lr = torch.randn(2, 3, 5, 6)
+
+    sr = generator(lr)
+
+    assert sr.shape == expected_shape
+
+
+def test_flexible_generator_rejects_invalid_configuration():
+    with pytest.raises(ValueError, match="scale must be one of"):
+        FlexibleGenerator(scale=1)
+
+    with pytest.raises(ValueError, match="block_type must be one of"):
+        FlexibleGenerator(block_type="unknown")
+
+
+@pytest.mark.parametrize(
+    "scale, expected_shape",
+    [
+        (1, (2, 2, 5, 6)),
+        (2, (2, 2, 10, 12)),
+        (4, (2, 2, 20, 24)),
+    ],
+)
+def test_esrgan_generator_forward_scales_output(scale, expected_shape):
+    generator = ESRGANGenerator(
+        in_channels=2,
+        out_channels=2,
+        n_features=8,
+        n_blocks=1,
+        growth_channels=4,
+        scale=scale,
+    )
+    lr = torch.randn(2, 2, 5, 6)
+
+    sr = generator(lr)
+
+    assert sr.shape == expected_shape
+
+
+def test_esrgan_generator_rejects_invalid_configuration():
+    with pytest.raises(ValueError, match="power-of-two scales"):
+        ESRGANGenerator(scale=3)
+
+    with pytest.raises(ValueError, match="at least one RRDB block"):
+        ESRGANGenerator(n_blocks=0)
+
+
 def test_stochastic_generator_forward_noise_paths():
     generator = StochasticGenerator(
         in_channels=2,