本文参考：
【python量化】将Transformer模型用于股票价格预测_蚂蚁爱Python的博客-CSDN博客_transformer 股票
一、Transformer初步架构图
二、transformer各组件信息
1、positional_encoder
X -> [positional_encoding(embedding + position)] -> X‘
2、multi_head_attention_forward
输入：
Query：(L, N, E)
Key：(S, N, E)
Value：(S, N, E)
输出：
Attn_output：（L，N，E）
Attn_output_weights：（N，L，S）
3、MultiheadAttention
输入：
Embed_dim
Num_heads
Multihead_attn = MultiheadAttention(embed_dim, num_heads)
Attn_output, attn_output_weights = multihead_attn(query, key, value)
4、TransformerEncoderLayer
输入：
d_model，
n_head
5、TransformerEncoder
输入：
Encoder_layer
Num_layers
Encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
Transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
6、TransformerDecoderLayer
输入：
d_model
n_head
7、TransformerDecoder
输入：
Decoder_layer
Num_layers
Decoder_layer = TransformerDecoder(d_model=512, nhead=8)
Transformer_decoder = TransformerDecoder(decoder_layer, num_layers=6)
8、transformer
输入：
Nhead
Num_encoder_layers
Transformer_model = Transformer(nhead=16, num_encoder_layers=12)
Src = torch.rand((10, 32, 512))
Tgt = torch.rand((20, 32, 512))
Out = transformer_model(src, tgt)
三、transformer进行序列预测示例
1、csv输入序列的shape为(126, 1)
2、切分训练集(88, )和测试集(38, )
3、滑动窗口input_window得到(序列值， 标签值)
得到的训练序列为(67, 2, 20)
67：窗口滑动次数；
2：序列+标签=2个
20：输入窗口长度
得到测试序列（17， 2， 20）
17：窗口滑动次数
2：序列+标签=2个
20：输入窗口长度
4、将input和target转成模型需要的格式（S=20，N=2， E=1）
5、模型训练至掩码生成(20, 20)
6、Pos_encoder
Src(20, 2, 1) + positional_encoder(20, 1, 250) => （20， 2， 250）
广播机制完成
7、transformer_encoder：（20， 2， 250）-> (20, 2, 250)
8、decoder：（20， 2， 250） -> (20, 2, 1) 通过Linear完成
9、loss值计算
因为有mask存在，所以每次预测下个值时的后续值是未知的。计算loss的过程如下：
四、完整代码（针对原文中错误修改过的）
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 = 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
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), :]
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 _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
def forward(self, src):
if self.src_mask is None or self.src_mask.shape[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 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)
def get_data():
series = pd.read_csv("D:\\temp\\0001_daily.csv", usecols=[0])
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)
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
def train(train_data):
model.train()
total_loss = 0
for batch_index, i in enumerate(range(0, len(train_data) - 1, batch_size)):
start_time = time.time()
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('transformer-epoch%d.png' % epoch)
plt.close()
return total_loss / i
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 = 60
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()
模型使用示例：
保存模型：
torch.save(model.state_dict(), 'transformer.pth.tar')
使用模型：
import os
model_path = os.path.join(os.getcwd(), 'transformer.pth.tar')
model_dict = torch.load(model_path)
model.load_state_dict(model_dict)
test_data = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0])
test_data = torch.from_numpy(test_data).view(-1, 1, 1).to('cuda').float() # seq_len, batch_size, embedding
test_output = model(test_data)
print(test_output[-1].item())