65 lines
3.4 KiB
Python
65 lines
3.4 KiB
Python
import argparse
|
||
import os
|
||
import torch
|
||
import pickle
|
||
|
||
from cross_exp.exp_crossformer import Exp_crossformer
|
||
from utils.tools import load_args, string_split
|
||
|
||
parser = argparse.ArgumentParser(description='CrossFormer')
|
||
|
||
parser.add_argument('--checkpoint_root', type=str, default='./checkpoints', help='location of the trained model')
|
||
parser.add_argument('--setting_name', type=str, default='Crossformer_ETTh1_il168_ol24_sl6_win2_fa10_dm256_nh4_el3_itr0', help='name of the experiment')
|
||
|
||
parser.add_argument('--num_workers', type=int, default=0, help='data loader num workers')
|
||
parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')
|
||
|
||
parser.add_argument('--different_split', action='store_true', help='use data split different from training process', default=False)
|
||
parser.add_argument('--data_split', type=str, default='0.7,0.1,0.2', help='data split of train, vali, test')
|
||
|
||
parser.add_argument('--inverse', action='store_true', help='inverse output data into the original scale', default=True)
|
||
parser.add_argument('--save_pred', action='store_true', help='whether to save the predicted future MTS', default=True)
|
||
|
||
parser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')
|
||
parser.add_argument('--gpu', type=int, default=0, help='gpu')
|
||
|
||
args = parser.parse_args()
|
||
|
||
args.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False
|
||
args.use_multi_gpu = False
|
||
|
||
args.checkpoint_dir = os.path.join(args.checkpoint_root, args.setting_name)
|
||
hyper_parameters = load_args(os.path.join(args.checkpoint_dir, 'args.json'))
|
||
|
||
#load the pre-trained model
|
||
args.data_dim = hyper_parameters['data_dim']; args.in_len = hyper_parameters['in_len']; args.out_len = hyper_parameters['out_len'];
|
||
args.seg_len = hyper_parameters['seg_len']; args.win_size = hyper_parameters['win_size']; args.factor = hyper_parameters['factor'];
|
||
args.d_model = hyper_parameters['d_model']; args.d_ff = hyper_parameters['d_ff']; args.n_heads = hyper_parameters['n_heads'];
|
||
args.e_layers = hyper_parameters['e_layers']; args.dropout = hyper_parameters['dropout']; args.baseline = hyper_parameters['baseline'];
|
||
exp = Exp_crossformer(args)
|
||
model_dict = torch.load(os.path.join(args.checkpoint_dir, 'checkpoint.pth'), map_location='cpu') # 加载的内容通常是一个字典，包含模型的权重（state_dict）以及可能的其他信息（如优化器状态、训练轮数等）。
|
||
exp.model.load_state_dict(model_dict)
|
||
|
||
#load the data
|
||
args.scale_statistic = pickle.load(open(os.path.join(args.checkpoint_dir, 'scale_statistic.pkl'), 'rb'))
|
||
args.root_path = hyper_parameters['root_path']; args.data_path = hyper_parameters['data_path'];
|
||
if args.different_split:
|
||
data_split = string_split(args.data_split)
|
||
args.data_split = data_split
|
||
else:
|
||
args.data_split = hyper_parameters['data_split']
|
||
|
||
mae, mse, rmse, mape, mspe = exp.eval(args.setting_name, args.save_pred, args.inverse)
|
||
|
||
folder_path = './results/' + args.setting_name +'/'
|
||
if not os.path.exists(folder_path):
|
||
os.makedirs(folder_path)
|
||
log_file = open(folder_path+'metric.log', 'w')
|
||
log_file.write('Data Path: {}\n'.format(os.path.join(args.root_path, args.data_path)))
|
||
log_file.write('Data Split: {}\n'.format(args.data_split))
|
||
log_file.write('Input Length:{} Output Length:{}\n'.format(args.in_len, args.out_len))
|
||
log_file.write('Inverse to original scale: {}\n\n'.format(args.inverse))
|
||
log_file.write('MAE:{}\nMSE:{}\nRMSE:{}\nMAPE:{}\nMSPE:{}\n'.format(mae, mse, rmse, mape, mspe))
|
||
log_file.close()
|
||
|