代码

  2. import copy
  3. import os
  5. import torch
  6. import torch.nn as nn
  7. import numpy as np
  8. import pandas as pd
  9. from sklearn.preprocessing import MinMaxScaler
  10. import datetime
  11. import joblib
  12. import matplotlib.pyplot as plt
  15. class LSTMRegression(nn.Module):
  16.     """
  17.     LSTM模型时间序列预测
  18.     """
  20.     def __init__(self, input_size, hidden_size,  num_layers, output_size=1,):
  21.         super().__init__()
  22.         # LSTM模型
  23.         self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
  24.         # RNN模型
  25.         # self.lstm = nn.RNN(input_size, hidden_size, num_layers)
  26.         self.fc = nn.Linear(hidden_size, output_size)
  27.         self.last_time = None
  28.         self.seq_step = []
  30.     def forward(self, _x):
  31.         x, _ = self.lstm(_x)  # _x is input, size (seq_len, batch, input_size)
  32.         x = self.fc(x)
  33.         return x
  36. class LSTMProduct:
  38.     def __init__(self, seq_step=5, hidden_size=12, num_layers=2):
  39.         self.seq_step = seq_step
  40.         self.sate = None
  41.         if os.path.exists('./parameter.pkl'):
  42.             self.sate = torch.load("./parameter.pkl")
  43.             # 预测模式
  44.             self.lstm_model = LSTMRegression(seq_step, hidden_size=hidden_size, num_layers=num_layers)
  45.             # 重新加载模型参数
  46.             self.lstm_model.load_state_dict(self.sate["model"])
  47.         else:
  48.             self.lstm_model = LSTMRegression(seq_step, hidden_size=hidden_size, num_layers=num_layers)
  51.     def load_data(self, path):
  52.         """
  53.         :param path: CSV文件路径, 时间间隔无法确定的情况下要考虑是否需要重采样
  54.         :return:
  55.         """
  56.         df = pd.read_csv(path)
  57.         df['time'] = pd.to_datetime(df['time'])
  58.         # 根据时间进行排序
  59.         df.sort_values(by='time', ascending=True, inplace=True)
  60.         # 将时间设置成索引
  61.         df2 = df.set_index(keys='time')
  62.         # 时间重采样 1H
  63.         df3 = df2.resample("1H").mean()
  64.         index = df3.index
  65.         data = df3['level'].values
  66.         return index, data
  68.     def create_dataset(self, data):
  69.         """
  70.         创建数据集
  71.         data格式: (index, values)
  72.         """
  73.         index, value = data
  74.         dataset_x, dataset_y, dataset_y_index = [], [], []
  75.         for i in range(len(value) - self.seq_step):
  76.             _x = value[i:(i + self.seq_step)]
  77.             dataset_x.append(_x)
  78.             dataset_y.append(value[i + self.seq_step])
  79.             dataset_y_index.append(index[i + self.seq_step])
  80.         return np.array(dataset_x), np.array(dataset_y), dataset_y_index
  82.     def data_parser(self, data):
  83.         """
  84.         数据处理: 将数据归一化处理,不然损失结果非常大
  85.         data格式: (index, values)
  86.         """
  87.         index, value = data
  88.         scaler = MinMaxScaler(feature_range=(-1, 1))
  89.         scaler_value = scaler.fit_transform(value.reshape(-1, 1)).reshape(-1, )
  90.         joblib.dump(scaler, "scaler.model")
  91.         return index, scaler_value
  93.     def result_data_parser(self, out):
  94.         """预测结果处理"""
  95.         scaler = joblib.load("scaler.model")
  96.         data = scaler.inverse_transform(np.array(out.view(-1).tolist()).reshape(-1, 1))
  97.         return data
  99.     def product(self, train_x):
  100.         out = self.lstm_model(train_x)
  101.         scaler = joblib.load("scaler.model")
  102.         scaler_value = scaler.inverse_transform(np.array(out.tolist()).reshape(-1, 1)).reshape(-1, )
  103.         return scaler_value.reshape(-1, ), out.view(-1).tolist()
  105.     def train_model(self, train_data, epoch=5000, lr=0.01):
  106.         """
  107.         训练模型
  108.         """
  109.         train_x, train_y = train_data
  110.         loss_function = nn.MSELoss(reduce=True, size_average=True)
  111.         optimizer = torch.optim.Adam(self.lstm_model.parameters(), lr=lr)
  112.         min_loss = 100
  113.         flag = 0
  114.         for i in range(epoch):
  115.             out = self.lstm_model(train_x)
  116.             loss = loss_function(out, train_y)
  117.             # 反向传播
  118.             loss.backward()
  119.             # 计算梯度
  120.             optimizer.step()
  121.             # 梯度0
  122.             optimizer.zero_grad()
  123.             # 损失值
  124.             loss_value = loss.item()
  125.             if (i+1) % 100 == 0:
  126.                 print('Epoch: {}, Loss:{:.6f}'.format(i+1, loss_value))
  127.             # 保存损失最小的模型
  128.             if min_loss > loss_value:
  129.                 self.sate = {"model": self.lstm_model.state_dict(), "last_seq": self.lstm_model.seq_step,
  130.                         "last_time": self.lstm_model.last_time}
  131.                 # 如果损失比较小提前结束训练
  132.                 if min_loss - loss_value <= 1e-8:
  133.                     flag += 1
  134.                 else:
  135.                     flag = 0
  136.                 if flag >= 3:
  137.                     torch.save(self.sate, './parameter.pkl')
  138.                     break
  139.                 min_loss = loss_value
  140.                 torch.save(self.sate, './parameter.pkl')
  142.     def train(self, data_path='./train_data.csv'):
  143.         """训练模型"""
  144.         # 加载数据
  145.         data = self.load_data(data_path)
  146.         # # 数据处理
  147.         parser_data = self.data_parser(data)
  148.         self.lstm_model.last_time = parser_data[0][-1]
  149.         self.lstm_model.seq_step = parser_data[-1][-self.seq_step:].tolist()
  150.         # 如果是训练模型就训练模型
  151.         dataset = self.create_dataset(parser_data)
  152.         # 将数据改变形状，RNN 读入的数据维度是 (seq_size, batch_size, feature_size)
  153.         train_x = dataset[0].reshape(1, -1, self.seq_step)
  154.         train_y = dataset[1].reshape(1, -1, 1)
  155.         # 转为pytorch的tensor对象
  156.         train_x = torch.from_numpy(train_x).to(torch.float32)
  157.         train_y = torch.from_numpy(train_y).to(torch.float32)
  158.         # 训练模型
  159.         self.train_model((train_x, train_y))
  160.         return self
  162.     def get_mse(self, data_1, data_2):
  163.         """
  164.         获取均方误差
  165.         :return:
  166.         """
  167.         input_ = torch.from_numpy(np.array(data_1))
  168.         output = torch.from_numpy(np.array(data_2))
  169.         loss_function = nn.MSELoss(reduce=True, size_average=True)
  170.         return loss_function(input_, output).item()
  172.     def forecast(self, num=5):
  173.         forecast_list = []
  174.         seq_step = copy.deepcopy(self.sate.get("last_seq"))
  175.         last_time = self.sate.get("last_time")
  176.         for i in range(num):
  177.             train_x = np.array(seq_step[-self.seq_step:]).reshape(1, -1, self.seq_step)
  178.             train_x = train_x.reshape(1, -1, self.seq_step)
  179.             train_x = torch.from_numpy(train_x).to(torch.float32)
  180.             product_data, scalar_product_data = self.product(train_x)
  181.             current_time = last_time + np.timedelta64(1, 'h')
  182.             last_time = current_time
  183.             forecast_list.append((current_time.strftime("%Y-%m-%d %H:%M:%S"), product_data[0]))
  184.             seq_step.extend(scalar_product_data)
  185.         return forecast_list
  187.     def test(self):
  188.         # 未来24小时数据用来计算测试误差
  189.         data_1 = [165, 142, 136,147,170,196,217,226,224,217,213,217,233,262,302,347,387,415, 423,410,377,333,283,235]
  190.         data = [i[-1] for i in self.forecast(24)]
  191.         print(self.get_mse(data, data_1))
  194. if __name__ == '__main__':
  195.     lstm_model = LSTMProduct(seq_step=60, hidden_size=2, num_layers=1)
  196.     # lstm_model.train(data_path="./train_data.csv")
  197.     lstm_model.test()
  198.     print(lstm_model.forecast(24))

数据

train_data.csv