【python量化】将Transformer模型用于股票价格预测_transformer进行股票预测

前言

下面的这篇文章主要教大家如何搭建一个基于Transformer的简单预测模型，并将其用于股票价格预测当中。原代码在文末进行获取。

1、Transformer模型

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。这篇文章的目的主要是将带大家通过Pytorch框架搭建一个基于Transformer的简单股票价格预测模型。

Transformer的基本架构

具体地，我们用到了上证指数的收盘价数据为例，进行预测t+1时刻的收盘价。需要注意的是，本文只是通过这样一个简单的基本模型，带大家梳理一下数据预处理，模型构建以及模型评估的流程。模型还有很多可以改进的地方，例如选择更有意义的特征，如何进行有效的多步预测等。

2、环境准备

本地环境：

Python 3.7
IDE:Pycharm
1
2

库版本：

numpy 1.18.1
pandas 1.0.3 
sklearn 0.22.2
matplotlib 3.2.1
torch 1.10.1
1
2
3
4
5

3、代码实现

3.1. 导入库以及定义超参

首先，需要导入用到库，以及模型的一些超参数的设置。其中，input_window和output_window分别用于设置输入数据的长度以及输出数据的长度。当然，这些参数大家也可以根据实际应用场景进行修改。

Python学习交流Q群：906715085###
import torch
import torch.nn as nn
import numpy as np
import time
import math
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
import pandas as pd

torch.manual_seed(0)
np.random.seed(0)

input_window = 20
output_window = 1
batch_size = 64
device = torch.
device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

3. 2. 模型构建

Transformer中很重要的一个组件是提出了一种新的位置编码的方式。我们知道，循环神经网络本身就是一种顺序结构，天生就包含了词在序列中的位置信息。当抛弃循环神经网络结构，完全采用Attention取而代之，这些词序信息就会丢失，模型就没有办法知道每个词在句子中的相对和绝对的位置信息。因此，有必要把词序信号加到词向量上帮助模型学习这些信息，位置编码（PositionalEncoding）就是用来解决这种问题的方法。它的原理是将生成的不同频率的正弦和余弦数据作为位置编码添加到输入序列中，从而使得模型可以捕捉输入变量的相对位置关系。

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, max_len=5000):     
       super(PositionalEncoding, self).__init__()       
        pe = torch.zeros(max_len, d_model)      
          position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)       
           div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))     
              pe[:, 0::2] = torch.sin(position * div_term)       
               pe[:, 1::2] = torch.cos(position * div_term)       
                pe = pe.unsqueeze(0).transpose(0, 1)      
                  self.register_buffer('pe', pe)
    def forward(self, x):        
    return x + self.pe[:x.size(0), :]
1
2
3
4
5
6
7
8
9
10
11
12
13

之后，搭建Transformer的基本结构，在Pytorch中有已经实现的封装好的Transformer组件，可以很方便地进行调用和修改。其中需要注意的是，文中并没有采用原论文中的Encoder-Decoder的架构，而是将Decoder用了一个全连接层进行代替，用于输出预测值。另外，其中的create_mask将输入进行mask，从而避免引入未来信息。

class TransAm(nn.Module):   
 def __init__(self, feature_size=250, num_layers=1, dropout=0.1):       
  super(TransAm, self).__init__()        
  self.model_type = 'Transformer'        
  self.src_mask = None        
  self.pos_encoder = PositionalEncoding(feature_size)        
  self.encoder_layer = nn.TransformerEncoderLayer(d_model=feature_size, nhead=10, dropout=dropout)        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)        
  self.decoder = nn.Linear(feature_size, 1)        
  self.init_weights()
    def init_weights(self):        
    initrange = 0.1        
    self.decoder.bias.data.zero_()        
    self.decoder.weight.data.uniform_(-initrange, initrange)
    def forward(self, src):       
     if self.src_mask is None or self.src_mask.size(0) != len(src):         
        device = src.device           
         mask = self._generate_square_subsequent_mask(len(src)).to(device)           
          self.src_mask = mask
        src = self
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18