train.py

# train.py
# Training script for masked distance map modeling

import os
import sys
import argparse
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
from tqdm import tqdm

from model import MaskedAutoencoderViT
from dataset import get_dataloaders

def parse_args():
    parser = argparse.ArgumentParser(description='Train masked autoencoder for protein distance maps')
    
    # Dataset parameters
    parser.add_argument('--data_dir', type=str, default='./protein_fragment_dataset',
                        help='Directory with the dataset')
    parser.add_argument('--output_dir', type=str, default='./output',
                        help='Directory to save model checkpoints and logs')
    
    # Model parameters
    parser.add_argument('--img_size', type=int, default=64,
                        help='Size of input distance maps')
    parser.add_argument('--patch_size', type=int, default=8,
                        help='Size of patches')
    parser.add_argument('--embed_dim', type=int, default=768,
                        help='Embedding dimension')
    parser.add_argument('--depth', type=int, default=12,
                        help='Depth of transformer encoder')
    parser.add_argument('--num_heads', type=int, default=12,
                        help='Number of attention heads')
    parser.add_argument('--decoder_embed_dim', type=int, default=512,
                        help='Decoder embedding dimension')
    parser.add_argument('--decoder_depth', type=int, default=8,
                        help='Depth of transformer decoder')
    parser.add_argument('--decoder_num_heads', type=int, default=16,
                        help='Number of decoder attention heads')
    parser.add_argument('--mask_ratio', type=float, default=0.75,
                        help='Ratio of patches to mask')
    
    # Training parameters
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Batch size for training')
    parser.add_argument('--epochs', type=int, default=200,
                        help='Number of training epochs')
    parser.add_argument('--lr', type=float, default=1.5e-4,
                        help='Learning rate')
    parser.add_argument('--weight_decay', type=float, default=0.05,
                        help='Weight decay')
    parser.add_argument('--warmup_epochs', type=int, default=10,
                        help='Number of warmup epochs')
    parser.add_argument('--min_lr', type=float, default=1e-6,
                        help='Minimum learning rate')
    
    # Other parameters
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed')
    parser.add_argument('--num_workers', type=int, default=4,
                        help='Number of data loading workers')
    parser.add_argument('--device', type=str, default='cuda',
                        help='Device to use for training (cuda or cpu)')
    parser.add_argument('--resume', type=str, default=None,
                        help='Resume from checkpoint')
    parser.add_argument('--visualize_every', type=int, default=10,
                        help='Visualize reconstruction every n epochs')
    
    return parser.parse_args()

def get_lr_scheduler(optimizer, args, steps_per_epoch):
    """
    Create learning rate scheduler with linear warmup and cosine decay.
    """
    def lr_lambda(step):
        # Compute epoch from step
        epoch = step / steps_per_epoch
        
        # Linear warmup for warmup_epochs
        if epoch < args.warmup_epochs:
            return epoch / args.warmup_epochs
        
        # Cosine decay after warmup
        decay_epochs = args.epochs - args.warmup_epochs
        if decay_epochs == 0:
            return 1.0
            
        # Compute cosine decay
        decay_steps = decay_epochs * steps_per_epoch
        step_after_warmup = step - args.warmup_epochs * steps_per_epoch
        cosine_decay = 0.5 * (1 + np.cos(np.pi * step_after_warmup / decay_steps))
        
        # Scale between min_lr and max_lr
        return args.min_lr / args.lr + (1 - args.min_lr / args.lr) * cosine_decay
    
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)

