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

added progressbar for batches

This commit is contained in:
Falko Victor Habel 2025-01-27 09:13:22 +01:00
parent b546f4ee27
commit 6c146f2767
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

@ -4,8 +4,8 @@ import pandas as pd
from aiia.model.config import AIIAConfig
from aiia.model import AIIABase
from aiia.data.DataLoader import AIIADataLoader
import os
import copy
from tqdm import tqdm
def pretrain_model(data_path1, data_path2, num_epochs=3):
# Read and merge datasets
@ -49,7 +49,7 @@ def pretrain_model(data_path1, data_path2, num_epochs=3):
aiia_loader = AIIADataLoader(
merged_df,
column="image_bytes",
batch_size=4,
batch_size=2,
pretraining=True,
collate_fn=safe_collate
)
@ -75,7 +75,7 @@ def pretrain_model(data_path1, data_path2, num_epochs=3):
model.train()
total_train_loss = 0.0
for batch in train_loader:
for batch in tqdm(train_loader):
if batch is None:
continue # Skip empty batches