File size: 1,005 Bytes
2adb41a | 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 | """
ml/training/train_sequence_vae.py
Train the 1D-CNN Sequence VAE on normal track geometry sequences (tc.v1 SOTA).
"""
from ml.scripts.train_sequence_vae import train_sequence_vae
__all__ = ["train_sequence_vae"]
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Train 1D-CNN Sequence VAE for Novel Geometry.")
parser.add_argument("--epochs", type=int, default=25, help="Number of training epochs")
parser.add_argument("--beta", type=float, default=0.01, help="Beta-VAE KL weight")
parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
parser.add_argument("--save-path", default="artifacts/checkpoints/geometry/sequence_vae.pt")
args = parser.parse_args()
train_sequence_vae(
epochs=args.epochs,
beta=args.beta,
batch_size=args.batch_size,
lr=args.lr,
save_path=args.save_path,
)
|