def visualize_reconstruction(model, val_loader, device, epoch, output_dir, num_samples=5):
    """
    Visualize original, masked, and reconstructed distance maps.
    """
    os.makedirs(os.path.join(output_dir, 'visualizations'), exist_ok=True)
    
    model.eval()
    
    # Get a batch of validation samples
    batch = next(iter(val_loader))
    distance_maps = batch['distance_map'].to(device)
    
    # Only use a subset for visualization
    distance_maps = distance_maps[:num_samples]
    
    # Forward pass with fixed mask ratio for consistency
    with torch.no_grad():
        reconstructions, masks = model(distance_maps, mask_ratio=0.75)
    
    # Convert tensors to numpy arrays
    distance_maps = distance_maps.cpu().numpy()
    reconstructions = reconstructions.cpu().numpy()
    
    # Create a figure
    fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4 * num_samples))
    
    # Ensure axes is 2D even with a single sample
    if num_samples == 1:
        axes = np.expand_dims(axes, axis=0)
    
    for i in range(num_samples):
        # Original distance map
        axes[i, 0].imshow(distance_maps[i, 0], cmap='viridis')
        axes[i, 0].set_title('Original')
        axes[i, 0].axis('off')
        
        # Get mask for this sample
        mask = masks[i].cpu().numpy()
        
        # Process mask to get correct shape
        h, w = distance_maps[i, 0].shape
        ph, pw = h // 8, w // 8  # Assuming patch size is 8
        
        # Reshape mask to match patch grid, accounting for CLS token
        if len(mask) == ph*pw + 1:
            mask = mask[1:]  # Remove CLS token
            
        mask_2d = mask.reshape(ph, pw)
        
        # Upsample mask to image size
        mask_full = np.kron(mask_2d, np.ones((8, 8)))
        
        # Create masked version for visualization (0 = masked)
        masked_map = distance_maps[i, 0].copy()
        masked_map[mask_full < 0.5] = 0  # Apply mask where value is < 0.5
        
        # Masked distance map
        axes[i, 1].imshow(masked_map, cmap='viridis')
        axes[i, 1].set_title('Masked (Input)')
        axes[i, 1].axis('off')
        
        # Reconstructed distance map
        axes[i, 2].imshow(reconstructions[i, 0], cmap='viridis')
        axes[i, 2].set_title('Reconstructed')
        axes[i, 2].axis('off')
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'visualizations', f'reconstruction_epoch_{epoch}.png'))
    plt.close()
    
    model.train()

def train_one_epoch(model, dataloader, criterion, optimizer, scheduler, device, epoch):
    """
    Train model for one epoch.
    """
    model.train()
    
    loss_meter = AverageMeter()
    pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
    
    for batch_idx, batch in enumerate(pbar):
        # Get data
        distance_maps = batch['distance_map'].to(device)
        
        # Forward pass
        reconstructions, masks = model(distance_maps)
        
        # Apply mask to loss computation (only compute loss on masked patches)
        target = distance_maps
        
        # Reshape masks to match the input dimensions
        batch_size, channels, h, w = target.shape
        ph, pw = h // model.patch_size, w // model.patch_size
        
        # Need to ensure mask has the right dimensions
        # Calculate number of patches (excluding CLS token)
        num_patches = ph * pw
        
        # Ensure mask has the correct size before reshaping
        if masks.size(1) != num_patches:
            # In case mask includes CLS token or has different dimensions
            # We need to resize it to match expected patch dimensions
            # First, verify what we're getting from the model
            if masks.size(1) == num_patches + 1:  # +1 for CLS token
                # If mask includes CLS token, remove it
                masks = masks[:, 1:]
                
        # Now reshape to 2D patch grid
        mask_resized = masks.reshape(batch_size, 1, ph, pw)
        mask_resized = mask_resized.repeat(1, channels, 1, 1)
        
        # Upsample mask to match image dimensions
        mask_full = torch.nn.functional.interpolate(
            mask_resized, 
            size=(h, w), 
            mode='nearest'
        )
        
        # Compute loss on masked patches only
        # We use (1 - mask) because 0 means masked in our implementation
        loss_mask = 1.0 - mask_full
        loss = criterion(reconstructions * loss_mask, target * loss_mask)
        
        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Update learning rate
        scheduler.step()
        
        # Update metrics
        loss_meter.update(loss.item(), distance_maps.size(0))
        
        # Update progress bar
        pbar.set_postfix({'loss': loss_meter.avg})
    
    return loss_meter.avg

def validate(model, dataloader, criterion, device):
    """
    Validate model on validation set.
    """
    model.eval()
    
    loss_meter = AverageMeter()
    
    with torch.no_grad():
        for batch in dataloader:
            # Get data
            distance_maps = batch['distance_map'].to(device)
            
            # Forward pass
            reconstructions, masks = model(distance_maps)
            
            # Apply mask to loss computation (only compute loss on masked patches)
            target = distance_maps
            
            # Reshape masks to match the input dimensions
            batch_size, channels, h, w = target.shape
            ph, pw = h // model.patch_size, w // model.patch_size
            
            # Need to ensure mask has the right dimensions
            # Calculate number of patches (excluding CLS token)
            num_patches = ph * pw
            
            # Ensure mask has the correct size before reshaping
            if masks.size(1) != num_patches:
                # In case mask includes CLS token or has different dimensions
                if masks.size(1) == num_patches + 1:  # +1 for CLS token
                    # If mask includes CLS token, remove it
                    masks = masks[:, 1:]
                    
            # Now reshape to 2D patch grid
            mask_resized = masks.reshape(batch_size, 1, ph, pw)
            mask_resized = mask_resized.repeat(1, channels, 1, 1)
            
            # Upsample mask to match image dimensions
            mask_full = torch.nn.functional.interpolate(
                mask_resized, 
                size=(h, w), 
                mode='nearest'
            )
            
            # Compute loss on masked patches only
            # We use (1 - mask) because 0 means masked in our implementation
            loss_mask = 1.0 - mask_full
            loss = criterion(reconstructions * loss_mask, target * loss_mask)
            
            # Update metrics
            loss_meter.update(loss.item(), distance_maps.size(0))
    
    return loss_meter.avg

