AIIA/src/pretrain.py

import torch
from torch import nn, utils
import pandas as pd
from aiia.model.config import AIIAConfig
from aiia.model import AIIABase
from aiia.data.DataLoader import AIIADataLoader
from tqdm import tqdm

class ProjectionHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_denoise = nn.Conv2d(512, 3, kernel_size=1)
        self.conv_rotate = nn.Conv2d(512, 4, kernel_size=1)  # 4 classes for 0, 90, 180, 270 degrees
        
    def forward(self, x, task='denoise'):
        if task == 'denoise':
            return self.conv_denoise(x)
        else:
            return self.conv_rotate(x).mean(dim=(2, 3))  # Global average pooling for rotation task

def pretrain_model(data_path1, data_path2, num_epochs=3):
    # Read and merge datasets
    df1 = pd.read_parquet(data_path1).head(10000)
    df2 = pd.read_parquet(data_path2).head(10000)
    merged_df = pd.concat([df1, df2], ignore_index=True)

    # Model configuration
    config = AIIAConfig(
        model_name="AIIA-Base-512x20k",
    )

    # Initialize model and projection head
    model = AIIABase(config)
    projection_head = ProjectionHead()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    projection_head.to(device)

    def safe_collate(batch):
        denoise_batch = []
        rotate_batch = []
        
        for sample in batch:
            try:
                noisy_img, target, task = sample
                if task == 'denoise':
                    denoise_batch.append({
                        'image': noisy_img,
                        'target': target,
                        'task': task
                    })
                else:  # rotate task
                    rotate_batch.append({
                        'image': noisy_img,
                        'target': target,
                        'task': task
                    })
            except Exception as e:
                print(f"Skipping sample due to error: {e}")
                continue
                
        if not denoise_batch and not rotate_batch:
            return None

        batch_data = {
            'denoise': None,
            'rotate': None
        }

        if denoise_batch:
            images = torch.stack([x['image'] for x in denoise_batch])
            targets = torch.stack([x['target'] for x in denoise_batch])
            batch_data['denoise'] = (images, targets)

        if rotate_batch:
            images = torch.stack([x['image'] for x in rotate_batch])
            targets = torch.stack([x['target'] for x in rotate_batch])
            batch_data['rotate'] = (images, targets)

        return batch_data

    aiia_loader = AIIADataLoader(
        merged_df,
        column="image_bytes",
        batch_size=2,
        pretraining=True,
        collate_fn=safe_collate
    )

    train_loader = aiia_loader.train_loader
    val_loader = aiia_loader.val_loader

    criterion_denoise = nn.MSELoss()
    criterion_rotate = nn.CrossEntropyLoss()
    
    # Update optimizer to include projection head parameters
    optimizer = torch.optim.AdamW(
        list(model.parameters()) + list(projection_head.parameters()), 
        lr=config.learning_rate
    )

    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch+1}/{num_epochs}")
        print("-" * 20)
        
        # Training phase
        model.train()
        projection_head.train()
        total_train_loss = 0.0
        batch_count = 0
        
        for batch_data in tqdm(train_loader):
            if batch_data is None:
                continue
            
            optimizer.zero_grad()
            batch_loss = 0
            
            # Handle denoise task
            if batch_data['denoise'] is not None:
                noisy_imgs, targets = batch_data['denoise']
                noisy_imgs = noisy_imgs.to(device)
                targets = targets.to(device)
                
                # Get features from base model
                features = model(noisy_imgs)
                # Project features back to image space
                outputs = projection_head(features, task='denoise')
                
                print(f"\nDenoising task shapes:")
                print(f"Input shape: {noisy_imgs.shape}")
                print(f"Target shape: {targets.shape}")
                print(f"Features shape: {features.shape}")
                print(f"Output shape: {outputs.shape}")
                
                loss = criterion_denoise(outputs, targets)
                batch_loss += loss
            
            # Handle rotate task
            if batch_data['rotate'] is not None:
                imgs, targets = batch_data['rotate']
                imgs = imgs.to(device)
                targets = targets.long().to(device)
                
                # Get features from base model
                features = model(imgs)
                # Project features to rotation predictions
                outputs = projection_head(features, task='rotate')
                
                print(f"\nRotation task shapes:")
                print(f"Input shape: {imgs.shape}")
                print(f"Target shape: {targets.shape}")
                print(f"Features shape: {features.shape}")
                print(f"Output shape: {outputs.shape}")
                
                loss = criterion_rotate(outputs, targets)
                batch_loss += loss
            
            if batch_loss > 0:
                batch_loss.backward()
                optimizer.step()
                total_train_loss += batch_loss.item()
                batch_count += 1
        
        avg_train_loss = total_train_loss / max(batch_count, 1)
        print(f"Training Loss: {avg_train_loss:.4f}")

        # Validation phase
        model.eval()
        projection_head.eval()
        val_loss = 0.0
        val_batch_count = 0
        
        with torch.no_grad():
            for batch_data in val_loader:
                if batch_data is None:
                    continue
                
                batch_loss = 0
                
                if batch_data['denoise'] is not None:
                    noisy_imgs, targets = batch_data['denoise']
                    noisy_imgs = noisy_imgs.to(device)
                    targets = targets.to(device)
                    
                    features = model(noisy_imgs)
                    outputs = projection_head(features, task='denoise')
                    loss = criterion_denoise(outputs, targets)
                    batch_loss += loss
                
                if batch_data['rotate'] is not None:
                    imgs, targets = batch_data['rotate']
                    imgs = imgs.to(device)
                    targets = targets.long().to(device)
                    
                    features = model(imgs)
                    outputs = projection_head(features, task='rotate')
                    loss = criterion_rotate(outputs, targets)
                    batch_loss += loss
                
                if batch_loss > 0:
                    val_loss += batch_loss.item()
                    val_batch_count += 1
        
        avg_val_loss = val_loss / max(val_batch_count, 1)
        print(f"Validation Loss: {avg_val_loss:.4f}")
        
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            # Save both model and projection head
            model.save("AIIA-base-512")
            print("Best model saved!")

if __name__ == "__main__":
    data_path1 = "/root/training_data/vision-dataset/images_checkpoint.parquet"
    data_path2 = "/root/training_data/vision-dataset/vec_images_dataset.parquet"
    pretrain_model(data_path1, data_path2, num_epochs=3)