feat/energy_efficenty #38

Merged
Fabel merged 6 commits from feat/energy_efficenty into develop 2025-04-17 10:52:25 +00:00
1 changed files with 8 additions and 12 deletions
Showing only changes of commit 7d24de1f7e - Show all commits

View File

@ -6,19 +6,15 @@ from aiia.pretrain import Pretrainer
config = AIIAConfig(model_name="AIIA-Base-512x20k") config = AIIAConfig(model_name="AIIA-Base-512x20k")
model = AIIABase(config) model = AIIABase(config)
# Initialize pretrainer with the model
pretrainer = Pretrainer(model, learning_rate=1e-4, config=config) pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)
# List of dataset paths # Set checkpoint directory
dataset_paths = [ checkpoint_dir = "checkpoints/my_model"
"/path/to/dataset1.parquet",
"/path/to/dataset2.parquet"
]
# Start training with multiple datasets # Start training (will automatically load checkpoint if available)
pretrainer.train( pretrainer.train(
dataset_paths=dataset_paths, dataset_paths=["path/to/dataset1.parquet", "path/to/dataset2.parquet"],
num_epochs=10, output_path="trained_models/my_model",
batch_size=2, checkpoint_dir=checkpoint_dir,
sample_size=10000 num_epochs=10
) )