feat/energy_efficenty #38
18
example.py
18
example.py
|
@ -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
|
||||||
)
|
)
|
Loading…
Reference in New Issue