下面的這篇文章主要教大家如何搭建一個基於Transformer的簡單預測模型，並將其用於股票價格預測當中。原代碼在文末進行獲取。
Transformer 是 Google 的團隊在 2017 年提出的一種 NLP 經典模型，現在比較火熱的 Bert 也是基於 Transformer。Transformer 模型使用了 Self-Attention 機制，不采用 RNN 的順序結構，使得模型可以並行化訓練，而且能夠擁有全局信息。這篇文章的目的主要是將帶大家通過Pytorch框架搭建一個基於Transformer的簡單股票價格預測模型。
具體地，我們用到了上證指數的收盤價數據為例，進行預測t+1時刻的收盤價。需要注意的是，本文只是通過這樣一個簡單的基本模型，帶大家梳理一下數據預處理，模型構建以及模型評估的流程。模型還有很多可以改進的地方，例如選擇更有意義的特征，如何進行有效的多步預測等。
本地環境：
Python 3.7
IDE:Pycharm
庫版本：
numpy 1.18.1
pandas 1.0.3
sklearn 0.22.2
matplotlib 3.2.1
torch 1.10.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)
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), :]
之後，搭建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.pos_encoder(src)
output = self.transformer_encoder(src, self.src_mask)
output = self.decoder(output)
return output
def _generate_square_subsequent_mask(self, sz):
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
接下來需要對數據進行預處理，首先定義一個窗口劃分的函數。它的作用是將輸入按照延遲output_windw的方式來劃分數據以及其標簽，文中是進行單步預測，所以假設輸入是1到20，則其標簽就是2到21，以適應Transformer的seq2seq的形式的輸出。
def create_inout_sequences(input_data, tw):
inout_seq = []
L = len(input_data)
for i in range(L - tw):
train_seq = input_data[i:i + tw]
train_label = input_data[i + output_window:i + tw + output_window]
inout_seq.append((train_seq, train_label))
return torch.FloatTensor(inout_seq)
之後劃分訓練集和測試集，其中前70%條數據用於模型訓練，後面的數據用於模型測試。具體地，我們用到了前input_window個收盤價來預測下一時刻的收盤價數據。
def get_data():
series = pd.read_csv('./000001_Daily.csv', usecols=['Close'])
# series = pd.read_csv('./daily-min-temperatures.csv', usecols=['Temp'])
scaler = MinMaxScaler(feature_range=(-1, 1))
series = scaler.fit_transform(series.values.reshape(-1, 1)).reshape(-1)
train_samples = int(0.7 * len(series)) train_data = series[:train_samples]
test_data = series[train_samples:]
train_sequence = create_inout_sequences(train_data, input_window)
train_sequence = train_sequence[:-output_window]
test_data = create_inout_sequences(test_data, input_window)
test_data = test_data[:-output_window]
return train_sequence.to(device), test_data.to(device)
接下來實現一個databatch generator，便於從數據中按照batch的形式進行讀取數據。
def get_batch(source, i, batch_size):
seq_len = min(batch_size, len(source) - 1 - i)
data = source[i:i + seq_len]
input = torch.stack(torch.stack([item[0] for item in data]).chunk(input_window, 1))
target = torch.stack(torch.stack([item[1] for item in data]).chunk(input_window, 1))
return input, target
下面是模型訓練的代碼。具體地，就是通過遍歷訓練集，通過既定的loss，對參數進行反向傳播，其中用到了梯度裁剪的技巧用於防止梯度爆炸，然後每間隔幾個間隔打印一下loss。
def train(train_data):
model.train()
for batch_index, i in enumerate(range(0, len(train_data) - 1, batch_size)):
start_time = time.time()
total_loss = 0
data, targets = get_batch(train_data, i, batch_size)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.7)
optimizer.step()
total_loss += loss.item()
log_interval = int(len(train_data) / batch_size / 5)
if batch_index % log_interval == 0 and batch_index > 0:
cur_loss = total_loss / log_interval
elapsed = time.time() - start_time
print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.6f} | {:5.2f} ms | loss {:5.5f} | ppl {:8.2f}'
.format(epoch, batch_index, len(train_data) // batch_size, scheduler.get_lr()[0], elapsed * 1000 / log_interval, cur_loss, math.exp(cur_loss)))
接下來是對模型進行評估的代碼。
def evaluate(eval_model, data_source):
eval_model.eval()
total_loss = 0
eval_batch_size = 1000
with torch.no_grad():
for i in range(0, len(data_source) - 1, eval_batch_size):
data, targets = get_batch(data_source, i, eval_batch_size)
output = eval_model(data)
total_loss += len(data[0]) * criterion(output, targets).cpu().item()
return total_loss / len(data_source)
最後，是模型運行過程的可視化。
def plot_and_loss(eval_model, data_source, epoch):
eval_model.eval()
total_loss = 0.
test_result = torch.Tensor(0)
truth = torch.Tensor(0)
with torch.no_grad():
for i in range(0, len(data_source) - 1):
data, target = get_batch(data_source, i, 1)
output = eval_model(data)
total_loss += criterion(output, target).item()
test_result = torch.cat((test_result, output[-1].view(-1).cpu()), 0)
truth = torch.cat((truth, target[-1].view(-1).cpu()), 0)
plt.plot(test_result, color="red") plt.plot(truth, color="blue")
plt.grid(True, which='both')
plt.axhline(y=0, color='k')
plt.savefig('graph/transformer-epoch%d.png' % epoch)
plt.close()
return total_loss / i
最後，對模型進行運行。其中用到了mse作為loss，adam作為優化器，以及設定學習率的調度器，最後運行200個epoch，每隔10個epoch在測試集上評估一下模型。
train_data, val_data = get_data()
model = TransAm().to(device)
criterion = nn.MSELoss()
lr = 0.005
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.95)
epochs = 200
for epoch in range(1, epochs + 1):
epoch_start_time = time.time()
train(train_data)
if (epoch % 10 is 0):
val_loss = plot_and_loss(model, val_data, epoch)
else:
val_loss = evaluate(model, val_data)
print('-' * 89)
print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.5f} | valid ppl {:8.2f}'.format(epoch, ( time.time() - epoch_start_time), val_loss, math.exp(val_loss)))
print('-' * 89) scheduler.step()
下面是運行的結果，可以看到loss明顯降低了。
cuda| epoch 1 | 2/ 10 batches | lr 0.005000 | 7.83 ms | loss 39.99368 | ppl 233902099994043520.00| epoch 1 | 4/ 10 batches | lr 0.005000 | 7.81 ms | loss 7.20889 | ppl 1351.39| epoch 1 | 6/ 10 batches | lr 0.005000 | 11.10 ms | loss 1.68758 | ppl 5.41| epoch 1 | 8/ 10 batches | lr 0.005000 | 9.35 ms | loss 0.00833 | ppl 1.01| epoch 1 | 10/ 10 batches | lr 0.005000 | 7.81 ms | loss 1.18041 | ppl 3.26-----------------------------------------------------------------------------------------| end of epoch 1 | time: 1.96s | valid loss 2.58557 | valid ppl 13.27
...
| end of epoch 198 | time: 0.30s | valid loss 0.00032 | valid ppl 1.00-----------------------------------------------------------------------------------------| epoch 199 | 2/ 10 batches | lr 0.000000 | 15.62 ms | loss 0.00057 | ppl 1.00| epoch 199 | 4/ 10 batches | lr 0.000000 | 15.62 ms | loss 0.00184 | ppl 1.00| epoch 199 | 6/ 10 batches | lr 0.000000 | 15.62 ms | loss 0.00212 | ppl 1.00| epoch 199 | 8/ 10 batches | lr 0.000000 | 7.81 ms | loss 0.00073 | ppl 1.00| epoch 199 | 10/ 10 batches | lr 0.000000 | 7.81 ms | loss 0.00057 | ppl 1.00-----------------------------------------------------------------------------------------| end of epoch 199 | time: 0.30s | valid loss 0.00032 | valid ppl 1.00-----------------------------------------------------------------------------------------| epoch 200 | 2/ 10 batches | lr 0.000000 | 15.62 ms | loss 0.00053 | ppl 1.00| epoch 200 | 4/ 10 batches | lr 0.000000 | 7.81 ms | loss 0.00177 | ppl 1.00| epoch 200 | 6/ 10 batches | lr 0.000000 | 7.81 ms | loss 0.00224 | ppl 1.00| epoch 200 | 8/ 10 batches | lr 0.000000 | 15.62 ms | loss 0.00069 | ppl 1.00| epoch 200 | 10/ 10 batches | lr 0.000000 | 7.81 ms | loss 0.00049 | ppl 1.00-----------------------------------------------------------------------------------------| end of epoch 200 | time: 0.62s | valid loss 0.00032 | valid ppl 1.00-----------------------------------------------------------------------------------------
最後是模型的擬合效果，從實驗結果中可以看出我們搭建的簡單的Transformer模型可以實現相對不錯的數據擬合效果。
在這篇文章中，我們介紹了如何基於Pytorch框架搭建一個基於Transformer的股票預測模型，並通過真實股票數據對模型進行了實驗，可以看出Transformer模型對股價預測具有一定的效果。另外，文中只是做了一個簡單的demo，其中仍然有很多可以改進的地方，如采用更多有意義的輸入數據，優化其中的一些組件等。除此之外，目前基於Transformer的模型層出不窮，其中也有很多值得我們去學習，大家也可以采用更先進的Transformer模型進行試驗。
先自我介紹一下，小編13年上師交大畢業，曾經在小公司待過，去過華為OPPO等大廠，18年進入阿裡，直到現在。深知大多數初中級java工程師，想要升技能，往往是需要自己摸索成長或是報班學習，但對於培訓機構動則近萬元的學費，著實壓力不小。自己不成體系的自學效率很低又漫長，而且容易碰到天花板技術停止不前。因此我收集了一份《java開發全套學習資料》送給大家，初衷也很簡單，就是希望幫助到想自學又不知道該從何學起的朋友，同時減輕大家的負擔。添加下方名片，即可獲取全套學習資料哦