class AverageMeter:
    """
    Compute and store the average and current value.
    """
    def __init__(self):
        self.reset()
        
    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def save_checkpoint(state, is_best, output_dir):
    """
    Save model checkpoint.
    """
    filename = os.path.join(output_dir, 'checkpoint.pth')
    torch.save(state, filename)
    if is_best:
        best_filename = os.path.join(output_dir, 'model_best.pth')
        torch.save(state, best_filename)

def main():
    # Parse arguments
    args = parse_args()
    
    # Set random seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Set device
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Create dataloaders
    print("Creating dataloaders...")
    dataloaders = get_dataloaders(
        args.data_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers
    )
    train_loader = dataloaders['train']
    val_loader = dataloaders['val']
    test_loader = dataloaders['test']
    
    # Create model
    print("Creating model...")
    model = MaskedAutoencoderViT(
        img_size=args.img_size,
        patch_size=args.patch_size,
        embed_dim=args.embed_dim,
        depth=args.depth,
        num_heads=args.num_heads,
        decoder_embed_dim=args.decoder_embed_dim,
        decoder_depth=args.decoder_depth,
        decoder_num_heads=args.decoder_num_heads,
        mask_ratio=args.mask_ratio
    ).to(device)
    
    # Print model statistics
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Number of learnable parameters: {n_parameters / 1e6:.2f}M")
    
    # Define loss function, optimizer, and scheduler
    criterion = nn.MSELoss()
    optimizer = optim.AdamW(
        model.parameters(),
        lr=args.lr,
        weight_decay=args.weight_decay
    )
    
    # Create scheduler with linear warmup and cosine decay
    steps_per_epoch = len(train_loader)
    scheduler = get_lr_scheduler(optimizer, args, steps_per_epoch)
    
    # Create tensorboard writer
    writer = SummaryWriter(log_dir=os.path.join(args.output_dir, 'tensorboard'))
    
    # Resume from checkpoint if specified
    start_epoch = 0
    best_val_loss = float('inf')
    if args.resume is not None:
        if os.path.isfile(args.resume):
            print(f"Loading checkpoint '{args.resume}'")
            checkpoint = torch.load(args.resume, map_location=device)
            start_epoch = checkpoint['epoch']
            best_val_loss = checkpoint['best_val_loss']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            scheduler.load_state_dict(checkpoint['scheduler'])
            print(f"Loaded checkpoint '{args.resume}' (epoch {checkpoint['epoch']})")
        else:
            print(f"No checkpoint found at '{args.resume}'")
    
    # Training loop
    print("Starting training...")
    for epoch in range(start_epoch, args.epochs):
        # Train for one epoch
        train_loss = train_one_epoch(
            model, train_loader, criterion, optimizer, scheduler, device, epoch
        )
        
        # Validate
        val_loss = validate(model, val_loader, criterion, device)
        
        # Visualize reconstructions periodically
        if (epoch + 1) % args.visualize_every == 0 or epoch == 0:
            visualize_reconstruction(model, val_loader, device, epoch + 1, args.output_dir)
        
        # Log metrics
        writer.add_scalar('Loss/train', train_loss, epoch)
        writer.add_scalar('Loss/val', val_loss, epoch)
        writer.add_scalar('LR', optimizer.param_groups[0]['lr'], epoch)
        
        # Print metrics
        print(f"Epoch: {epoch+1}/{args.epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
        
        # Save checkpoint
        is_best = val_loss < best_val_loss
        best_val_loss = min(val_loss, best_val_loss)
        
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_val_loss': best_val_loss,
            'optimizer': optimizer.state_dict(),
            'scheduler': scheduler.state_dict(),
            'args': args,
        }, is_best, args.output_dir)
    
    # Final evaluation on test set
    test_loss = validate(model, test_loader, criterion, device)
    print(f"Test Loss: {test_loss:.4f}")
    
    # Close tensorboard writer
    writer.close()

if __name__ == '__main__':
    main()