Logic/mamba_diffusion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add configurable dt sampling and loss toggles

Browse Source
This commit is contained in:
gameloader
2026-01-21 15:14:04 +08:00
parent c58a73ae26
commit cac3236f9d
3 changed files with 207 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
main.py
View File

@@ -6,7 +6,6 @@ from as_mamba import TrainConfig, run_training_and_plot
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Train AS-Mamba on sphere-to-sphere flow.")
parser.add_argument("--epochs", type=int, default=None)
parser.add_argument("--warmup-epochs", type=int, default=None)
parser.add_argument("--batch-size", type=int, default=None)
parser.add_argument("--steps-per-epoch", type=int, default=None)
parser.add_argument("--seq-len", type=int, default=None)
@@ -15,6 +14,27 @@ def build_parser() -> argparse.ArgumentParser:
parser.add_argument("--output-dir", type=str, default=None)
parser.add_argument("--project", type=str, default=None)
parser.add_argument("--run-name", type=str, default=None)
parser.add_argument("--dt-alpha", type=float, default=None)
parser.add_argument("--dt-min", type=float, default=None)
parser.add_argument("--dt-max", type=float, default=None)
parser.add_argument("--lambda-flow", type=float, default=None)
parser.add_argument("--lambda-pos", type=float, default=None)
parser.add_argument("--lambda-dt", type=float, default=None)
parser.add_argument(
"--use-flow-loss",
action=argparse.BooleanOptionalAction,
default=None,
)
parser.add_argument(
"--use-pos-loss",
action=argparse.BooleanOptionalAction,
default=None,
)
parser.add_argument(
"--use-dt-loss",
action=argparse.BooleanOptionalAction,
default=None,
)
parser.add_argument("--val-every", type=int, default=None)
parser.add_argument("--val-samples", type=int, default=None)
parser.add_argument("--val-plot-samples", type=int, default=None)