File size: 4,702 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Resolution-curriculum helpers shared by the DDPM and Flow adapters.
The adapters deliberately use conservative batch caps instead of trying an
out-of-memory probe in a real training job. A user can still turn the policy
off and enter every batch setting manually.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from adam.executor import ToolExecutionError
@dataclass(frozen=True, slots=True)
class ResolutionStage:
resolution: int
epochs: int
_BATCH_CAPS = {
"ddpm": {64: 16, 128: 12, 256: 4, 384: 2, 512: 1},
"flow": {64: 12, 128: 8, 256: 4, 384: 2, 512: 1},
}
def parse_stages(value: Any, *, trainer: str, total_epochs: int) -> list[ResolutionStage]:
"""Validate a JSON-friendly progressive-resolution schedule.
Stages are intentionally a list of small dictionaries so plans remain easy
to inspect and edit in saved job JSON.
"""
if not isinstance(value, list) or len(value) < 2:
raise ToolExecutionError("Progressive training needs at least two resolution stages.")
multiple = 16 if trainer == "flow" else 8
stages: list[ResolutionStage] = []
previous = 0
for raw in value:
if not isinstance(raw, dict):
raise ToolExecutionError("Each progressive stage must include resolution and epochs.")
try:
resolution = int(raw.get("resolution", 0))
epochs = int(raw.get("epochs", 0))
except (TypeError, ValueError) as exc:
raise ToolExecutionError("Progressive stage resolution and epochs must be whole numbers.") from exc
if not 64 <= resolution <= 512 or resolution % multiple:
raise ToolExecutionError(
f"{trainer.upper()} progressive resolutions must be 64–512 and divisible by {multiple}."
)
if resolution <= previous:
raise ToolExecutionError("Progressive stages must increase from lower to higher resolution.")
if epochs < 1:
raise ToolExecutionError("Each progressive stage needs at least one epoch.")
stages.append(ResolutionStage(resolution, epochs))
previous = resolution
if sum(stage.epochs for stage in stages) != int(total_epochs):
raise ToolExecutionError(
f"Progressive stage epochs total {sum(stage.epochs for stage in stages):,}, "
f"but training length is {int(total_epochs):,}."
)
return stages
def suggested_stages(final_resolution: int, total_epochs: int) -> list[ResolutionStage]:
"""Return an editable low-to-high schedule that always matches the budget."""
resolutions = [size for size in (64, 128, 256, 384, 512) if size <= int(final_resolution)]
if len(resolutions) < 2:
return [ResolutionStage(int(final_resolution), int(total_epochs))]
if int(total_epochs) < len(resolutions):
resolutions = resolutions[-int(total_epochs):]
return [ResolutionStage(resolution, 1) for resolution in resolutions]
# Front-load inexpensive structure learning while reserving final-resolution
# refinement. Normalizing lets the same policy work for any epoch budget.
weights = [0.60, 0.20, 0.10, 0.06, 0.04][-len(resolutions):]
allocation = [max(1, round(total_epochs * weight / sum(weights))) for weight in weights]
difference = int(total_epochs) - sum(allocation)
allocation[0] += difference
return [ResolutionStage(resolution, epochs) for resolution, epochs in zip(resolutions, allocation)]
def stage_batch_settings(
*, trainer: str, stage_resolution: int, final_resolution: int,
final_batch_size: int, base_accumulation: int, auto_batch: bool,
) -> tuple[int, int]:
"""Choose a conservative physical batch and matching accumulation count."""
requested_batch = max(1, min(64, int(final_batch_size)))
requested_accumulation = max(1, min(64, int(base_accumulation)))
if not auto_batch:
return requested_batch, requested_accumulation
caps = _BATCH_CAPS[trainer]
cap = caps[min(caps, key=lambda size: abs(size - int(stage_resolution)))]
# Scale from the user's final-stage batch according to image area, then
# enforce a trainer-specific safe cap. This never raises the 512px batch.
scaled = round(requested_batch * (int(final_resolution) / int(stage_resolution)) ** 2)
batch_size = max(1, min(cap, 64, scaled))
effective_batch = requested_batch * requested_accumulation
accumulation = max(1, min(64, -(-effective_batch // batch_size)))
return batch_size, accumulation
def stage_summary(stages: list[ResolutionStage]) -> str:
return ", ".join(f"{stage.resolution}px × {stage.epochs}" for stage in stages)
|