test.py

import scipy.io
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import math
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
import time
from torch.utils.data import TensorDataset, DataLoader
import random
import os
from model.Caputo_FractionalRNN import Caputo_FractionalRNN
from model.StandardRNN import StandardRNN
from model.LSTMModel import LSTMModel
from model.GRUModel import GRUModel
from model.TransformerModel import Transformer
from model.TransformerxlModel import TransformerXL
from model.InformerModel import InformerModel
from model.utils import count_parameters, set_seed, load_data_to_tensorloader
from model.utils import model_size_in_mb, print_model_info
from model.FractionalRNNFX1Model import FractionalRNN1
from model.FractionalRNNFX2Model import FractionalRNN2
from model.utils import train_model, evaluate_model
import pandas as pd

def main_test(
        model_choice = "Caputo_FractionalRNN", 
        save_name = 'results_MG',
        data_name = 'MG',
        in_len = 7,
        out_len = 0, 
        in_dim = 1 , 
        out_dim = 1
):
    set_seed(42)
    hidden_sizes = [64,]
    alpha = 0.5
    memory_length = 30
    input_history_length = in_len 
    output_history_length = out_len 
    input_dim = in_dim 
    output_dim = out_dim
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    seq_len = input_history_length + output_history_length
    _, X_test, y_test, output_scaler = load_data_to_tensorloader(f"./data/{data_name}.mat",data_name=data_name, batch_size=128, input_len=input_history_length, output_len= output_history_length)

    model_choice = model_choice
    if model_choice == "Caputo_FractionalRNN":
        model = Caputo_FractionalRNN(input_dim, hidden_sizes, output_dim, alpha, memory_length)
    elif model_choice == "StandardRNN":
        model = StandardRNN(input_dim, hidden_sizes, output_dim)
    elif model_choice == "LSTM":
        model = LSTMModel(input_dim, hidden_sizes, output_dim)
    elif model_choice == "GRU":
        model = GRUModel(input_dim, hidden_sizes, output_dim)
    elif model_choice == "FX1":
        model = FractionalRNN1(input_dim, hidden_sizes, output_dim, alpha=0.8)
    elif model_choice == "FX2":
        model = FractionalRNN2(input_dim, hidden_sizes, output_dim, gamma=0.3)

    elif model_choice == "Transformer":
        model = Transformer(
            input_dim=input_dim,
            d_model=64,
            num_layers=2,
            num_heads=4,
            output_dim=output_dim
        )
    elif model_choice == "TransformerXL":
        model = TransformerXL(
            input_dim=input_dim,
            output_dim=output_dim,
            d_model=64,
            n_heads=4,
            n_layers=4,
            d_ff=256,
            mem_len=20
        )
    elif model_choice == "Informer":
        model = InformerModel(
            input_dim=input_dim,
            output_dim=output_dim,
            seq_len=seq_len,
            pred_len=1,
            d_model=64,
            n_heads=4,
            e_layers=2,
            d_layers=1,
            d_ff=256,
        )
    else:
        raise ValueError(f"未知模型类型: {model_choice}")
    if model_choice == 'Caputo_FractionalRNN':
        model.load_state_dict(torch.load(f"./{save_name}/best_models/best_Caputo.pth"))
    else:
        model.load_state_dict(torch.load(f"./{save_name}/best_models/{model_choice}.pth"))
    model = model.to(device)
    model.eval()
    X_test = X_test.to(device)
    y_test = y_test.to(device)
    
    y_pred_original,y_test_original, mse, rmse, mae, mape = evaluate_model(model, model_choice, X_test, y_test,device=device, output_scaler=output_scaler,save_name=save_name)

    csv_path = f"./{save_name}/predictions/{model_choice}_prediction.csv"
    num_outputs = y_test_original.shape[1]
    data_dict = {}
    for i in range(num_outputs):
        data_dict[f"True_{i+1}"] = y_test_original[:, i]
        data_dict[f"Pred_{i+1}"] = y_pred_original[:, i]
    df = pd.DataFrame(data_dict)
    df.to_csv(csv_path, index=False)

    txt_path = f"./{save_name}/metrics/{model_choice}_metrics.txt"

    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(f"Model: {model_choice}\n")
        f.write(f"MSE:  {mse:.15f}\n")
        f.write(f"RMSE: {rmse:.15f}\n")
        f.write(f"MAE:  {mae:.15f}\n")
        f.write(f"MAPE: {mape:.15f}%\n")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_choice", type=str, required=True)
    parser.add_argument("--save_name", type=str, default="results")
    parser.add_argument("--data_name", type=str, default="MG")
    parser.add_argument("--in_len", type=int, default=5)
    parser.add_argument("--out_len", type=int, default=0)
    parser.add_argument("--in_dim", type=int, default=1)
    parser.add_argument("--out_dim", type=int, default=1)

    args = parser.parse_args()

    main_test(
        model_choice=args.model_choice,
        save_name=args.save_name,
        data_name=args.data_name,
        in_len=args.in_len,
        out_len=args.out_len,
        in_dim=args.in_dim,
        out_dim=args.out_dim
    )