train.py

import torch
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sentence_transformers import SentenceTransformer

from hypernet import HyperNet
from loss import OrthogonalSVDLoss

class LoRADataset(Dataset):
    def __init__(self, data_path, r=8, device="cpu"):
        self.data = torch.load(data_path, map_location="cpu", weights_only=False)
        self.r = r
        
        # Phase 2: The Conditioning Vector (z)
        print("Encoding task descriptions using frozen SentenceTransformer...")
        encoder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2").to(device)
        encoder.eval()
        descriptions = [item["task_description_string"] for item in self.data]
        with torch.no_grad():
            self.z_vectors = encoder.encode(descriptions, convert_to_tensor=True).cpu()
            
        print(f"Loaded {len(self.data)} training examples.")
        
        # Identify target shapes dynamically from the first item
        self.target_shapes = {}
        sub_dict = self.data[0]["tensors"]
        for name in sub_dict["U_dict"].keys():
            d_out = sub_dict["U_dict"][name].shape[0]
            # V represents V^T here mathematically, so its shape is (r, d_in)
            d_in = sub_dict["V_dict"][name].shape[1]
            self.target_shapes[name] = (d_out, d_in)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        return {
            "z": self.z_vectors[idx],
            "U_dict": item["tensors"]["U_dict"],
            "log_Sigma_dict": item["tensors"]["log_Sigma_dict"],
            "V_dict": item["tensors"]["V_dict"]
        }

def collate_fn(batch):
    z = torch.stack([x["z"] for x in batch])
    
    first_item = batch[0]
    keys = first_item["U_dict"].keys()
    
    U_dict, log_Sigma_dict, V_dict = {}, {}, {}
    for k in keys:
        U_dict[k] = torch.stack([x["U_dict"][k] for x in batch])
        log_Sigma_dict[k] = torch.stack([x["log_Sigma_dict"][k] for x in batch])
        V_dict[k] = torch.stack([x["V_dict"][k] for x in batch])
        
    return z, U_dict, log_Sigma_dict, V_dict

def set_seed(seed=42):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def main():
    set_seed()
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # Configuration / Hyperparameters
    data_path = "lora_svd_dataset.pt"
    r = 8
    hidden_dim = 2048
    epochs = 100
    batch_size = 4
    lr = 1e-4

    # The dataset might not exist if the script hasn't been run yet.
    import os
    if not os.path.exists(data_path):
        print(f"Dataset {data_path} not found. Please run extract_svd.py first.")
        return

    dataset = LoRADataset(data_path, r=r, device=device)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    
    target_shapes = dataset.target_shapes
    
    # Phase 3: Instantiate Hypernetwork
    model = HyperNet(target_shapes, hidden_dim=hidden_dim, r=r).to(device)
    
    # Phase 4: Instantiate Geometrically Constrained Loss Function
    criterion = OrthogonalSVDLoss(lambda_1=0.1, lambda_2=0.1, r=r)
    
    optimizer = optim.AdamW(model.parameters(), lr=lr)
    
    print("Starting Training Loop...")
    model.train()
    for epoch in range(1, epochs + 1):
        total_loss, total_mse, total_ortho = 0.0, 0.0, 0.0
        
        for batch_z, batch_U, batch_log_Sigma, batch_V in dataloader:
            batch_z = batch_z.to(device)
            
            U_gen_dict, log_Sigma_gen_dict, V_gen_dict = model(batch_z)
            
            loss, l_mse, l_ortho = criterion(
                U_gen_dict, log_Sigma_gen_dict, V_gen_dict,
                batch_U, batch_log_Sigma, batch_V
            )
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            total_mse += l_mse.item()
            total_ortho += l_ortho.item()
            
        avg_loss = total_loss / len(dataloader)
        avg_mse = total_mse / len(dataloader)
        avg_ortho = total_ortho / len(dataloader)
        
        print(f"Epoch {epoch:03d}/{epochs} | Loss: {avg_loss:.4f} "
              f"(MSE: {avg_mse:.4f}, Ortho: {avg_ortho:.4f})")
              
    # Phase 5 context: Save the trained Hypernetwork for inference
    torch.save(model.state_dict(), "hypernet_final.pt")
    print("Training complete. Model saved to hypernet_final.pt")

if __name__ == "__main__":
    main()