autoencoder.py

import torch, torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import scanpy as sc
import numpy as np
import os
import matplotlib.pyplot as plt
from pathlib import Path
import glob
import time
from datetime import datetime
import torch.nn.functional as F

class AE(nn.Module):
    def __init__(self, n_genes, z_dim=128, drop=0.2):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Linear(n_genes, 1024), 
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Dropout(drop),
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, z_dim))
        self.dec = nn.Sequential(
            nn.Linear(z_dim, 512), 
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(drop),
            nn.Linear(512, 1024), 
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, n_genes))
    def forward(self, x):
        z = self.enc(x)
        return self.dec(z), z
    def __repr__(self):
        return ""
        
        
def train_model(model, train_loader, val_loader, loss_fn, optimizer, scheduler, 
                n_epochs, device, start_epoch=0, patience=20, 
                log_file_path='training_log.txt', model_save_dir='checkpoints'):
    os.makedirs(model_save_dir, exist_ok=True)
    best_ckpt = os.path.join(model_save_dir, 'best_model.pth')
    latest_ckpt = os.path.join(model_save_dir, 'latest_model.pth')
    final_ckpt = os.path.join(model_save_dir, 'final_model.pth')   

    best_val_loss = float('inf')
    patience_counter = 0
    best_model_state = None
    best_epoch_info = None    

    train_losses = []
    val_losses = []
    learning_rates = []    

    log_file = open(log_file_path, 'a')
    log_file.write(f"\n{'='*60}\n")
    log_file.write(f"Training started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    log_file.write(f"Total epochs: {n_epochs}, Start epoch: {start_epoch}, Patience: {patience}\n")
    log_file.write(f"{'='*60}\n")
    start_time = time.time()    

    for epoch in range(start_epoch, n_epochs):
        epoch_start_time = time.time()
        model.train()
        running_loss = 0.0
        train_samples = 0        
        for batch_idx, (xx,) in enumerate(train_loader):
            xx = xx.to(device)
            rec, z = model(xx)
            loss = loss_fn(rec, xx)            

            optimizer.zero_grad()
            loss.backward()            

            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()
            running_loss += loss.item() * xx.size(0)
            train_samples += xx.size(0)            

        epoch_loss = running_loss / len(train_loader.dataset)
        train_losses.append(epoch_loss)
        if scheduler is not None:
            if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(epoch_loss)
            else:
                scheduler.step()
        current_lr = optimizer.param_groups[0]['lr']
        learning_rates.append(current_lr)        

        model.eval()
        val_loss = 0.0        
        with torch.no_grad():
            for xx, in val_loader:
                xx = xx.to(device)
                rec, _ = model(xx)
                val_loss += loss_fn(rec, xx).item() * xx.size(0)        

        val_loss /= len(val_loader.dataset)
        val_losses.append(val_loss)        

        if val_loss < best_val_loss:
            improvement = (best_val_loss - val_loss) / max(best_val_loss, 1e-8)
            best_val_loss = val_loss
            patience_counter = 0           

            best_model_state = {
                'epoch': epoch + 1,
                'model_state_dict': {k: v.cpu().clone() for k, v in model.state_dict().items()},
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
                'train_loss': epoch_loss,
                'val_loss': val_loss,
                'learning_rate': current_lr,
            }            
            best_epoch_info = {
                'epoch': epoch + 1,
                'train_loss': epoch_loss,
                'val_loss': val_loss,
                'improvement': improvement,
            }            
            print(f'? Epoch {epoch+1:03d}: Validation loss improved to {val_loss:.4f} '
                  f'(improvement: {improvement:.2%})')
        else:
            patience_counter += 1
            improvement = 0.0
            print(f'  Epoch {epoch+1:03d}: Validation loss did not improve '
                  f'({patience_counter}/{patience})')        

        if (epoch + 1) % 50 == 0 or epoch == n_epochs - 1:
            torch.save({
                'epoch': epoch + 1,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
                'train_loss': epoch_loss,
                'val_loss': val_loss,
                'best_val_loss': best_val_loss,
                'train_losses': train_losses,
                'val_losses': val_losses,
                'learning_rates': learning_rates,
            }, latest_ckpt)
            print(f'  Latest model saved at epoch {epoch+1}')       

        epoch_time = time.time() - epoch_start_time
        log_file.write(f'epoch {epoch+1:03d}/{n_epochs} | '
                      f'time: {epoch_time:.1f}s | '
                      f'train_loss: {epoch_loss:.4f} | '
                      f'val_loss: {val_loss:.4f} | '
                      f'lr: {current_lr:.2e} | '
                      f'patience: {patience_counter}/{patience}\n')
        log_file.flush()        

        print(f'Epoch {epoch+1:03d}/{n_epochs} | '
              f'Train Loss: {epoch_loss:.4f} | '
              f'Val Loss: {val_loss:.4f} | '
              f'LR: {current_lr:.2e} | '
              f'Time: {epoch_time:.1f}s')        

        if patience_counter >= patience:
            print(f'\n{"="*60}')
            print(f'Early stopping triggered at epoch {epoch + 1}')
            print(f'Best model was at epoch {best_epoch_info["epoch"]} '
                  f'with val_loss={best_epoch_info["val_loss"]:.4f}')
            print(f'{"="*60}\n')
            break    

    total_time = time.time() - start_time   

    torch.save({
        'epoch': epoch + 1,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
        'train_loss': epoch_loss,
        'val_loss': val_loss,
        'best_val_loss': best_val_loss,
        'train_losses': train_losses,
        'val_losses': val_losses,
        'learning_rates': learning_rates,
        'total_training_time': total_time,
    }, final_ckpt)   

    if best_model_state is not None:
        torch.save(best_model_state, best_ckpt)
        print(f'\n? Best model saved to: {best_ckpt}')
        print(f'  - From epoch: {best_epoch_info["epoch"]}')
        print(f'  - Validation loss: {best_epoch_info["val_loss"]:.4f}')
        print(f'  - Train loss: {best_epoch_info["train_loss"]:.4f}')    

    log_file.write(f'\n{"="*60}\n')
    log_file.write(f"Training completed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    log_file.write(f"Total training time: {total_time:.2f}s ({total_time/3600:.2f} hours)\n")
    log_file.write(f"Total epochs trained: {epoch + 1}\n")
    log_file.write(f"Best epoch: {best_epoch_info['epoch'] if best_epoch_info else 'N/A'}\n")
    log_file.write(f"Best validation loss: {best_val_loss:.4f}\n")
    log_file.write(f"Early stopped: {'Yes' if patience_counter >= patience else 'No'}\n")
    log_file.write(f'\n{"="*60}\n')
    log_file.close()    

    print(f"\n{'='*60}")
    print("TRAINING SUMMARY")
    print(f"{'='*60}")
    print(f"Total training time: {total_time:.2f}s ({total_time/3600:.2f} hours)")
    print(f"Total epochs: {epoch + 1}")
    if best_epoch_info:
        print(f"Best model - Epoch {best_epoch_info['epoch']}:")
        print(f"  Train loss: {best_epoch_info['train_loss']:.4f}")
        print(f"  Val loss:   {best_epoch_info['val_loss']:.4f}")
    print(f"Final model saved to: {final_ckpt}")
    print(f"Best model saved to: {best_ckpt}")
    print(f"Latest model saved to: {latest_ckpt}")
    print(f"Training log saved to: {log_file_path}")
    print(f"{'='*60}")    

    return {
        'best_epoch': best_epoch_info['epoch'] if best_epoch_info else None,
        'best_val_loss': best_val_loss,
        'best_model_path': best_ckpt,
        'final_model_path': final_ckpt,
        'train_losses': train_losses,
        'val_losses': val_losses,
        'learning_rates': learning_rates,
        'total_training_time': total_time,
        'early_stopped': patience_counter >= patience,
    }

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    model = AE(X.shape[1], 128).to(device) 
    BATCH = 256
    train_loader = DataLoader(train_set, batch_size=BATCH, shuffle=True)
    val_loader   = DataLoader(val_set,   batch_size=BATCH, shuffle=False)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.5, verbose=False)
    results = train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        loss_fn=loss_fn,
        optimizer=optimizer,
        scheduler=scheduler,
        n_epochs=500,
        device='cuda' if torch.cuda.is_available() else 'cpu',
        start_epoch=0,
        patience=20,
        log_file_path='./model/training_log.txt',
        model_save_dir='./model/'
    )    

    def load_best_model_for_inference(model_path, model):
        checkpoint = torch.load(model_path, map_location='cpu')
        model.load_state_dict(checkpoint['model_state_dict'])
        model.eval()
        print(f"Loaded best model from epoch {checkpoint['epoch']}")
        print(f"Validation loss: {checkpoint['val_loss']:.4f}")
        return model
    best_model = load_best_model_for_inference(results['best_model_path'], model)

if __name__ == "__main__":
    adata = sc.read_h5ad('./data/MouseSC.h5ad')
    sc.pp.normalize_total(adata, target_sum=1e4)
    sc.pp.log1p(adata)
    sc.pp.highly_variable_genes(adata, n_top_genes=2000, subset=True)  
    sc.pp.scale(adata, max_value=10) 
    X = adata.X.astype('float32')
    if hasattr(X, 'toarray'):
        X = X.toarray()    
    X = torch.from_numpy(X)

    perm = torch.randperm(len(X))
    train_idx = perm[:int(0.9*len(X))]
    val_idx   = perm[int(0.9*len(X)):]    
    train_set = TensorDataset(X[train_idx])
    val_set   = TensorDataset(X[val_idx])
    main()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    checkpoint_path = './model/best_model.pt'
    checkpoint = torch.load(checkpoint_path, map_location=device)  
    model = AE(X.shape[1], 128).to(device)  
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    with torch.no_grad():
        Z = model.enc(X.to(device)).cpu().numpy()

    adata.obsm['X_ae128'] = Z
    adata.write('./data/MouseSC.h5ad', compression='gzip')