transformer序列预测示例

本文参考:

【python量化】将Transformer模型用于股票价格预测_蚂蚁爱Python的博客-CSDN博客_transformer 股票

一、Transformer初步架构图

transformer序列预测示例_第1张图片

二、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进行序列预测示例

transformer序列预测示例_第2张图片

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)

transformer序列预测示例_第3张图片

transformer序列预测示例_第4张图片

transformer序列预测示例_第5张图片

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的过程如下:

transformer序列预测示例_第6张图片

 四、完整代码(针对原文中错误修改过的)

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())

你可能感兴趣的:(transformer,深度学习,人工智能)