Function main

tasks/AutoTPPR/code/experiment.py:1375–1416

(data_path='./data', out_dir='./saved_models', device='cuda:0')

def main(data_path='./data', out_dir='./saved_models', device='cuda:0'):
    os.makedirs(data_path, exist_ok=True)
    os.makedirs(out_dir, exist_ok=True)

    os.environ["WANDB_SILENT"] = "true" 
    os.environ["WANDB_ERROR_REPORTING"] = "false"

    print_sys("=== data loading ===")
    pert_data = PertData(data_path)
    
    pert_data.load(data_name='norman')
    
    pert_data.prepare_split(split='simulation', seed=1)
    pert_data.get_dataloader(batch_size=32, test_batch_size=128)

    print_sys("\n=== model traing ===")
    gears_model = GEARS(
        pert_data,
        device=device,
        weight_bias_track=True,
        proj_name='GEARS',
        exp_name='gears_norman'
    )
    gears_model.model_initialize(hidden_size = 64)
    
    gears_model.train(epochs=args.epochs, lr=1e-3)
    
    gears_model.save_model(os.path.join(out_dir, 'norman_full_model'))
    print_sys(f"model saved to {out_dir}")
    gears_model.load_pretrained(os.path.join(out_dir, 'norman_full_model'))

    final_infos = {
            "Gears":{
                "means":{
                    "Test Top 20 DE MSE": float(gears_model.test_metrics['mse_de'].item())
                }
            }
        }
    
    with open(os.path.join(out_dir, 'final_info.json'), 'w') as f:
        json.dump(final_infos, f, indent=4)
    print_sys("final info saved.")
    
def print_sys(s):
    """system print