├── .gitignore ├── README.md ├── config ├── attention_lstm_params.json ├── cnn_lstm_params.json ├── cnn_params.json ├── gru_params.json ├── lstm_params.json ├── rnn_params.json └── seq2seq_params.json ├── main.py ├── models ├── AttentionLSTM │ └── attention_lstm.py ├── BiLSTM │ └── bilstm.py ├── CNN │ └── cnn.py ├── CNN_LSTM │ └── cnn_lstm.py ├── GRU │ └── gru.py ├── LSTM │ └── lstm.py ├── RNN │ └── rnn.py └── SEQ2SEQ │ └── seq2seq.py └── trainPEMS.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.pyc 3 | *.pyo 4 | *.pyd 5 | *.pyw 6 | *.pyz 7 | data/* 8 | img/* 9 | save_model/* -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # BRTdata_predict 2 | RNN、LSTM、CNN+LSTM、GRU 3 | -------------------------------------------------------------------------------- /config/attention_lstm_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "num_layers": 1, 6 | "bidirectional": "False", 7 | "lr": 0.001, 8 | "epochs": 300 9 | } -------------------------------------------------------------------------------- /config/cnn_lstm_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "lr": 0.001, 5 | "epochs": 150 6 | } -------------------------------------------------------------------------------- /config/cnn_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "lr": 0.001, 6 | "epochs": 150 7 | } -------------------------------------------------------------------------------- /config/gru_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "num_layers": 1, 6 | "bidirectional": "False", 7 | "lr": 0.001, 8 | "epochs": 150 9 | } -------------------------------------------------------------------------------- /config/lstm_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "num_layers": 1, 6 | "bidirectional": "True", 7 | "lr": 0.005, 8 | "epochs": 250 9 | } -------------------------------------------------------------------------------- /config/rnn_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "num_layers": 1, 6 | "bidirectional": "False", 7 | "lr": 0.001, 8 | "epochs": 150 9 | } -------------------------------------------------------------------------------- /config/seq2seq_params.json: -------------------------------------------------------------------------------- 1 | { 2 | "input_size": 3, 3 | "output_size": 1, 4 | "hidden_size": 64, 5 | "num_layers": 2, 6 | "lr": 0.002, 7 | "epochs": 500 8 | } -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from models.LSTM.lstm import LSTMmodel 5 | from models.CNN_LSTM.cnn_lstm import CNN_LSTM_Model 6 | from models.RNN.rnn import RNNmodel 7 | from models.CNN.cnn import CNNmodel 8 | from models.GRU.gru import GRUmodel 9 | from models.SEQ2SEQ.seq2seq import Encoder, Decoder, Seq2Seq 10 | from models.AttentionLSTM.attention_lstm import AttentionLSTMmodel 11 | import matplotlib.pyplot as plt 12 | import seaborn as sns 13 | from matplotlib.font_manager import FontProperties 14 | import os 15 | import json 16 | import time 17 | import pandas as pd 18 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' 19 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 20 | X = None 21 | Y = None 22 | testX = None 23 | testY = None 24 | lr = 0.0001 25 | epochs = 100 26 | model = None 27 | min_test_data = None 28 | max_test_data = None # 测试集的最大最小值 29 | data_path = "./data/BRT/厦门北站_brtdata.npz" # 数据集路径 30 | data_name = "厦门北站" # 站点数据集名称 31 | save_mod_dir = "./save_model" # 模型保存路径 32 | 33 | '''load_data''' 34 | def load_data(cut_point, data_path, input_length=7*24, output_length=1*24, days=61, hours=24): 35 | global min_test_data, max_test_data 36 | ''' 37 | INPUT: 38 | cut_point, 训练和测试数据集划分点 39 | days, 客流数据的总天数 40 | hours, 每天小时数 41 | input_length, 模型输入的过去7天的数据作为一个序列 42 | output_length, 预测未来1天的数据 43 | OUTPUT: 44 | train_data 45 | test_data 46 | ''' 47 | data = np.load(data_path)['data'] 48 | data = data.astype(np.float32) 49 | data = data[:, np.newaxis] 50 | 51 | seq_days = np.arange(days) 52 | seq_hours = np.arange(hours) 53 | seq_day_hour = np.transpose( 54 | [np.repeat(seq_days, len(seq_hours)), 55 | np.tile(seq_hours, len(seq_days))] 56 | ) # Cartesian Product 57 | 58 | # 按照列方向拼接 --> [客流量, days, hours] 59 | data = np.concatenate((data, seq_day_hour), axis=1) 60 | 61 | train_data = data[:cut_point] 62 | test_data = data[cut_point:] 63 | 64 | # 分别进行min-max scaling 65 | train_data, min_train_data, max_train_data = min_max_normalise_numpy(train_data) 66 | test_data, min_test_data, max_test_data = min_max_normalise_numpy(test_data) 67 | 68 | ''' 69 | input_data 模型输入 连续 input_length // 24 = 7天的数据 70 | output_data 模型输出 未来 output_length // 24 = 1天的数据 71 | ''' 72 | input_data = [] 73 | output_data = [] 74 | 75 | # 滑动窗口法 76 | for i in range(len(train_data) - input_length - output_length + 1): 77 | input_seq = train_data[i : i + input_length] 78 | output_seq = train_data[i + input_length : i + input_length + output_length, 0] 79 | input_data.append(input_seq) 80 | output_data.append(output_seq) 81 | 82 | # 转为torch.tensor 83 | X = torch.tensor(input_data, dtype=torch.float32, device=device) 84 | Y = torch.tensor(output_data, dtype=torch.float32, device=device) 85 | Y = Y.unsqueeze(-1) # 在最后一个维度（维度索引为1）上增加一个维度 86 | # X = torch.tensor([item.cpu().detach().numpy() for item in input_data], dtype=torch.float32, device=device) 87 | # Y = torch.tensor([item.cpu().detach().numpy() for item in output_data], dtype=torch.float32, device=device) 88 | 89 | test_inseq = [] 90 | test_outseq = [] 91 | # 滑动窗口法 92 | for i in range(len(test_data) - input_length - output_length + 1): 93 | input_seq = test_data[i : i + input_length] 94 | output_seq = test_data[i + input_length : i + input_length + output_length, 0] 95 | 96 | test_inseq.append(input_seq) 97 | test_outseq.append(output_seq) 98 | 99 | # 转为torch.tensor 100 | testX = torch.tensor(test_inseq, dtype=torch.float32, device=device) 101 | testY = torch.tensor(test_outseq, dtype=torch.float32, device=device) 102 | testY = testY.unsqueeze(-1) # 在最后一个维度（维度索引为1）上增加一个维度 103 | # testX = torch.tensor([item.cpu().detach().numpy() for item in test_inseq], dtype=torch.float32, device=device) 104 | # testY = torch.tensor([item.cpu().detach().numpy() for item in test_outseq], dtype=torch.float32, device=device) 105 | 106 | # 输出数据形状 107 | print("数据集处理完毕：") 108 | print("data - 原数据集 shape:", data.shape) 109 | print("traindata - Input shape:", X.shape) 110 | print("traindata - Output shape:", Y.shape) 111 | print("testdata - Input shape:", testX.shape) 112 | print("testdata - Output shape:", testY.shape) 113 | return X, Y, testX, testY 114 | 115 | '''min-max Scaling''' 116 | def min_max_normalise_numpy(x): 117 | # shape: [sequence_length, features] 118 | min_vals = np.min(x, axis=0) 119 | max_vals = np.max(x, axis=0) 120 | # [features] -> shape: [1, features] 121 | min_vals = np.expand_dims(min_vals, axis=0) 122 | max_vals = np.expand_dims(max_vals, axis=0) 123 | # 归一化 -> [-1, 1] 124 | normalized_data = 2 * (x - min_vals) / (max_vals - min_vals) - 1 125 | # 归一化 -> [0, 1] 126 | # normalized_data = (x - min_vals) / (max_vals - min_vals) 127 | return normalized_data, min_vals, max_vals 128 | 129 | def inverse_min_max_normalise_numpy(normalised_x, min_vals, max_vals): 130 | x = (normalised_x + 1) / 2 * (max_vals - min_vals) + min_vals 131 | return x 132 | 133 | def min_max_normalise_tensor(x): 134 | # shape: [samples, sequence_length, features] 135 | min_vals = x.min(dim=1).values.unsqueeze(1) 136 | max_vals = x.max(dim=1).values.unsqueeze(1) 137 | # data ->[-1, 1] 138 | normalise_x = 2 * (x - min_vals) / (max_vals - min_vals) - 1 139 | return normalise_x, min_vals, max_vals 140 | 141 | def inverse_min_max_normalise_tensor(x, min_vals, max_vals): 142 | min_vals = torch.tensor(min_vals).to(device) 143 | max_vals = torch.tensor(max_vals).to(device) 144 | # shape: [1, features] -> [1, 1, features] 145 | min_vals = min_vals.unsqueeze(0) 146 | max_vals = max_vals.repeat(x.shape[0], 1, 1) 147 | # [1, 1, features] -> [samples, 1, features] 148 | min_vals = min_vals.repeat(x.shape[0], 1, 1) 149 | max_vals = max_vals.repeat(x.shape[0], 1, 1) 150 | 151 | x = (x + 1) / 2 * (max_vals - min_vals) + min_vals 152 | return x 153 | 154 | # MAPE: 平均绝对百分比误差 155 | def MAPELoss(y_hat, y): 156 | x = torch.tensor(1e-6, dtype=torch.float32).to(device) 157 | y_new = torch.where(y==0, x, y) # 防止分母为0 158 | abs_error = torch.abs((y - y_hat) / y_new) 159 | mape = 100. * torch.mean(abs_error) 160 | return mape 161 | 162 | def cnn(params): 163 | global lr, epochs, model 164 | '''超参数加载''' 165 | input_size = params["input_size"] # 输入特征数 166 | hidden_size = params["hidden_size"] # CNN output_channels 167 | output_size = params["output_size"] # 输出特征数 168 | lr = params["lr"] # 学习率 169 | epochs = params["epochs"] # 训练轮数 170 | model = CNNmodel(input_size, hidden_size, output_size).to(device) 171 | 172 | def rnn(params): 173 | global lr, epochs, model 174 | '''超参数加载''' 175 | input_size = params["input_size"] # 输入特征数 176 | hidden_size = params["hidden_size"] # RNN隐藏层神经元数 177 | num_layers = params["num_layers"] # RNN层数 178 | output_size = params["output_size"] # 输出特征数 179 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 180 | lr = params["lr"] # 学习率 181 | epochs = params["epochs"] # 训练轮数 182 | model = RNNmodel(input_size, hidden_size, num_layers, output_size, bidirectional=bidirectional).to(device) 183 | 184 | def lstm(params): 185 | global lr, epochs, model 186 | '''超参数加载''' 187 | input_size = params["input_size"] # 输入特征数 188 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 189 | num_layers = params["num_layers"] # LSTM层数 190 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 191 | output_size = params["output_size"] # 输出特征数 192 | lr = params["lr"] # 学习率 193 | epochs = params["epochs"] # 训练轮数 194 | 195 | model = LSTMmodel(input_size, hidden_size, num_layers, output_size, bidirectional=bidirectional).to(device) 196 | 197 | def cnn_lstm(params): 198 | global lr, epochs, model 199 | '''超参数加载''' 200 | input_size = params["input_size"] # 输入特征数 201 | output_size = params["output_size"] # 输出特征数 202 | lr = params["lr"] # 学习率 203 | epochs = params["epochs"] # 训练轮数 204 | model = CNN_LSTM_Model(input_size=input_size, output_size=output_size).to(device) 205 | 206 | def gru(params): 207 | global lr, epochs, model 208 | '''超参数加载''' 209 | input_size = params["input_size"] # 输入特征数 210 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 211 | num_layers = params["num_layers"] # LSTM层数 212 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 213 | output_size = params["output_size"] # 输出特征数 214 | lr = params["lr"] # 学习率 215 | epochs = params["epochs"] # 训练轮数 216 | model = GRUmodel(input_size, hidden_size, num_layers, output_size, bidirectional=bidirectional).to(device) 217 | 218 | def attention_lstm(params): 219 | global lr, epochs, model 220 | '''超参数加载''' 221 | input_size = params["input_size"] # 输入特征数 222 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 223 | num_layers = params["num_layers"] # LSTM层数 224 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 225 | output_size = params["output_size"] # 输出特征数 226 | lr = params["lr"] # 学习率 227 | epochs = params["epochs"] # 训练轮数 228 | model = AttentionLSTMmodel(input_size, hidden_size, num_layers, output_size, bidirectional=bidirectional, device=device).to(device) 229 | 230 | def seq2seq(params): 231 | global lr, epochs, model 232 | '''超参数加载''' 233 | input_size = params["input_size"] # 输入特征数 234 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 235 | num_layers = params["num_layers"] # LSTM层数 236 | output_size = params["output_size"] # 输出特征数 237 | lr = params["lr"] # 学习率 238 | epochs = params["epochs"] # 训练轮数 239 | encoder = Encoder(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers).to(device) 240 | decoder = Decoder(output_size=output_size, hidden_size=hidden_size, num_layers=num_layers).to(device) 241 | model = Seq2Seq(encoder, decoder).to(device) 242 | 243 | def train(model_name, save_mod=False): 244 | loss_function = nn.MSELoss() 245 | optimizer = torch.optim.Adam(model.parameters(), lr=lr) 246 | loss_list = [] # 保存训练过程中loss数据 247 | # 开始计时 248 | start_time = time.time() # 记录模型训练开始时间 249 | 250 | # train model 251 | for epoch in range(epochs): 252 | optimizer.zero_grad() # 在每个epoch开始时清零梯度 253 | if(model_name=="SEQ2SEQ"): 254 | Y_hat = model(X,Y,1) 255 | else: 256 | Y_hat = model(X) 257 | # print('Y_hat.size: [{},{},{}]'.format(Y_hat.size(0), Y_hat.size(1), Y_hat.size(2))) 258 | # print('Y.size: [{},{},{}]'.format(Y.size(0), Y.size(1), Y.size(2))) 259 | loss = loss_function(Y_hat, Y) 260 | loss.backward() 261 | nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=4, norm_type=2) 262 | optimizer.step() 263 | print(f'Epoch [{epoch+1}/{epochs}], MSE-Loss: {loss.item()}') 264 | loss_list.append(loss.item()) 265 | end_time = time.time() # 记录模型训练结束时间 266 | total_time = end_time - start_time # 计算总耗时 267 | print(f"本次模型训练总耗时: {total_time} 秒，超越了全国99.2%的单片机，太棒啦！") 268 | if save_mod == True: 269 | # save model -> ./**/model_name_eopch{n}.pth 270 | torch.save(model.state_dict(), '{}/{}_epoch{}.pth'.format(save_mod_dir, model_name, epochs)) 271 | print("Trained Model have saved in:", '{}/{}_epoch{}.pth'.format(save_mod_dir, model_name, epochs)) 272 | 273 | # save training loss graph 274 | epoch_index = list(range(1,epochs+1)) 275 | plt.figure(figsize=(10, 6)) 276 | sns.set(style='darkgrid') 277 | plt.plot(epoch_index, loss_list, marker='.', label='MSE Loss', color='red') 278 | plt.xlabel('Epoch') # x 轴标签 279 | plt.ylabel('MSE Loss') # y 轴标签 280 | plt.title(f'{data_name}车站,{model_name}模型,训练过程MSE Loss曲线图',fontproperties='SimHei', fontsize=20) 281 | plt.legend() 282 | plt.grid(True) # 添加网格背景 283 | plt.savefig('./img/{}_{}_training_MSEloss_epoch{}.png'.format(data_name, model_name, epochs)) 284 | plt.close() 285 | 286 | def predict(model_name): 287 | # predict 288 | with torch.no_grad(): 289 | if(model_name=="SEQ2SEQ"): 290 | # 关闭教师强制策略 291 | testY_hat = model(testX, testY, teacher_forcing_ratio=0) 292 | else: 293 | testY_hat = model(testX) 294 | predict = testY_hat 295 | real = testY 296 | 297 | # cal_loss 298 | MAEloss = nn.L1Loss() 299 | MSEloss = nn.MSELoss() 300 | 301 | # MAE: 平均绝对误差 302 | mae_val = MAEloss(predict, real) 303 | # MAPE: 平均绝对百分比误差 304 | mape_val = MAPELoss(predict, real) 305 | # MSE:均方误差 306 | mse_val = MSEloss(predict, real) 307 | # RMSE:均方根误差 308 | rmse_val = torch.sqrt(mse_val) 309 | 310 | losses = [ 311 | {'Loss Type': 'MAE Loss', 'Loss Value': mae_val.cpu().item()}, 312 | {'Loss Type': 'MAPE Loss', 'Loss Value': mape_val.cpu().item()}, 313 | {'Loss Type': 'RMSE Loss', 'Loss Value': rmse_val.cpu().item()} 314 | ] 315 | losses_df = pd.DataFrame(losses) 316 | print(losses_df) 317 | # predict = inverse_min_max_normalise_numpy(predict, min_vals=min_test_data[:,0], max_vals=max_test_data[:,0]) 318 | # real = inverse_min_max_normalise_numpy(real, min_vals=min_test_data[:,0], max_vals=max_test_data[:,0]) 319 | predict = inverse_min_max_normalise_tensor(testY_hat, min_vals=min_test_data, max_vals=max_test_data) 320 | real = inverse_min_max_normalise_tensor(testY, min_vals=min_test_data, max_vals=max_test_data) 321 | # draw 322 | predict = predict[0, :, 0].cpu().data.numpy() 323 | real = real[0, :, 0].cpu().data.numpy() 324 | # 生成时间序列，时间间隔为1小时，共24小时 325 | time_series = list(range(24)) 326 | # 定义时间标签 327 | time_labels = [f"{i}:00" for i in range(24)] 328 | font = FontProperties(family='SimHei', size=20) 329 | plt.figure(figsize=(12, 6)) 330 | sns.set(style='darkgrid') 331 | plt.ylabel('客流量', fontproperties=font) 332 | plt.plot(time_series, predict, marker='o', color='red', label='预测值') 333 | plt.plot(time_series, real, marker='o', color='blue', label='真实值') 334 | plt.xlabel('时间', fontproperties=font) 335 | plt.ylabel('客流量', fontproperties=font) 336 | plt.title(f'{data_name}车站2023年5月31日客流量,{model_name}模型,预测值效果对比图(epoch={epochs})', fontproperties=font) 337 | plt.legend(prop=font) 338 | plt.xticks(time_series, time_labels) # 设置时间标签 339 | plt.savefig('./img/{}_{}_prediction_epoch{}.png'.format(data_name, model_name, epochs)) 340 | plt.show() 341 | 342 | def main(model_name, save_mod): 343 | '''加载数据集''' 344 | global X, Y, testX, testY 345 | X, Y, testX, testY = load_data(cut_point=1272, data_path=data_path) 346 | 347 | '''读取超参数配置文件''' 348 | params = None 349 | with open(f"./config/{model_name.lower()}_params.json", 'r', encoding='utf-8') as file: 350 | params = json.load(file) 351 | 352 | if model_name == "RNN": 353 | rnn(params) 354 | elif model_name == "LSTM": 355 | lstm(params) 356 | elif model_name == "CNN_LSTM": 357 | cnn_lstm(params) 358 | elif model_name == "GRU": 359 | gru(params) 360 | elif model_name == "CNN": 361 | cnn(params) 362 | elif model_name == "Attention_LSTM": 363 | attention_lstm(params) 364 | elif model_name == "SEQ2SEQ": 365 | seq2seq(params) 366 | 367 | train(model_name=model_name, save_mod=save_mod) 368 | predict(model_name=model_name) 369 | 370 | if __name__ == '__main__': 371 | '''CHOOSE DATAPATH''' 372 | # (1464, 45, 1) ——> 45个BRT站点连续61days共计1464个小时的客流量数据 373 | station_list = ['第一码头', '开禾路口', '思北', '斗西路', 374 | '二市', '文灶', '金榜公园', '火车站', 375 | '莲坂', '龙山桥', '卧龙晓城', '东芳山庄', 376 | '洪文', '前埔枢纽站', '蔡塘', '金山', '市政务服务中心', 377 | '双十中学', '县后', '高崎机场', 'T4候机楼', '嘉庚体育馆', 378 | '诚毅学院', '华侨大学', '大学城', '产业研究院', '中科院', 379 | '东宅', '田厝', '厦门北站', '凤林', '东安', '后田', '东亭', 380 | '美峰', '蔡店', '潘涂', '滨海新城西柯枢纽', '官浔', '轻工食品园', 381 | '四口圳', '工业集中区', '第三医院', '城南', '同安枢纽'] 382 | # 创建一个字典，将数字与车站名称进行映射 383 | station_dict = {i: station_list[i-1] for i in range(1, len(station_list)+1)} 384 | # 创建 DataFrame 385 | df = pd.DataFrame(list(station_dict.items()), columns=['编号', '车站名称']) 386 | # 设置显示宽度以保持对齐 387 | pd.set_option('display.max_colwidth', 28) 388 | print(df.to_string(index=False)) 389 | user_input = input("请输入你要训练的车站数据集编号（1~45，否则随机选择）：") 390 | if int(user_input) in station_dict: 391 | data_name = station_dict[int(user_input)] 392 | else: 393 | # 随机选择一个站点(1-45之间的一个站点) 394 | random_station = np.random.randint(1, 46) 395 | data_name = station_list[random_station] 396 | data_path = f"./data/BRT/{data_name}_brtdata.npz" 397 | print("数据加载中, 请稍后……", data_path) 398 | 399 | '''INPUT YOUR MODEL NAME''' 400 | name_list = ["CNN", "RNN", "LSTM", "CNN_LSTM", "GRU", "Attention_LSTM", "SEQ2SEQ"] 401 | model_name = input("请输入要使用的模型【1: CNN 2: RNN 3: LSTM 4: CNN_LSTM 5: GRU 6: Attention_LSTM 7: SEQ2SEQ】\n") 402 | if model_name.isnumeric() and int(model_name) <= len(name_list): 403 | model_name = name_list[int(model_name) - 1] 404 | '''SAVE MODE''' 405 | save_mod = input("是否要保存训练后的模型？（输入 '1' 保存，否则不保存）\n") 406 | if int(save_mod) == 1: 407 | save_mod = True 408 | else: 409 | save_mod = False 410 | '''main()''' 411 | main(model_name, save_mod) -------------------------------------------------------------------------------- /models/AttentionLSTM/attention_lstm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class AttentionLSTMmodel(nn.Module): 6 | def __init__(self, 7 | input_size, 8 | hidden_size, 9 | num_layers, 10 | output_size, 11 | output_length=1*24, 12 | batch_first=True, 13 | bidirectional=False, 14 | device=None): 15 | super().__init__() 16 | self.output_length = output_length 17 | self.hidden_size = hidden_size 18 | 19 | # LSTM层 20 | # IN: (batch_size, seq_length, input_size) 21 | # OUT: (batch_size, seq_length, hidden_size) 22 | self.lstm = nn.LSTM( 23 | input_size=input_size, 24 | hidden_size=hidden_size, 25 | num_layers=num_layers, 26 | batch_first=batch_first, 27 | bidirectional=bidirectional 28 | ) 29 | 30 | # Attention 层 31 | # IN: (batch_size, seq_length, hidden_size) 32 | # OUT: (batch_size, output_length, hidden_size) 33 | self.query = nn.Linear(hidden_size, hidden_size) 34 | self.key = nn.Linear(hidden_size, hidden_size) 35 | self.value = nn.Linear(hidden_size, hidden_size) 36 | 37 | self.scale = torch.sqrt(torch.FloatTensor([hidden_size])).to(device) 38 | 39 | # MLP 层 40 | self.mlp = nn.Sequential( 41 | nn.Linear(hidden_size, hidden_size), 42 | nn.ReLU(), 43 | nn.Linear(hidden_size, output_size) 44 | ) 45 | 46 | self.activation = nn.ReLU() 47 | 48 | def attention(self, lstm_output): 49 | """ 50 | 计算注意力权重并应用到LSTM输出 51 | Args: 52 | lstm_output: shape (batch_size, seq_length, hidden_size) 53 | Returns: 54 | context_vector: shape (batch_size, 1, hidden_size) 55 | """ 56 | # 计算Q、K、V 57 | # (batch_size, output_length, hidden_size) 58 | Q = self.activation(self.query(lstm_output[:, -self.output_length:, :])) 59 | # (batch_size, seq_length, hidden_size) 60 | K = self.activation(self.key(lstm_output)) 61 | # (batch_size, seq_length, hidden_size) 62 | V = self.activation(self.value(lstm_output)) 63 | 64 | # 计算注意力分数 65 | # (batch_size, output_length, hidden_size) × (batch_size, hidden_size, seq_length) 66 | # -> (batch_size, output_length, seq_length) 67 | attention_weights = torch.bmm(Q, K.transpose(1, 2)) / self.scale 68 | attention_weights = torch.softmax(attention_weights, dim=2) 69 | 70 | # 将注意力权重应用到V 71 | # (batch_size, output_length, seq_length) × (batch_size, seq_length, hidden_size) 72 | # -> (batch_size, output_length, hidden_size) 73 | context_vector = torch.bmm(attention_weights, V) 74 | 75 | return context_vector 76 | 77 | def forward(self, x): 78 | """ 79 | 前向传播 80 | Args: 81 | x: 输入张量, shape (batch_size, seq_length, input_size) 82 | Returns: 83 | output: shape (batch_size, output_length, output_size) 84 | """ 85 | # LSTM 层 86 | # (batch_size, seq_length, input_size) -> (batch_size, seq_length, hidden_size) 87 | lstm_out, _ = self.lstm(x) 88 | 89 | # Attention 层 90 | # (batch_size, seq_length, hidden_size) -> (batch_size, output_length, hidden_size) 91 | context_vector = self.attention(lstm_out) 92 | 93 | # MLP 层 + residual connection 94 | mlp_input = context_vector + lstm_out[:, -self.output_length:, :] 95 | output = self.mlp(mlp_input) 96 | 97 | return output -------------------------------------------------------------------------------- /models/BiLSTM/bilstm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class LSTMmodel(nn.Module): 4 | def __init__(self, input_size, hidden_size, num_layers, output_size, output_length=1*24, batch_first=True): 5 | super().__init__() 6 | self.output_length=output_length 7 | self.lstm = nn.LSTM( 8 | input_size=input_size, 9 | hidden_size=hidden_size, 10 | num_layers=num_layers, 11 | batch_first=batch_first, 12 | bidirectional=True 13 | ) 14 | self.fc = nn.Sequential( 15 | nn.Linear(2*hidden_size, hidden_size), 16 | nn.ReLU(), 17 | nn.Linear(hidden_size, output_size), 18 | ) 19 | 20 | def forward(self, x): 21 | out, _ = self.lstm(x) 22 | out = self.fc(out[:, -self.output_length:, :]) 23 | return out -------------------------------------------------------------------------------- /models/CNN/cnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class CNNmodel(nn.Module): 4 | def __init__(self, input_size, hidden_size, output_size, output_length=1*24) -> None: 5 | super().__init__() 6 | self.output_length = output_length 7 | self.conv = nn.Conv1d(in_channels=input_size, out_channels=hidden_size, kernel_size=3, stride=1, padding=1) 8 | self.fc = nn.Sequential( 9 | nn.Linear(hidden_size, hidden_size), 10 | nn.ReLU(), 11 | nn.Linear(hidden_size, output_size), 12 | ) 13 | self.relu = nn.ReLU() 14 | 15 | def forward(self, x): 16 | # 输入数据形状为(samples, sequence_length, features) 17 | x = x.permute(0, 2, 1) # (samples, sequence_length, features) -> (samples, features, sequence_length) 18 | x = self.conv(x) 19 | x = self.relu(x) 20 | x = x.permute(0, 2, 1) # (samples, features, sequence_length) -> (samples, sequence_length, features) 21 | x = self.fc(x[:, -self.output_length:, :]) 22 | return x -------------------------------------------------------------------------------- /models/CNN_LSTM/cnn_lstm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class CNN_LSTM_Model(nn.Module): 4 | def __init__(self, input_size, output_size, output_length=1*24, batch_first=True) -> None: 5 | super().__init__() 6 | self.output_length = output_length 7 | self.conv1 = nn.Conv1d(in_channels=input_size, out_channels=32, kernel_size=3, stride=1, padding=1) 8 | self.conv2 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1) 9 | self.maxpool = nn.MaxPool1d(kernel_size=3, stride=1) 10 | self.lstm = nn.LSTM(input_size=32, hidden_size=64, num_layers=2, batch_first=batch_first) 11 | self.fc = nn.Linear(in_features=64, out_features=output_size) 12 | self.relu = nn.ReLU() 13 | 14 | def forward(self, x): 15 | # 输入数据形状为(samples, sequence_length, features) [B, S, F] 16 | x = x.permute(0, 2, 1) # [B, S, F] -> [B, F, S] 17 | x = self.conv1(x) # [B, F, S] -> [B, out_channels_1, S] 18 | x = self.relu(x) 19 | x = self.maxpool(x) # [B, out_channels_1, S] -> [B, out_channels_1, S-2] 20 | 21 | x = self.conv2(x) # [B, out_channels_1, S-2] -> [B, out_channels_2, S-2] 22 | x = self.relu(x) 23 | x = self.maxpool(x) # [B, out_channels_2, S-2] -> [B, out_channels, S-4] 24 | 25 | x = x.permute(0, 2, 1) 26 | x, _ = self.lstm(x) 27 | x = self.fc(x[:, -self.output_length:, :]) 28 | return x 29 | 30 | 31 | -------------------------------------------------------------------------------- /models/GRU/gru.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class GRUmodel(nn.Module): 4 | def __init__(self, input_size, hidden_size, num_layers, output_size, output_length=1*24, batch_first=True, bidirectional=False): 5 | super().__init__() 6 | self.output_length=output_length 7 | self.gru = nn.GRU( 8 | input_size=input_size, 9 | hidden_size=hidden_size, 10 | num_layers=num_layers, 11 | batch_first=batch_first, 12 | bidirectional=bidirectional 13 | ) 14 | self.fc = nn.Sequential( 15 | nn.Linear(hidden_size, hidden_size), 16 | nn.ReLU(), 17 | nn.Linear(hidden_size, output_size), 18 | ) 19 | 20 | def forward(self, x): 21 | out, _ = self.gru(x) 22 | out = self.fc(out[:, -self.output_length:, :]) 23 | return out 24 | -------------------------------------------------------------------------------- /models/LSTM/lstm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class LSTMmodel(nn.Module): 4 | def __init__(self, input_size, hidden_size, num_layers, output_size, output_length=1*24, batch_first=True, bidirectional=False): 5 | super().__init__() 6 | self.output_length=output_length 7 | self.lstm = nn.LSTM( 8 | input_size=input_size, 9 | hidden_size=hidden_size, 10 | num_layers=num_layers, 11 | batch_first=batch_first, 12 | bidirectional=bidirectional 13 | ) 14 | if(bidirectional==True): # 双向LSTM，隐状态翻倍 15 | self.fc = nn.Sequential( 16 | nn.Linear(2*hidden_size, hidden_size), 17 | nn.ReLU(), 18 | nn.Linear(hidden_size, output_size), 19 | ) 20 | else: 21 | self.fc = nn.Sequential( 22 | nn.Linear(hidden_size, hidden_size), 23 | nn.ReLU(), 24 | nn.Linear(hidden_size, output_size), 25 | ) 26 | 27 | def forward(self, x): 28 | out, _ = self.lstm(x) 29 | out = self.fc(out[:, -self.output_length:, :]) 30 | return out -------------------------------------------------------------------------------- /models/RNN/rnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class RNNmodel(nn.Module): 4 | def __init__(self, input_size, hidden_size, num_layers, output_size, output_length=1 * 24, batch_first=True, bidirectional=False): 5 | super().__init__() 6 | self.output_length = output_length 7 | self.rnn = nn.RNN( 8 | input_size=input_size, 9 | hidden_size=hidden_size, 10 | num_layers=num_layers, 11 | batch_first=batch_first, 12 | bidirectional=bidirectional 13 | ) 14 | self.fc1 = nn.Linear(in_features=hidden_size, out_features=hidden_size) 15 | self.fc2 = nn.Linear(in_features=hidden_size, out_features=output_size) 16 | self.relu = nn.ReLU() 17 | 18 | 19 | def forward(self, x): 20 | out, _ = self.rnn(x) 21 | out = self.fc1(out[:, -self.output_length:, :]) 22 | out = self.relu(out) 23 | out = self.fc2(out) 24 | return out -------------------------------------------------------------------------------- /models/SEQ2SEQ/seq2seq.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class Encoder(nn.Module): 5 | def __init__(self, input_size, hidden_size, num_layers): 6 | super(Encoder, self).__init__() 7 | self.hidden_size = hidden_size 8 | self.num_layers = num_layers 9 | self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) 10 | 11 | def forward(self, x): 12 | # 初始化隐藏状态和单元状态 13 | h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) 14 | c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) 15 | 16 | # LSTM前向传播 17 | out, (hidden_state, cell_state) = self.lstm(x, (h0, c0)) 18 | return out, (hidden_state, cell_state) 19 | 20 | class Decoder(nn.Module): 21 | def __init__(self, output_size, hidden_size, num_layers): 22 | super(Decoder, self).__init__() 23 | self.hidden_size = hidden_size 24 | self.num_layers = num_layers 25 | self.lstm = nn.LSTM(output_size, hidden_size, num_layers, batch_first=True) 26 | self.fc = nn.Linear(hidden_size, output_size) 27 | 28 | def forward(self, x, hidden_state, cell_state): 29 | out, (hidden, cell) = self.lstm(x, (hidden_state, cell_state)) 30 | out = self.fc(out) 31 | return out, (hidden, cell) 32 | 33 | class Seq2Seq(nn.Module): 34 | def __init__(self, encoder, decoder): 35 | super(Seq2Seq, self).__init__() 36 | self.encoder = encoder 37 | self.decoder = decoder 38 | 39 | def forward(self, input_sequence, target_sequence, teacher_forcing_ratio=1): 40 | encoder_output, (hidden_state, cell_state) = self.encoder(input_sequence) 41 | # 初始化解码器的输入 42 | decoder_input = torch.zeros(input_sequence.size(0), 1, target_sequence.size(2)).to(input_sequence.device) 43 | # decoder_input = torch.randn(input_sequence.size(0), 1, target_sequence.size(2)).to(input_sequence.device) 44 | # decoder_input = target_sequence[:, 0:1, :] 45 | 46 | # 存储预测结果 47 | predicted_outputs = [] 48 | 49 | # 是否使用教师强制 50 | use_teacher_forcing = True if torch.rand(1).item() < teacher_forcing_ratio else False 51 | 52 | # 初始化解码器的隐藏状态和细胞状态 53 | decoder_hidden = hidden_state 54 | decoder_cell = cell_state 55 | 56 | # 进行解码器的前向传播 57 | 58 | for t in range(target_sequence.size(1)): 59 | decoder_output, (decoder_hidden, decoder_cell) = self.decoder(decoder_input, decoder_hidden, decoder_cell) 60 | predicted_outputs.append(decoder_output) 61 | 62 | # 如果使用教师强制，则下一个时间步的解码器输入为真实目标序列；否则使用模型生成的输出 63 | decoder_input = target_sequence[:, t:t+1, :] if use_teacher_forcing else decoder_output 64 | 65 | # 将预测结果转换为三维张量 66 | predicted_outputs = torch.cat(predicted_outputs, dim=1) 67 | 68 | return predicted_outputs -------------------------------------------------------------------------------- /trainPEMS.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from models.LSTM.lstm import LSTMmodel 5 | from models.CNN_LSTM.cnn_lstm import CNN_LSTM_Model 6 | from models.RNN.rnn import RNNmodel 7 | from models.CNN.cnn import CNNmodel 8 | from models.GRU.gru import GRUmodel 9 | from models.AttentionLSTM.attention_lstm import AttentionLSTMmodel 10 | import matplotlib.pyplot as plt 11 | import seaborn as sns 12 | from matplotlib.font_manager import FontProperties 13 | import os 14 | import json 15 | import time 16 | import pandas as pd 17 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' 18 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 19 | X = None 20 | Y = None 21 | testX = None 22 | testY = None 23 | lr = 0.0001 24 | epochs = 100 25 | model = None 26 | min_test_data = None 27 | max_test_data = None # 测试集的最大最小值 28 | 29 | data_path = "./data/PEMS03/new_pems03_num88.npz" # 数据集路径 30 | # 5760 = 5min * 12 * 24 * 20 = 20天 31 | save_mod_dir = "./models/save" # 模型保存路径 32 | 33 | '''load_data''' 34 | def load_data(cut_point, data_path, input_length=12*24*7, output_length=12*24): 35 | global min_test_data, max_test_data 36 | ''' 37 | INPUT: 38 | cut_point, 训练和测试数据集划分点 39 | days, 客流数据的总天数 40 | hours, 每天小时数 41 | input_length, 模型输入的过去7天的数据作为一个序列 42 | output_length, 预测未来1天的数据 43 | OUTPUT: 44 | train_data 45 | test_data 46 | ''' 47 | data = np.load(data_path)['data'] 48 | data = data.astype(np.float32) 49 | data = data[:, np.newaxis] 50 | 51 | train_data = data[:cut_point] 52 | test_data = data[cut_point:] 53 | 54 | # 分别进行min-max scaling 55 | train_data, min_train_data, max_train_data = min_max_normalise_numpy(train_data) 56 | test_data, min_test_data, max_test_data = min_max_normalise_numpy(test_data) 57 | 58 | ''' 59 | input_data 模型输入 连续 input_length//7 天的数据 60 | output_data 模型输出 未来 output_length//7 天的数据 61 | ''' 62 | input_data = [] 63 | output_data = [] 64 | 65 | # 滑动窗口法 66 | for i in range(len(train_data) - input_length - output_length + 1): 67 | input_seq = train_data[i : i + input_length] 68 | output_seq = train_data[i + input_length : i + input_length + output_length, 0] 69 | 70 | input_data.append(input_seq) 71 | output_data.append(output_seq) 72 | 73 | # 转为torch.tensor 74 | X = torch.tensor(input_data, dtype=torch.float32, device=device) 75 | Y = torch.tensor(output_data, dtype=torch.float32, device=device) 76 | Y = Y.unsqueeze(-1) # 在最后一个维度（维度索引为1）上增加一个维度 77 | # X = torch.tensor([item.cpu().detach().numpy() for item in input_data], dtype=torch.float32, device=device) 78 | # Y = torch.tensor([item.cpu().detach().numpy() for item in output_data], dtype=torch.float32, device=device) 79 | 80 | test_inseq = [] 81 | test_outseq = [] 82 | # 滑动窗口法 83 | for i in range(len(test_data) - input_length - output_length + 1): 84 | input_seq = test_data[i : i + input_length] 85 | output_seq = test_data[i + input_length : i + input_length + output_length, 0] 86 | 87 | test_inseq.append(input_seq) 88 | test_outseq.append(output_seq) 89 | 90 | # 转为torch.tensor 91 | testX = torch.tensor(test_inseq, dtype=torch.float32, device=device) 92 | testY = torch.tensor(test_outseq, dtype=torch.float32, device=device) 93 | testY = testY.unsqueeze(-1) # 在最后一个维度（维度索引为1）上增加一个维度 94 | # testX = torch.tensor([item.cpu().detach().numpy() for item in test_inseq], dtype=torch.float32, device=device) 95 | # testY = torch.tensor([item.cpu().detach().numpy() for item in test_outseq], dtype=torch.float32, device=device) 96 | 97 | # 输出数据形状 98 | print("数据集处理完毕：") 99 | print("data - 原数据集 shape:", data.shape) 100 | print("traindata - Input shape:", X.shape) 101 | print("traindata - Output shape:", Y.shape) 102 | print("testdata - Input shape:", testX.shape) 103 | print("testdata - Output shape:", testY.shape) 104 | return X, Y, testX, testY 105 | 106 | '''min-max Scaling''' 107 | def min_max_normalise_numpy(x): 108 | # shape: [sequence_length, features] 109 | min_vals = np.min(x, axis=0) 110 | max_vals = np.max(x, axis=0) 111 | # [features] -> shape: [1, features] 112 | min_vals = np.expand_dims(min_vals, axis=0) 113 | max_vals = np.expand_dims(max_vals, axis=0) 114 | # 归一化 -> [-1, 1] 115 | normalized_data = 2 * (x - min_vals) / (max_vals - min_vals) - 1 116 | # 归一化 -> [0, 1] 117 | # normalized_data = (x - min_vals) / (max_vals - min_vals) 118 | return normalized_data, min_vals, max_vals 119 | 120 | def inverse_min_max_normalise_numpy(normalised_x, min_vals, max_vals): 121 | x = (normalised_x + 1) / 2 * (max_vals - min_vals) + min_vals 122 | return x 123 | 124 | def min_max_normalise_tensor(x): 125 | # shape: [samples, sequence_length, features] 126 | min_vals = x.min(dim=1).values.unsqueeze(1) 127 | max_vals = x.max(dim=1).values.unsqueeze(1) 128 | # data ->[-1, 1] 129 | normalise_x = 2 * (x - min_vals) / (max_vals - min_vals) - 1 130 | return normalise_x, min_vals, max_vals 131 | 132 | def inverse_min_max_normalise_tensor(x, min_vals, max_vals): 133 | min_vals = torch.tensor(min_vals).to(device) 134 | max_vals = torch.tensor(max_vals).to(device) 135 | # shape: [1, features] -> [1, 1, features] 136 | min_vals = min_vals.unsqueeze(0) 137 | max_vals = max_vals.repeat(x.shape[0], 1, 1) 138 | # [1, 1, features] -> [samples, 1, features] 139 | min_vals = min_vals.repeat(x.shape[0], 1, 1) 140 | max_vals = max_vals.repeat(x.shape[0], 1, 1) 141 | 142 | x = (x + 1) / 2 * (max_vals - min_vals) + min_vals 143 | return x 144 | 145 | # MAPE: 平均绝对百分比误差 146 | def MAPELoss(y_hat, y): 147 | x = torch.tensor(0.0001, dtype=torch.float32).to(device) 148 | y_new = torch.where(y==0, x, y) # 防止分母为0 149 | abs_error = torch.abs((y - y_hat) / y_new) 150 | mape = 100. * torch.mean(abs_error) 151 | return mape 152 | 153 | def cnn(params): 154 | global lr, epochs, model 155 | '''超参数加载''' 156 | input_size = params["input_size"] # 输入特征数 157 | hidden_size = params["hidden_size"] # CNN output_channels 158 | output_size = params["output_size"] # 输出特征数 159 | lr = params["lr"] # 学习率 160 | epochs = params["epochs"] # 训练轮数 161 | model = CNNmodel(input_size, hidden_size, output_size, output_length=24*12).to(device) 162 | 163 | def rnn(params): 164 | global lr, epochs, model 165 | '''超参数加载''' 166 | input_size = params["input_size"] # 输入特征数 167 | hidden_size = params["hidden_size"] # RNN隐藏层神经元数 168 | num_layers = params["num_layers"] # RNN层数 169 | output_size = params["output_size"] # 输出特征数 170 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 171 | lr = params["lr"] # 学习率 172 | epochs = params["epochs"] # 训练轮数 173 | model = RNNmodel(input_size, hidden_size, num_layers, output_size, output_length=24*12, bidirectional=bidirectional).to(device) 174 | 175 | def lstm(params): 176 | global lr, epochs, model 177 | '''超参数加载''' 178 | input_size = params["input_size"] # 输入特征数 179 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 180 | num_layers = params["num_layers"] # LSTM层数 181 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 182 | output_size = params["output_size"] # 输出特征数 183 | lr = params["lr"] # 学习率 184 | epochs = params["epochs"] # 训练轮数 185 | 186 | model = LSTMmodel(input_size, hidden_size, num_layers, output_size, output_length=24*12, bidirectional=bidirectional).to(device) 187 | 188 | def cnn_lstm(params): 189 | global lr, epochs, model 190 | '''超参数加载''' 191 | input_size = params["input_size"] # 输入特征数 192 | output_size = params["output_size"] # 输出特征数 193 | lr = params["lr"] # 学习率 194 | epochs = params["epochs"] # 训练轮数 195 | model = CNN_LSTM_Model(input_size=1, output_size=output_size, output_length=24*12).to(device) 196 | 197 | def gru(params): 198 | global lr, epochs, model 199 | '''超参数加载''' 200 | input_size = params["input_size"] # 输入特征数 201 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 202 | num_layers = params["num_layers"] # LSTM层数 203 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 204 | output_size = params["output_size"] # 输出特征数 205 | lr = params["lr"] # 学习率 206 | epochs = params["epochs"] # 训练轮数 207 | model = GRUmodel(input_size, hidden_size, num_layers, output_size, output_length=24*12, bidirectional=bidirectional).to(device) 208 | 209 | def attention_lstm(params): 210 | global lr, epochs, model 211 | '''超参数加载''' 212 | input_size = params["input_size"] # 输入特征数 213 | hidden_size = params["hidden_size"] # LSTM隐藏层神经元数 214 | num_layers = params["num_layers"] # LSTM层数 215 | bidirectional = params["bidirectional"].lower() == "true" # 是否双向 216 | output_size = params["output_size"] # 输出特征数 217 | lr = params["lr"] # 学习率 218 | epochs = params["epochs"] # 训练轮数 219 | model = AttentionLSTMmodel(input_size, hidden_size, num_layers, output_size, output_length=24*12, bidirectional=bidirectional).to(device) 220 | 221 | def train(model_name, save_mod=False): 222 | loss_function = nn.MSELoss() 223 | optimizer = torch.optim.Adam(model.parameters(), lr=lr) 224 | 225 | loss_list = [] # 保存训练过程中loss数据 226 | 227 | # 开始计时 228 | start_time = time.time() # 记录模型训练开始时间 229 | # train model 230 | for epoch in range(epochs): 231 | Y_hat = model(X) 232 | loss = loss_function(Y_hat, Y) 233 | optimizer.zero_grad() 234 | loss.backward() 235 | nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=10, norm_type=2) 236 | optimizer.step() 237 | print(f'Epoch [{epoch+1}/{epochs}], MSE-Loss: {loss.item()}') 238 | loss_list.append(loss.item()) 239 | end_time = time.time() # 记录模型训练结束时间 240 | total_time = end_time - start_time # 计算总耗时 241 | print(f"本次模型训练总耗时: {total_time} 秒，超越了全国99.2%的单片机，太棒啦！") 242 | if save_mod == True: 243 | # save model -> ./**/model_name_eopch{n}.pth 244 | torch.save(model.state_dict(), '{}/PEMS_{}_epoch{}.pth'.format(save_mod_dir, model_name, epochs)) 245 | print("Trained Model have saved in:", '{}/{}_epoch{}.pth'.format(save_mod_dir, model_name, epochs)) 246 | 247 | # save training loss graph 248 | epoch_index = list(range(1,epochs+1)) 249 | plt.figure(figsize=(10, 6)) 250 | sns.set(style='darkgrid') 251 | plt.plot(epoch_index, loss_list, marker='.', label='MSE Loss', color='red') 252 | plt.xlabel('Epoch') # x 轴标签 253 | plt.ylabel('MSE Loss') # y 轴标签 254 | plt.title(f'PEMS03数据集,{model_name}模型,训练过程MSE Loss曲线图',fontproperties='SimHei', fontsize=20) 255 | plt.legend() 256 | plt.grid(True) # 添加网格背景 257 | plt.savefig('./img/PEMS03_{}_training_MSEloss_epoch{}.png'.format(model_name, epochs)) 258 | plt.close() 259 | 260 | def predict(model_name): 261 | # predict 262 | with torch.no_grad(): 263 | testY_hat = model(testX) 264 | predict = testY_hat 265 | real = testY 266 | # cal_loss 267 | MAEloss = nn.L1Loss() 268 | MSEloss = nn.MSELoss() 269 | 270 | # MAE: 平均绝对误差 271 | mae_val = MAEloss(predict, real) 272 | # MAPE: 平均绝对百分比误差 273 | mape_val = MAPELoss(predict, real) 274 | # MSE:均方误差 275 | mse_val = MSEloss(predict, real) 276 | # RMSE:均方根误差 277 | rmse_val = torch.sqrt(mse_val) 278 | 279 | losses = [ 280 | {'Loss Type': 'MAE Loss', 'Loss Value': mae_val.cpu().item()}, 281 | {'Loss Type': 'MAPE Loss', 'Loss Value': mape_val.cpu().item()}, 282 | {'Loss Type': 'RMSE Loss', 'Loss Value': rmse_val.cpu().item()} 283 | ] 284 | losses_df = pd.DataFrame(losses) 285 | print(losses_df) 286 | 287 | # predict = inverse_min_max_normalise_tensor(predict, min_vals=min_test_data, max_vals=max_test_data) 288 | # real = inverse_min_max_normalise_tensor(real, min_vals=min_test_data, max_vals=max_test_data) 289 | 290 | # draw 291 | predict_np = predict[10, :, 0].cpu().data.numpy() 292 | real_np = real[10, :, 0].cpu().data.numpy() 293 | 294 | predict_np = inverse_min_max_normalise_numpy(predict_np, min_vals=min_test_data[:,0], max_vals=max_test_data[:,0]) 295 | real_np = inverse_min_max_normalise_numpy(real_np, min_vals=min_test_data[:,0], max_vals=max_test_data[:,0]) 296 | 297 | font = FontProperties(family='SimHei', size=20) 298 | plt.figure(figsize=(12, 6)) 299 | sns.set(style='darkgrid') 300 | plt.plot(predict_np, color='red', label='预测值') 301 | plt.plot(real_np, color='blue', label='真实值') 302 | plt.xlabel('时间', fontproperties=font) 303 | plt.ylabel('车流量', fontproperties=font) 304 | plt.title(f'PEMS03车流量数据集,{model_name}模型,预测值效果对比图(epoch={epochs})', fontproperties=font) 305 | plt.legend(prop=font) 306 | plt.savefig('./img/PEMS03_{}_prediction_epoch{}.png'.format(model_name, epochs)) 307 | plt.show() 308 | 309 | def main(model_name, save_mod): 310 | '''加载数据集''' 311 | global X, Y, testX, testY 312 | X, Y, testX, testY = load_data(cut_point=3168, data_path=data_path) 313 | 314 | '''读取超参数配置文件''' 315 | params = None 316 | with open(f"./config/{model_name.lower()}_params.json", 'r', encoding='utf-8') as file: 317 | params = json.load(file) 318 | 319 | if model_name == "RNN": 320 | rnn(params) 321 | elif model_name == "LSTM": 322 | lstm(params) 323 | elif model_name == "CNN_LSTM": 324 | cnn_lstm(params) 325 | elif model_name == "GRU": 326 | gru(params) 327 | elif model_name == "CNN": 328 | cnn(params) 329 | elif model_name == "Attention_LSTM": 330 | attention_lstm(params) 331 | 332 | 333 | train(model_name=model_name, save_mod=save_mod) 334 | predict(model_name=model_name) 335 | 336 | def continue_main(model_name): 337 | '''加载数据集''' 338 | global X, Y, testX, testY 339 | X, Y, testX, testY = load_data(cut_point=3168, data_path=data_path) 340 | 341 | '''读取超参数配置文件''' 342 | params = None 343 | with open(f"./config/{model_name.lower()}_params.json", 'r', encoding='utf-8') as file: 344 | params = json.load(file) 345 | 346 | if model_name == "RNN": 347 | rnn(params) 348 | elif model_name == "LSTM": 349 | lstm(params) 350 | elif model_name == "CNN_LSTM": 351 | cnn_lstm(params) 352 | elif model_name == "GRU": 353 | gru(params) 354 | elif model_name == "CNN": 355 | cnn(params) 356 | elif model_name == "Attention_LSTM": 357 | attention_lstm(params) 358 | 359 | # 设置之前保存模型的路径 360 | saved_model_path = './models/save/PEMS_CNN_LSTM_epoch70.pth' 361 | # 加载模型权重 362 | model.load_state_dict(torch.load(saved_model_path)) 363 | predict(model_name=model_name) 364 | 365 | if __name__ == '__main__': 366 | '''INPUT YOUR MODEL NAME''' 367 | name_list = ["CNN", "RNN", "LSTM", "CNN_LSTM", "GRU", "Attention_LSTM"] 368 | model_name = input("请输入要使用的模型【1: CNN 2: RNN 3: LSTM 4: CNN_LSTM 5: GRU 6: Attention_LSTM】\n") 369 | if model_name.isnumeric() and int(model_name) <= len(name_list): 370 | model_name = name_list[int(model_name) - 1] 371 | '''SAVE MODE''' 372 | save_mod = input("是否要保存训练后的模型？（输入 '1' 保存，否则不保存）\n") 373 | if int(save_mod) == 1: 374 | save_mod = True 375 | else: 376 | save_mod = False 377 | 378 | # continue_main(model_name=model_name) 379 | '''main()''' 380 | main(model_name, save_mod) --------------------------------------------------------------------------------