Upload New File

INR/regression_coral.py

+import os
+import sys
+from pathlib import Path
+import matplotlib.pyplot as plt
+from sklearn.metrics import mean_squared_error
+sys.path.append(str(Path(__file__).parents[1]))
+import time
+import hydra
+import numpy as np
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+from coral.utils.models.load_inr import create_inr_instance
+from dataset import *
+
+
+# Function to load a model
+def load_model(model_weights, cfg,input_dim, output_dim):
+    # ... load the INR model from model_path using the configuration cfg
+    # load inr weights
+    
+    inr_in = create_inr_instance(
+        cfg, input_dim=input_dim, output_dim=output_dim, device="cuda"
+    )
+    inr_in.load_state_dict(model_weights)
+    inr_in.eval()
+
+    return inr_in
+
+
+
+
+@hydra.main(config_path="", config_name="regression.yaml")
+def main(cfg: DictConfig) -> None:
+    inr_cfg = cfg.inr_cfg
+
+    target_field = inr_cfg.target_field
+
+    print(f"Processing target field: {target_field}")
+
+    dir = os.path.dirname(__file__)
+    input_inr = torch.load(os.path.join(dir,inr_cfg.model_path))
+
+    cfg_single = input_inr['cfg']
+    num_points = None 
+    add_cond_in = False if cfg_single.inr_in.add_cond_in is None else cfg_single.inr_in.add_cond_in
+
+    # Create dataset instance
+    dataset = TransonicRAE(cfg.data_directory, target_field, add_cond_in=add_cond_in, num_points=num_points)
+    dataset.create_splits(train_ratio=0.9, val_ratio=0.05, test_ratio=0.05, seed=42)
+
+    ntest = len(dataset.test_dataset)
+
+     # Improved folder naming with timestamp and all output target fields
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    folder_name = f"training_resilt_{timestamp}" 
+    results_directory = os.path.join(dir + '/trainings', folder_name)
+    os.makedirs(results_directory, exist_ok=True)
+    cfg_save_path = os.path.join(results_directory, 'config.yaml')
+    torch.save(cfg, cfg_save_path)
+
+    # Predict outputs for test dataset
+    result_test = {'predictions': {}, 'targets' : {}, 'mse' : {}, 'cond' : {}}
+    cfg_single = torch.load(os.path.join(dir,inr_cfg.model_path))['cfg']
+    
+    if cfg_single.inr_in.global_norm:
+        min_val, max_val = dataset.coef_norm['min'],dataset.coef_norm['max']
+    else:
+         mean_val, std_val = dataset.coef_norm['mean'],dataset.coef_norm['std']
+        
+    inr_model = load_model(torch.load(os.path.join(dir,inr_cfg.model_path))["inr_in"], cfg_single, inr_cfg.input_dim, inr_cfg.output_dim)
+    inr_model.eval()
+    result_test['predictions'][inr_cfg.target_field] = []
+    result_test['targets'][inr_cfg.target_field] = []
+    result_test['mse'][inr_cfg.target_field] = []
+    result_test['cond'][inr_cfg.target_field] = []
+
+    for substep, graph in enumerate(dataset.test_dataset):
+        if cfg_single.inr_in.global_norm:
+            pred = inr_model.modulated_forward(graph.input.cuda(), graph.cond.cuda()).detach().cpu() * (max_val - min_val) + min_val
+            target = graph.output * (max_val - min_val) + min_val
+        else:
+            pred =inr_model.modulated_forward(graph.input.cuda(), graph.cond.cuda()).detach().cpu() * (std_val) + mean_val
+            target = graph.output * (std_val) + mean_val
+
+        mse = ((pred.numpy() -target.numpy())**2).mean(axis = 0)
+       
+        result_test['predictions'][inr_cfg.target_field].append(pred.numpy())
+        result_test['targets'][inr_cfg.target_field].append(target.numpy())
+        result_test['mse'][inr_cfg.target_field].append(mse)
+        result_test['cond'][inr_cfg.target_field].append(graph.cond.numpy())
+
+
+    # Save the outputs           
+    outputs_save_path = os.path.join(results_directory, 'result_test.pt')
+    torch.save(result_test, outputs_save_path)
+
+    print("Outputs predicted and saved at:", outputs_save_path)
+    print('Mean mse total:', sum(result_test['mse'][inr_cfg.target_field])/len(result_test['mse'][inr_cfg.target_field]))
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file