Add compressor_with_embeddings.py

src/compressor_with_embeddings.py

@@ -0,0 +1,278 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
+import json
+import numpy as np
+from tqdm import tqdm
+
+# ---------------- Hyperparameters ----------------
+ESM_DIM      = 1280  # ESM-2 hidden dim (esm2_t33_650M_UR50D)
+COMP_RATIO   = 16    # compression factor
+COMP_DIM     = ESM_DIM // COMP_RATIO
+MAX_SEQ_LEN  = 50    # Actual sequence length from final_sequence_encoder.py
+BATCH_SIZE   = 32
+EPOCHS       = 30
+BASE_LR       = 1e-3  # initial learning rate
+LR_MIN        = 8e-5  # minimum learning rate for cosine schedule
+WARMUP_STEPS = 10_000
+DEPTH        = 4     # total transformer layers (2 pre-pool, 2 post-pool)
+HEADS        = 8     # attention heads
+DIM_FF       = ESM_DIM * 4
+POOLING      = True  # enforce ProtFlow hourglass pooling
+
+# ---------------- Dataset for Pre-computed Embeddings ----------------
+class PrecomputedEmbeddingDataset(Dataset):
+    def __init__(self, embeddings_path):
+        """
+        Load pre-computed embeddings from the final_sequence_encoder.py output.
+        Args:
+            embeddings_path: Path to the directory containing individual .pt embedding files
+        """
+        print(f"Loading pre-computed embeddings from {embeddings_path}...")
+        
+        # Load all individual embedding files
+        import glob
+        import os
+        
+        embedding_files = glob.glob(os.path.join(embeddings_path, "*.pt"))
+        embedding_files = [f for f in embedding_files if not f.endswith('metadata.json') and not f.endswith('sequence_ids.json')]
+        
+        print(f"Found {len(embedding_files)} embedding files")
+        
+        # Load and stack all embeddings
+        embeddings_list = []
+        for file_path in embedding_files:
+            try:
+                embedding = torch.load(file_path)
+                if embedding.dim() == 2:  # (seq_len, hidden_dim)
+                    embeddings_list.append(embedding)
+                else:
+                    print(f"Warning: Skipping {file_path} - unexpected shape {embedding.shape}")
+            except Exception as e:
+                print(f"Warning: Could not load {file_path}: {e}")
+        
+        if not embeddings_list:
+            raise ValueError("No valid embeddings found!")
+        
+        self.embeddings = torch.stack(embeddings_list)
+        print(f"Loaded {len(self.embeddings)} embeddings with shape {self.embeddings.shape}")
+        
+        # Ensure embeddings are the right shape
+        if len(self.embeddings.shape) != 3:
+            raise ValueError(f"Expected 3D tensor, got shape {self.embeddings.shape}")
+        
+        if self.embeddings.shape[1] != MAX_SEQ_LEN:
+            print(f"Warning: Expected sequence length {MAX_SEQ_LEN}, got {self.embeddings.shape[1]}")
+        
+        if self.embeddings.shape[2] != ESM_DIM:
+            print(f"Warning: Expected embedding dim {ESM_DIM}, got {self.embeddings.shape[2]}")
+
+    def __len__(self): 
+        return len(self.embeddings)
+    
+    def __getitem__(self, idx): 
+        return self.embeddings[idx]
+
+# ---------------- Compressor ----------------
+class Compressor(nn.Module):
+    def __init__(self, in_dim=ESM_DIM, out_dim=COMP_DIM):
+        super().__init__()
+        self.norm = nn.LayerNorm(in_dim)
+        layer = lambda: nn.TransformerEncoderLayer(
+            d_model=in_dim, nhead=HEADS, dim_feedforward=DIM_FF,
+            batch_first=True)
+        # two layers before pool, two after
+        self.pre_tr  = nn.TransformerEncoder(layer(), num_layers=DEPTH//2)
+        self.post_tr = nn.TransformerEncoder(layer(), num_layers=DEPTH//2)
+        self.proj    = nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, out_dim),
+            nn.Tanh()
+        )
+        self.pooling = POOLING
+
+    def forward(self, x, stats=None):
+        if stats:
+            m, s, mn, mx = stats['mean'], stats['std'], stats['min'], stats['max']
+            # Move stats to the same device as x
+            m = m.to(x.device)
+            s = s.to(x.device)
+            mn = mn.to(x.device)
+            mx = mx.to(x.device)
+            x = torch.clamp((x - m) / s, -4, 4)
+            x = torch.clamp((x - mn) / (mx - mn + 1e-8), 0, 1)
+        x = self.norm(x)
+        x = self.pre_tr(x)                 # [B, L, D]
+        if self.pooling:
+            B, L, D = x.shape
+            if L % 2: x = x[:, :-1, :]
+            x = x.view(B, L//2, 2, D).mean(2)  # halve sequence length
+        x = self.post_tr(x)                # [B, L' , D]
+        return self.proj(x)                # [B, L', COMP_DIM]
+
+# ---------------- Decompressor ----------------
+class Decompressor(nn.Module):
+    def __init__(self, in_dim=COMP_DIM, out_dim=ESM_DIM):
+        super().__init__()
+        self.proj    = nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, out_dim)
+        )
+        layer = lambda: nn.TransformerEncoderLayer(
+            d_model=out_dim, nhead=HEADS, dim_feedforward=DIM_FF,
+            batch_first=True)
+        self.decoder = nn.TransformerEncoder(layer(), num_layers=DEPTH//2)
+        self.pooling = POOLING
+
+    def forward(self, z):
+        x = self.proj(z)                   # [B, L', D]
+        if self.pooling:
+            x = x.repeat_interleave(2, dim=1)  # unpool to full length
+        return self.decoder(x)             # [B, L, out_dim]
+
+# ---------------- Training Loop ----------------
+def train_with_precomputed_embeddings(embeddings_path, device='cuda'):
+    """
+    Train compressor using pre-computed embeddings from final_sequence_encoder.py
+    """
+    # Load dataset
+    ds = PrecomputedEmbeddingDataset(embeddings_path)
+    
+    # Compute normalization statistics
+    print("Computing normalization statistics...")
+    flat = ds.embeddings.view(-1, ESM_DIM)
+    stats = {
+        'mean': flat.mean(0),
+        'std':  flat.std(0) + 1e-8,
+        'min':  torch.clamp((flat - flat.mean(0)) / (flat.std(0) + 1e-8), -4,4).min(0)[0],
+        'max':  torch.clamp((flat - flat.mean(0)) / (flat.std(0) + 1e-8), -4,4).max(0)[0]
+    }
+    
+    # Save statistics for later use
+    torch.save(stats, 'normalization_stats.pt')
+    print("Saved normalization statistics to normalization_stats.pt")
+    
+    # Create data loader
+    dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True)
+    
+    # Initialize models
+    comp = Compressor().to(device)
+    decomp = Decompressor().to(device)
+    
+    # Initialize optimizer
+    opt = optim.AdamW(list(comp.parameters()) + list(decomp.parameters()), lr=BASE_LR)
+    
+    # LR scheduling: warmup -> cosine
+    warmup_sched = LinearLR(opt, start_factor=1e-8, end_factor=1.0, total_iters=WARMUP_STEPS)
+    cosine_sched = CosineAnnealingLR(opt, T_max=EPOCHS*len(dl), eta_min=LR_MIN)
+    sched = SequentialLR(opt, [warmup_sched, cosine_sched], milestones=[WARMUP_STEPS])
+
+    print(f"Starting training for {EPOCHS} epochs...")
+    print(f"Device: {device}")
+    print(f"Batch size: {BATCH_SIZE}")
+    print(f"Total batches per epoch: {len(dl)}")
+
+    # Training loop
+    for epoch in range(1, EPOCHS+1):
+        total_loss = 0
+        comp.train()
+        decomp.train()
+        
+        for batch_idx, x in enumerate(tqdm(dl, desc=f"Epoch {epoch}/{EPOCHS}")):
+            x = x.to(device)
+            z = comp(x, stats)
+            xr = decomp(z)
+            loss = (x - xr).pow(2).mean()
+            
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            sched.step()
+            
+            total_loss += loss.item()
+            
+            # Print progress every 100 batches
+            if batch_idx % 100 == 0:
+                print(f"  Batch {batch_idx}/{len(dl)} - Loss: {loss.item():.6f}")
+        
+        avg_loss = total_loss / len(dl)
+        print(f"Epoch {epoch}/{EPOCHS} — Average MSE: {avg_loss:.6f}")
+        
+        # Save checkpoint every 5 epochs
+        if epoch % 5 == 0:
+            torch.save({
+                'epoch': epoch,
+                'compressor_state_dict': comp.state_dict(),
+                'decompressor_state_dict': decomp.state_dict(),
+                'optimizer_state_dict': opt.state_dict(),
+                'loss': avg_loss,
+            }, f'checkpoint_epoch_{epoch}.pth')
+
+    # Save final models
+    torch.save(comp.state_dict(), 'compressor_final.pth')
+    torch.save(decomp.state_dict(), 'decompressor_final.pth')
+    print("Training completed! Models saved as compressor_final.pth and decompressor_final.pth")
+
+# ---------------- Utility Functions ----------------
+def load_and_test_models(compressor_path, decompressor_path, embeddings_path, device='cuda'):
+    """
+    Load trained models and test reconstruction quality
+    """
+    print("Loading trained models...")
+    comp = Compressor().to(device)
+    decomp = Decompressor().to(device)
+    
+    comp.load_state_dict(torch.load(compressor_path))
+    decomp.load_state_dict(torch.load(decompressor_path))
+    
+    comp.eval()
+    decomp.eval()
+    
+    # Load test data
+    ds = PrecomputedEmbeddingDataset(embeddings_path)
+    test_loader = DataLoader(ds, batch_size=16, shuffle=False)
+    
+    # Load normalization stats
+    stats = torch.load('normalization_stats.pt')
+    
+    print("Testing reconstruction quality...")
+    total_mse = 0
+    total_samples = 0
+    
+    with torch.no_grad():
+        for batch in tqdm(test_loader, desc="Testing"):
+            x = batch.to(device)
+            z = comp(x, stats)
+            xr = decomp(z)
+            mse = (x - xr).pow(2).mean()
+            total_mse += mse.item() * len(x)
+            total_samples += len(x)
+    
+    avg_mse = total_mse / total_samples
+    print(f"Average reconstruction MSE: {avg_mse:.6f}")
+    
+    return avg_mse
+
+# ---------------- Entrypoint ----------------
+if __name__ == '__main__':
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Train protein compressor with pre-computed embeddings')
+    parser.add_argument('--embeddings', type=str, default='/data2/edwardsun/flow_project/compressor_dataset/peptide_embeddings.pt',
+                       help='Path to pre-computed embeddings from final_sequence_encoder.py')
+    parser.add_argument('--device', type=str, default='cuda', help='Device to use (cuda/cpu)')
+    parser.add_argument('--test', action='store_true', help='Test existing models instead of training')
+    
+    args = parser.parse_args()
+    
+    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    if args.test:
+        # Test existing models
+        load_and_test_models('compressor_final.pth', 'decompressor_final.pth', args.embeddings, device)
+    else:
+        # Train new models
+        train_with_precomputed_embeddings(args.embeddings, device)