Machine-Learning
/
AIIA
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

downszied batchsize

This commit is contained in:
Falko Victor Habel 2025-01-27 09:02:22 +01:00
parent 32526c3c30
commit 8dad1d7150
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/pretrain.py
View File

@ -9,13 +9,13 @@ import copy
def pretrain_model(data_path1, data_path2, num_epochs=3): def pretrain_model(data_path1, data_path2, num_epochs=3):
# Read and merge datasets # Read and merge datasets
df1 = pd.read_parquet(data_path1).head(2500) df1 = pd.read_parquet(data_path1).head(10000)
df2 = pd.read_parquet(data_path2).head(2500) df2 = pd.read_parquet(data_path2).head(10000)
merged_df = pd.concat([df1, df2], ignore_index=True) merged_df = pd.concat([df1, df2], ignore_index=True)
# Model configuration # Model configuration
config = AIIAConfig( config = AIIAConfig(
model_name="AIIA-Base-512x5k", model_name="AIIA-Base-512x20k",
) )
# Initialize model and data loader # Initialize model and data loader
@ -49,7 +49,7 @@ def pretrain_model(data_path1, data_path2, num_epochs=3):
aiia_loader = AIIADataLoader( aiia_loader = AIIADataLoader(
merged_df, merged_df,
column="image_bytes", column="image_bytes",
batch_size=32, batch_size=8,
pretraining=True, pretraining=True,
collate_fn=safe_collate collate_fn=safe_collate
) )