File size: 3,087 Bytes
e3c7370 |
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 |
#!/usr/bin/env python3
"""
OmniCoreX Training CLI Script
Launches the training process for OmniCoreX with command-line options
to configure parameters. Supports distributed and mixed precision training.
"""
import argparse
import os
from utils import set_seed, load_config_file, setup_logging
from model import OmniCoreXModel
from data_loader import create_omncorex_dataloader
from trainer import Trainer
def parse_args():
parser = argparse.ArgumentParser(description="Train OmniCoreX Model")
parser.add_argument("--config", type=str, default="config.yaml", help="Path to config file")
parser.add_argument("--epochs", type=int, help="Number of epochs to train")
parser.add_argument("--batch_size", type=int, help="Batch size override")
parser.add_argument("--seed", type=int, help="Random seed override")
parser.add_argument("--log_file", type=str, help="Log file path")
parser.add_argument("--no_mixed_precision", action="store_true", help="Disable mixed precision training")
return parser.parse_args()
def main():
args = parse_args()
config = load_config_file(args.config)
if args.seed:
set_seed(args.seed)
else:
set_seed(config.get("seed", 42))
logger = setup_logging(log_file=args.log_file)
logger.info("Starting OmniCoreX training")
data_cfg = config["data"]
training_cfg = config.get("training", {})
batch_size = args.batch_size or data_cfg.get("batch_size", 16)
mixed_precision = not args.no_mixed_precision and training_cfg.get("mixed_precision", True)
train_loader = create_omncorex_dataloader(
metadata_path=data_cfg["metadata_path"],
modalities=data_cfg["modalities"],
tokenizer=None, # Assume custom tokenizer later
batch_size=batch_size,
shuffle=data_cfg.get("shuffle", True),
num_workers=data_cfg.get("num_workers", 4),
augmentation=data_cfg.get("augmentation", False)
)
model_cfg = config["model"]
model = OmniCoreXModel(
stream_configs=model_cfg["streams"],
embed_dim=model_cfg.get("architecture", {}).get("embed_dim", 768),
num_layers=model_cfg.get("architecture", {}).get("num_layers", 24),
num_heads=model_cfg.get("architecture", {}).get("num_heads", 12),
dropout=model_cfg.get("architecture", {}).get("dropout", 0.1)
)
trainer = Trainer(
model=model,
train_loader=train_loader,
valid_loader=None,
save_dir=training_cfg.get("save_dir", "./checkpoints"),
lr=training_cfg.get("learning_rate", 5e-5),
weight_decay=training_cfg.get("weight_decay", 0.01),
max_grad_norm=training_cfg.get("max_grad_norm", 1.0),
accumulation_steps=training_cfg.get("accumulation_steps", 1),
total_steps=training_cfg.get("total_steps", 100000),
warmup_steps=training_cfg.get("warmup_steps", 1000),
device=None,
mixed_precision=mixed_precision
)
epochs = args.epochs or training_cfg.get("epochs", 1)
trainer.fit(epochs=epochs)
if __name__ == "__main__":
main()
|