Predict_LSTM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "#### 导入相关库"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import logging\n",
    "import warnings\n",
    "from logging.handlers import RotatingFileHandler\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "from torch.nn import Module, LSTM, Linear\n",
    "from torch.utils.data import DataLoader, TensorDataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 数据集模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "outputs": [],
   "source": [
    "class Data:\n",
    "    def __init__(self, config):\n",
    "        self.config = config\n",
    "        self.data, self.data_column_name = self.read_data()\n",
    "        self.data_num = self.data.shape[0]\n",
    "        self.train_num = int(self.data_num * self.config.train_data_rate)\n",
    "        self.mean = np.mean(self.data, axis=0)\n",
    "        self.std = np.std(self.data, axis=0)\n",
    "        self.norm_data = (self.data - self.mean) / self.std  # 归一化，去量纲\n",
    "        self.start_num_in_test = 0  # 测试集中前几天的数据会被删掉，因为它不够一个time_step\n",
    "\n",
    "    def read_data(self):\n",
    "        if self.config.debug_mode:\n",
    "            init_data = pd.read_csv(self.config.train_data_path, nrows=self.config.debug_num,\n",
    "                                    usecols=self.config.feature_columns)\n",
    "        else:\n",
    "            init_data = pd.read_csv(self.config.train_data_path, usecols=self.config.feature_columns)\n",
    "        return init_data.values, init_data.columns.tolist()  # .columns.tolist() 是获取列名\n",
    "\n",
    "    def get_train_and_valid_data(self):\n",
    "        feature_data = self.norm_data[:self.train_num]\n",
    "        label_data = self.norm_data[self.config.predict_day: self.config.predict_day + self.train_num,\n",
    "                     self.config.label_in_feature_index]  # 将延后几天的数据作为label\n",
    "        if not self.config.do_continue_train:\n",
    "            # 在非连续训练模式下，每time_step行数据会作为一个样本，两个样本错开一行\n",
    "            # 比如：1-20行，2-21行···\n",
    "            train_x = [feature_data[i:i + self.config.time_step] for i in range(self.train_num - self.config.time_step)]\n",
    "            train_y = [label_data[i:i + self.config.time_step] for i in range(self.train_num - self.config.time_step)]\n",
    "        else:\n",
    "            # 在连续训练模式下，每time_step行数据会作为一个样本，两个样本错开time_step行，\n",
    "            # 比如：1-20行，21-40行···到数据末尾，然后又是 2-21行，22-41行。。。到数据末尾，……\n",
    "            train_x = [\n",
    "                feature_data[start_index + i * self.config.time_step: start_index + (i + 1) * self.config.time_step]\n",
    "                for start_index in range(self.config.time_step)\n",
    "                for i in range((self.train_num - start_index) // self.config.time_step)]\n",
    "            train_y = [\n",
    "                label_data[start_index + i * self.config.time_step: start_index + (i + 1) * self.config.time_step]\n",
    "                for start_index in range(self.config.time_step)\n",
    "                for i in range((self.train_num - start_index) // self.config.time_step)]\n",
    "        train_x, train_y = np.array(train_x), np.array(train_y)\n",
    "        # 划分训练和验证集，并打乱\n",
    "        train_x, valid_x, train_y, valid_y = train_test_split(train_x, train_y, test_size=self.config.valid_data_rate,\n",
    "                                                              random_state=self.config.random_seed,\n",
    "                                                              shuffle=self.config.shuffle_train_data)\n",
    "        return train_x, valid_x, train_y, valid_y\n",
    "\n",
    "    def get_test_data(self, return_label_data=False):\n",
    "        feature_data = self.norm_data[self.train_num:]\n",
    "        sample_interval = min(feature_data.shape[0], self.config.time_step)  # 防止time_step大于测试集数量\n",
    "        self.start_num_in_test = feature_data.shape[0] % sample_interval  # 这些天的数据不够一个sample_interval\n",
    "        time_step_size = feature_data.shape[0] // sample_interval\n",
    "        # 在测试数据中，每time_step行数据会作为一个样本，两个样本错开time_step行\n",
    "        # 比如：1-20行，21-40行···到数据末尾。\n",
    "        test_x = [feature_data[\n",
    "                  self.start_num_in_test + i * sample_interval: self.start_num_in_test + (i + 1) * sample_interval]\n",
    "                  for i in range(time_step_size)]\n",
    "        if return_label_data:  # 实际应用中的测试集是没有label数据的\n",
    "            label_data = self.norm_data[self.train_num + self.start_num_in_test:, self.config.label_in_feature_index]\n",
    "            return np.array(test_x), label_data\n",
    "        return np.array(test_x)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 建立LSTM时间序列预测模型"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "outputs": [],
   "source": [
    "class Net(Module):\n",
    "    '''\n",
    "    pytorch预测模型，包括LSTM时序预测层和Linear回归输出层\n",
    "    '''\n",
    "    def __init__(self, config):\n",
    "        super(Net, self).__init__()\n",
    "        self.lstm = LSTM(input_size=config.input_size, hidden_size=config.hidden_size,\n",
    "                         num_layers=config.lstm_layers, batch_first=True, dropout=config.dropout_rate)\n",
    "        self.linear = Linear(in_features=config.hidden_size, out_features=config.output_size)\n",
    "\n",
    "    def forward(self, x, hidden=None):\n",
    "        lstm_out, hidden = self.lstm(x, hidden)\n",
    "        linear_out = self.linear(lstm_out)\n",
    "        return linear_out, hidden"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 模型训练模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "outputs": [],
   "source": [
    "def train(config, logger, train_and_valid_data):\n",
    "    if config.do_train_visualized:\n",
    "        import visdom\n",
    "        vis = visdom.Visdom(env='model_pytorch')\n",
    "    train_X, train_Y, valid_X, valid_Y = train_and_valid_data\n",
    "    train_X, train_Y = torch.from_numpy(train_X).float(), torch.from_numpy(train_Y).float()\n",
    "    train_loader = DataLoader(TensorDataset(train_X, train_Y), batch_size=config.batch_size)\n",
    "    valid_X, valid_Y = torch.from_numpy(valid_X).float(), torch.from_numpy(valid_Y).float()\n",
    "    valid_loader = DataLoader(TensorDataset(valid_X, valid_Y), batch_size=config.batch_size)\n",
    "    device = torch.device(\"cuda:0\" if config.use_cuda and torch.cuda.is_available() else \"cpu\")\n",
    "    model = Net(config).to(device)\n",
    "    if config.add_train:\n",
    "        model.load_state_dict(torch.load(config.model_save_path + config.model_name))\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)\n",
    "    criterion = torch.nn.MSELoss()\n",
    "    valid_loss_min = float(\"inf\")\n",
    "    bad_epoch = 0\n",
    "    global_step = 0\n",
    "    for epoch in range(config.epoch):\n",
    "        logger.info(\"Epoch {}/{}\".format(epoch, config.epoch))\n",
    "        model.train()\n",
    "        train_loss_array = []\n",
    "        hidden_train = None\n",
    "        for i, _data in enumerate(train_loader):\n",
    "            _train_X, _train_Y = _data[0].to(device),_data[1].to(device)\n",
    "            optimizer.zero_grad()\n",
    "            pred_Y, hidden_train = model(_train_X, hidden_train)\n",
    "            if not config.do_continue_train:\n",
    "                hidden_train = None\n",
    "            else:\n",
    "                h_0, c_0 = hidden_train\n",
    "                h_0.detach_(), c_0.detach_()    # 去掉梯度信息\n",
    "                hidden_train = (h_0, c_0)\n",
    "            loss = criterion(pred_Y, _train_Y)  # 计算loss\n",
    "            loss.backward()                     # 将loss反向传播\n",
    "            optimizer.step()                    # 用优化器更新参数\n",
    "            train_loss_array.append(loss.item())\n",
    "            global_step += 1\n",
    "            if config.do_train_visualized and global_step % 100 == 0:\n",
    "                vis.line(X=np.array([global_step]), Y=np.array([loss.item()]), win='Train_Loss',\n",
    "                         update='append' if global_step > 0 else None, name='Train', opts=dict(showlegend=True))\n",
    "        # 以下为早停机制，当模型训练连续config.patience个epoch都没有使验证集预测效果提升时，就停止，防止过拟合\n",
    "        model.eval()\n",
    "        valid_loss_array = []\n",
    "        hidden_valid = None\n",
    "        for _valid_X, _valid_Y in valid_loader:\n",
    "            _valid_X, _valid_Y = _valid_X.to(device), _valid_Y.to(device)\n",
    "            pred_Y, hidden_valid = model(_valid_X, hidden_valid)\n",
    "            if not config.do_continue_train: hidden_valid = None\n",
    "            loss = criterion(pred_Y, _valid_Y)\n",
    "            valid_loss_array.append(loss.item())\n",
    "        train_loss_cur = np.mean(train_loss_array)\n",
    "        valid_loss_cur = np.mean(valid_loss_array)\n",
    "        logger.info(\"The train loss is {:.6f}. \".format(train_loss_cur) +\n",
    "              \"The valid loss is {:.6f}.\".format(valid_loss_cur))\n",
    "        if config.do_train_visualized:\n",
    "            vis.line(X=np.array([epoch]), Y=np.array([train_loss_cur]), win='Epoch_Loss',\n",
    "                     update='append' if epoch > 0 else None, name='Train', opts=dict(showlegend=True))\n",
    "            vis.line(X=np.array([epoch]), Y=np.array([valid_loss_cur]), win='Epoch_Loss',\n",
    "                     update='append' if epoch > 0 else None, name='Eval', opts=dict(showlegend=True))\n",
    "        if valid_loss_cur < valid_loss_min:\n",
    "            valid_loss_min = valid_loss_cur\n",
    "            bad_epoch = 0\n",
    "            torch.save(model.state_dict(), config.model_save_path + config.model_name)\n",
    "        else:\n",
    "            bad_epoch += 1\n",
    "            # 如果验证集指标连续patience个epoch没有提升，就停掉训练\n",
    "            if bad_epoch >= config.patience:\n",
    "                logger.info(\" The training stops early in epoch {}\".format(epoch))\n",
    "                break"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 模型预测模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "outputs": [],
   "source": [
    "def predict(config, test_X):\n",
    "    # 获取测试数据\n",
    "    test_X = torch.from_numpy(test_X).float()\n",
    "    test_set = TensorDataset(test_X)\n",
    "    test_loader = DataLoader(test_set, batch_size=1)\n",
    "    # 加载模型\n",
    "    device = torch.device(\"cuda:0\" if config.use_cuda and torch.cuda.is_available() else \"cpu\")\n",
    "    model = Net(config).to(device)\n",
    "    model.load_state_dict(torch.load(config.model_save_path + config.model_name))   # 加载模型参数\n",
    "    # 先定义一个tensor保存预测结果\n",
    "    result = torch.Tensor().to(device)\n",
    "    # 预测过程\n",
    "    model.eval()\n",
    "    hidden_predict = None\n",
    "    for _data in test_loader:\n",
    "        data_X = _data[0].to(device)\n",
    "        pred_X, hidden_predict = model(data_X, hidden_predict)\n",
    "        cur_pred = torch.squeeze(pred_X, dim=0)\n",
    "        result = torch.cat((result, cur_pred), dim=0)\n",
    "    return result.detach().cpu().numpy()    # 先去梯度信息，如果在gpu要转到cpu，最后要返回numpy数据"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 项目配置模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "outputs": [],
   "source": [
    "class Config:\n",
    "    # 数据参数\n",
    "    feature_columns = list(range(1, 15))\n",
    "    label_columns = [14]\n",
    "    label_in_feature_index = (lambda x, y: [x.index(i) for i in y])(feature_columns, label_columns)\n",
    "    predict_day = 5  # 预测未来多少天\n",
    "    # 网络参数\n",
    "    input_size = len(feature_columns)\n",
    "    output_size = len(label_columns)\n",
    "    hidden_size = 64\n",
    "    lstm_layers = 4\n",
    "    dropout_rate = 0.2\n",
    "    time_step = 10\n",
    "    # 训练参数\n",
    "    do_train = False\n",
    "    do_predict = not do_train\n",
    "    add_train = False\n",
    "    shuffle_train_data = True\n",
    "    use_cuda = True\n",
    "    train_data_rate = 0.95\n",
    "    valid_data_rate = 0.2\n",
    "    batch_size = 256\n",
    "    learning_rate = 0.001\n",
    "    epoch = 3000\n",
    "    patience = 800\n",
    "    random_seed = 42\n",
    "    do_continue_train = False\n",
    "    continue_flag = \"\"\n",
    "    if do_continue_train:\n",
    "        shuffle_train_data = False\n",
    "        batch_size = 1\n",
    "        continue_flag = \"continue_\"\n",
    "    if do_predict:\n",
    "        train_data_rate = 0\n",
    "    # 训练模式\n",
    "    debug_mode = False\n",
    "    debug_num = 500\n",
    "    # 框架参数\n",
    "    used_frame = \"pytorch\"\n",
    "    model_name = \"model_\" + continue_flag + \"pytorch.pth\"\n",
    "    # 路径参数\n",
    "    train_data_path = \"Data.csv\"\n",
    "    model_save_path = \"./checkpoint/\"\n",
    "    figure_save_path = \"./figure/\"\n",
    "    log_save_path = \"./log/\"\n",
    "    do_log_print_to_screen = True\n",
    "    do_log_save_to_file = True\n",
    "    do_figure_save = True\n",
    "    do_train_visualized = False\n",
    "    if not os.path.exists(model_save_path):\n",
    "        os.makedirs(model_save_path)\n",
    "    if not os.path.exists(figure_save_path):\n",
    "        os.mkdir(figure_save_path)\n",
    "    if do_train and (do_log_save_to_file or do_train_visualized):\n",
    "        cur_time = time.strftime(\"%Y_%m_%d_%H_%M_%S\", time.localtime())\n",
    "        log_save_path = log_save_path + cur_time + \"/\"\n",
    "        os.makedirs(log_save_path)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### log日志记录模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "outputs": [],
   "source": [
    "def load_logger(config):\n",
    "    logger = logging.getLogger()\n",
    "    logger.setLevel(level=logging.DEBUG)\n",
    "    # StreamHandler\n",
    "    if config.do_log_print_to_screen:\n",
    "        stream_handler = logging.StreamHandler(sys.stdout)\n",
    "        stream_handler.setLevel(level=logging.INFO)\n",
    "        formatter = logging.Formatter(datefmt='%Y/%m/%d %H:%M:%S', fmt='[ %(asctime)s ] %(message)s')\n",
    "        stream_handler.setFormatter(formatter)\n",
    "        logger.addHandler(stream_handler)\n",
    "    # FileHandler\n",
    "    if config.do_log_save_to_file:\n",
    "        file_handler = RotatingFileHandler(config.log_save_path + \"out.log\", maxBytes=1024000, backupCount=5, encoding='utf-8')\n",
    "        file_handler.setLevel(level=logging.INFO)\n",
    "        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n",
    "        file_handler.setFormatter(formatter)\n",
    "        logger.addHandler(file_handler)\n",
    "        # 把config信息也记录到log 文件中\n",
    "        config_dict = {}\n",
    "        for key in dir(config):\n",
    "            if not key.startswith(\"_\"):\n",
    "                config_dict[key] = getattr(config, key)\n",
    "        config_str = str(config_dict)\n",
    "        config_list = config_str[1:-1].split(\", '\")\n",
    "        config_save_str = \"\\nConfig:\\n\" + \"\\n'\".join(config_list)\n",
    "        logger.info(config_save_str)\n",
    "    return logger"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 绘图模块"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "outputs": [],
   "source": [
    "def draw(config: Config, origin_data: Data, logger, predict_norm_data: np.ndarray):\n",
    "    label_data = origin_data.data[origin_data.train_num + origin_data.start_num_in_test:, config.label_in_feature_index]\n",
    "    predict_data = predict_norm_data * origin_data.std[config.label_in_feature_index] + \\\n",
    "                   origin_data.mean[config.label_in_feature_index]\n",
    "    assert label_data.shape[0] == predict_data.shape[0], \"The element number in origin and predicted data is different\"\n",
    "    label_name = [origin_data.data_column_name[i] for i in config.label_in_feature_index]\n",
    "    label_column_num = len(config.label_columns)\n",
    "    # label 和 predict 是错开config.predict_day天的数据的\n",
    "    loss = np.mean((label_data[config.predict_day:] - predict_data[:-config.predict_day]) ** 2, axis=0)\n",
    "    loss_norm = loss / (origin_data.std[config.label_in_feature_index] ** 2)\n",
    "    logger.info(\"The mean squared error of stock {} is \".format(label_name) + str(loss_norm))\n",
    "    label_X = range(origin_data.data_num - origin_data.train_num - origin_data.start_num_in_test)\n",
    "    predict_X = [x + config.predict_day for x in label_X]\n",
    "    for i in range(label_column_num):\n",
    "        plt.figure(i + 1)\n",
    "        plt.plot(label_X, label_data[:, i], label='真实值', color='red')\n",
    "        plt.plot(predict_X, predict_data[:, i], label='预测值', color='blue')\n",
    "        plt.title(\"{}预测图\".format(label_name[i]), fontname=\"SimHei\")\n",
    "        plt.legend(loc=\"upper left\")\n",
    "        logger.info(\"The predicted stock {} for the next {} day(s) is: \".format(label_name[i], config.predict_day) +\n",
    "                    str(np.squeeze(predict_data[-config.predict_day:, i])))\n",
    "        if config.do_figure_save:\n",
    "            plt.savefig(config.figure_save_path + \"{}predict_{}.png\".format(config.continue_flag, label_name[i]))\n",
    "    plt.show()\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 主文件"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    warnings.filterwarnings(\"ignore\")\n",
    "    plt.style.use('seaborn')\n",
    "    plt.rcParams['font.sans-serif'] = 'Microsoft Yahei'\n",
    "    config = Config()\n",
    "    logger = load_logger(config)\n",
    "    try:\n",
    "        np.random.seed(config.random_seed)\n",
    "        data_gainer = Data(config)\n",
    "        if config.do_train:\n",
    "            train_X, valid_X, train_Y, valid_Y = data_gainer.get_train_and_valid_data()\n",
    "            train(config, logger, [train_X, train_Y, valid_X, valid_Y])\n",
    "        if config.do_predict:\n",
    "            test_X, test_Y = data_gainer.get_test_data(return_label_data=True)\n",
    "            pred_result = predict(config, test_X)\n",
    "            draw(config, data_gainer, logger, pred_result)\n",
    "    except Exception:\n",
    "        logger.error(\"Run Error\", exc_info=True)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}