File size: 3,351 Bytes
f8c73f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from __future__ import annotations
import pytest
from adam.executor import ToolContext, ToolExecutionError
from adam.progressive_training import parse_stages, stage_batch_settings, suggested_stages
from adam.registry import ToolSpec
from adam.tools import ddpm_adapter
def _context(root):
import threading
running = threading.Event()
running.set()
return ToolContext(
root=root, job_id="PROGRESSIVE", tool=ToolSpec("ddpm_trainer", "DDPM", "", "Training", "train_ddpm"),
cancel_event=threading.Event(), run_event=running,
progress_callback=lambda *_args, **_kwargs: None, log_callback=lambda _message: None,
)
def test_progressive_stages_require_an_increasing_complete_budget() -> None:
stages = parse_stages(
[{"resolution": 64, "epochs": 30}, {"resolution": 128, "epochs": 10}],
trainer="ddpm", total_epochs=40,
)
assert [(stage.resolution, stage.epochs) for stage in stages] == [(64, 30), (128, 10)]
with pytest.raises(ToolExecutionError, match="total"):
parse_stages(
[{"resolution": 64, "epochs": 30}, {"resolution": 128, "epochs": 10}],
trainer="ddpm", total_epochs=41,
)
with pytest.raises(ToolExecutionError, match="increase"):
parse_stages(
[{"resolution": 128, "epochs": 10}, {"resolution": 64, "epochs": 10}],
trainer="ddpm", total_epochs=20,
)
def test_auto_stage_batch_never_keeps_an_unsafe_final_batch() -> None:
batch, accumulation = stage_batch_settings(
trainer="flow", stage_resolution=512, final_resolution=512,
final_batch_size=12, base_accumulation=1, auto_batch=True,
)
assert (batch, accumulation) == (1, 12)
low_batch, low_accumulation = stage_batch_settings(
trainer="flow", stage_resolution=64, final_resolution=512,
final_batch_size=1, base_accumulation=12, auto_batch=True,
)
assert low_batch == 12
assert low_accumulation == 1
def test_suggested_schedule_uses_the_entire_epoch_budget() -> None:
stages = suggested_stages(512, 485)
assert stages[0].resolution == 64
assert stages[-1].resolution == 512
assert sum(stage.epochs for stage in stages) == 485
def test_ddpm_curriculum_hands_weights_forward_and_publishes_only_final_model(tmp_path, monkeypatch) -> None:
calls: list[tuple[int, str, str]] = []
def fake_stage(_context, _dataset, _name, stage_epochs, stage_output, resume_from, *args):
path = __import__("pathlib").Path(stage_output)
path.mkdir(parents=True)
(path / "model_index.json").write_text("{}", encoding="utf-8")
calls.append((stage_epochs, str(path), str(resume_from)))
return {}
monkeypatch.setattr(ddpm_adapter, "_train_ddpm_stage", fake_stage)
final_output = tmp_path / "output" / "Example"
result = ddpm_adapter.train_ddpm(
_context(tmp_path), dataset_dir=str(tmp_path), model_name="Example", epochs=5,
output_dir=str(final_output), resolution=128, batch_size=1,
progressive_stages=[{"resolution": 64, "epochs": 3}, {"resolution": 128, "epochs": 2}],
)
assert len(calls) == 2
assert calls[0][2] == ""
assert calls[1][2] == calls[0][1]
assert (final_output / "model_index.json").is_file()
assert result["output_folder"] == str(final_output)
|