export_models.py

import torch
import os

# Import our custom architectures from the training scripts
from train_cnn import AetherisResNet
from train_lstm import AetherisPredictor

def export_models():
    print("[Export] Initializing model exportation to LibTorch...")
    
    # ---------------------------------------------------------
    # 1. Export the Spatial CNN (AetherisResNet)
    # ---------------------------------------------------------
    cnn_weights_path = 'aetheris_resnet18_best.pth'
    if not os.path.exists(cnn_weights_path):
        print(f"[Error] CNN weights not found at {cnn_weights_path}")
    else:
        print(f"[Export] Loading CNN weights from {cnn_weights_path}...")
        # Instantiate model architecture
        cnn_model = AetherisResNet(num_classes=50)
        # Load weights (map to CPU to ensure safe tracing regardless of training device)
        cnn_model.load_state_dict(torch.load(cnn_weights_path, map_location='cpu'))
        cnn_model.eval()
        
        # Create a dummy input tensor matching the C++ inference input
        # Shape: (batch_size=1, channels=1, height=160, width=160)
        dummy_cnn_input = torch.randn(1, 1, 160, 160)
        
        print("[Export] Tracing CNN with TorchScript...")
        traced_cnn = torch.jit.trace(cnn_model, dummy_cnn_input)
        
        cnn_out_path = 'cnn_traced.pt'
        traced_cnn.save(cnn_out_path)
        print(f"[Export] Successfully saved traced CNN to {cnn_out_path}")
        
    # ---------------------------------------------------------
    # 2. Export the Temporal LSTM (AetherisPredictor)
    # ---------------------------------------------------------
    lstm_weights_path = 'aetheris_lstm_best.pth'
    if not os.path.exists(lstm_weights_path):
        print(f"[Error] LSTM weights not found at {lstm_weights_path}")
    else:
        print(f"\n[Export] Loading LSTM weights from {lstm_weights_path}...")
        # Instantiate model architecture
        lstm_model = AetherisPredictor(input_size=50, hidden_size=256, num_layers=2)
        lstm_model.load_state_dict(torch.load(lstm_weights_path, map_location='cpu'))
        lstm_model.eval()
        
        # Create a dummy input tensor for the sequence
        # Shape: (batch_size=1, seq_len=5, features=50)
        dummy_lstm_input = torch.randn(1, 5, 50)
        
        print("[Export] Tracing LSTM with TorchScript...")
        traced_lstm = torch.jit.trace(lstm_model, dummy_lstm_input)
        
        lstm_out_path = 'lstm_traced.pt'
        traced_lstm.save(lstm_out_path)
        print(f"[Export] Successfully saved traced LSTM to {lstm_out_path}")
        
    print("\n[DONE] Export pipeline complete. Models are ready for C++ deployment.")

if __name__ == '__main__':
    export_models()