src/aiunn/finetune.py

@@ -7,13 +7,15 @@ from aiia import AIIABase
 import csv
 from tqdm import tqdm
 import base64
+from torch.amp import autocast, GradScaler
+
 
 class UpscaleDataset(Dataset):
     def __init__(self, parquet_files: list, transform=None):
         combined_df = pd.DataFrame()
         for parquet_file in parquet_files:
             # Load data with chunking for memory efficiency
-            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(10000)
+            df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(5000)
             combined_df = pd.concat([combined_df, df], ignore_index=True)
         
         # Validate data format
@@ -80,8 +82,6 @@ transform = transforms.Compose([
 ])
 
 
-import torch
-
 # Replace with your actual pretrained model path
 pretrained_model_path = "/root/vision/AIIA/AIIA-base-512"  
 
@@ -94,7 +94,7 @@ from torch.utils.data import DataLoader
 
 # Create your dataset and dataloader
 dataset = UpscaleDataset(["/root/training_data/vision-dataset/image_upscaler.parquet", "/root/training_data/vision-dataset/image_vec_upscaler.parquet"], transform=transform)
-data_loader = DataLoader(dataset, batch_size=4, shuffle=True)
+data_loader = DataLoader(dataset, batch_size=2, shuffle=True)
 
 # Define a loss function and optimizer
 criterion = nn.MSELoss()
@@ -113,29 +113,34 @@ with open(csv_file, mode='a', newline='') as file:
     if file.tell() == 0:
         writer.writerow(['Epoch', 'Train Loss'])
 
+# Create a gradient scaler (for scaling gradients when using AMP)
+scaler = GradScaler()
+
 for epoch in range(num_epochs):
     epoch_loss = 0.0
-    # Wrap the data_loader with tqdm for progress tracking
     data_loader_with_progress = tqdm(data_loader, desc=f"Epoch {epoch + 1}")
-    print(f"Epoche: {epoch}")
     for low_res, high_res in data_loader_with_progress:
-        low_res = low_res.to(device)
+        low_res = low_res.to(device, non_blocking=True)
-        high_res = high_res.to(device)
+        high_res = high_res.to(device, non_blocking=True)
 
         optimizer.zero_grad()
+        
+        # Use automatic mixed precision context
+        with autocast():
             outputs = model(low_res)
             loss = criterion(outputs, high_res)
-        loss.backward()
-        optimizer.step()
-        epoch_loss += loss.item()
         
-    avg_epoch_loss = epoch_loss / len(data_loader)
+        scaler.scale(loss).backward()
-    print(f"Epoch {epoch + 1}, Loss: {avg_epoch_loss}")
+        scaler.step(optimizer)
+        scaler.update()
+
+        epoch_loss += loss.item()
+    print(f"Epoch {epoch + 1}, Loss: {epoch_loss}")
 
     # Append the training loss to the CSV file
     with open(csv_file, mode='a', newline='') as file:
         writer = csv.writer(file)
-        writer.writerow([epoch + 1, avg_epoch_loss])
+        writer.writerow([epoch + 1, epoch_loss])
     
 # Optionally, save the finetuned model to a new directory
 finetuned_model_path = "aiuNN"