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

added tqdm to improve code

This commit is contained in:
Falko Victor Habel 2025-02-08 12:52:48 +01:00
parent a48bad5a49
commit 20dd9f68ed
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/aiunn/finetune.py
View File

@ -15,7 +15,7 @@ from aiia import AIIA, AIIAConfig, AIIABase, AIIABaseShared, AIIAmoe, AIIAchunke
class aiuNNDataset(torch.utils.data.Dataset): class aiuNNDataset(torch.utils.data.Dataset):
def __init__(self, parquet_path): def __init__(self, parquet_path):
# Read the Parquet file # Read the Parquet file
self.df = pd.read_parquet(parquet_path).head(2500) self.df = pd.read_parquet(parquet_path).head(1250)
# Data augmentation pipeline without Resize as it's redundant # Data augmentation pipeline without Resize as it's redundant
self.augmentation = Compose([ self.augmentation = Compose([
@ -124,12 +124,15 @@ def finetune_model(model: AIIA, datasets:list[str], batch_size=2, epochs=10):
best_val_loss = float('inf') best_val_loss = float('inf')
from tqdm import tqdm
for epoch in range(epochs): for epoch in range(epochs):
model.train() model.train()
train_loss = 0.0 train_loss = 0.0
for batch_idx, batch in enumerate(train_loader): for batch_idx, batch in enumerate(tqdm(train_loader)):
# Your training code here
low_res = batch['low_res'].to(device) low_res = batch['low_res'].to(device)
high_res = batch['high_res'].to(device) high_res = batch['high_res'].to(device)
@ -155,7 +158,7 @@ def finetune_model(model: AIIA, datasets:list[str], batch_size=2, epochs=10):
val_loss = 0.0 val_loss = 0.0
with torch.no_grad(): with torch.no_grad():
for batch in val_loader: for batch in tqdm(val_loader, desc="Validation"):
low_res = batch['low_res'].to(device) low_res = batch['low_res'].to(device)
high_res = batch['high_res'].to(device) high_res = batch['high_res'].to(device)