深度学习算法transformer（时序预测）_transformer预测代码

 transformer代码部分（关于原理网上资料已经很多，这里只展示代码，数据集为开源ETT数据集）

本文采用nn.Transformer，nn.Transformer没有加入embedding、位置embedding、linear+softmax

本文使用的是时序数据，embedding使用的是nn.Linear；位置embedding使用的是传统常见的位置embedding，详情见代码；nn.Transformer之后再接一个nn.Linear做预测

encoder的输入序列长度是seq_length = 96，decoder的输入序列长度label_length + pred_length = 72

讲的比较好的transformer原理Transformer模型详解（图解最完整版） - 知乎

1. 导入必须要的包

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import plotly.express as px
from sklearn import metrics
from torch.utils.data import Dataset, DataLoader

2. 定义transformer网络

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, device, max_len=5000):
        super(PositionalEncoding, self).__init__()
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.pe = pe.unsqueeze(0).transpose(1, 0).to(device)
 
    def forward(self, x):
        return self.pe[:x.size(0), :]
 
class TransformerTimeSeriesModel(nn.Module):
    def __init__(self, input_size, d_model, device, nhead,
                 num_encoder_layers, num_decoder_layers, dim_feedforward, dropout=0.1):
        super(TransformerTimeSeriesModel, self).__init__()
 
        self.value_encoding = nn.Linear(input_size, d_model)
 
        self.positional_encoding = PositionalEncoding(d_model, device)
        self.transformer = nn.Transformer(d_model=d_model, nhead=nhead,
                                          num_encoder_layers=num_encoder_layers,
                                          num_decoder_layers=num_decoder_layers,
                                          dim_feedforward=dim_feedforward,
                                          dropout=dropout)
        self.fc_out = nn.Linear(d_model, 1)
 
    def forward(self, src, tgt, device, tgt_mask=None):
 
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(0)).to(device)
 
        src = self.value_encoding(src) + self.positional_encoding(src)
        tgt = self.value_encoding(tgt) + self.positional_encoding(tgt)
 
        output = self.transformer(src, tgt, tgt_mask=tgt_mask)
 
        return self.fc_out(output)

3. 定义有关Dataset的数据整理

# 定义Dataset
class get_dataset(Dataset):
    def __init__(self, data_path, seq_length, label_length, pred_length,
                 features, train_split, mode):
        self.mode = mode
        self.data_path = data_path
        self.features = features
        self.seq_length = seq_length
 
        self.label_length = label_length
        self.pred_length = pred_length
 
        self.data, self.data_max, self.data_min = self.get_data()
        #         print(self.data)
        #         print(self.data[0, :-1, :])
        #         print(self.data[0, -1, -1])
        #         print(self.data[0, -1, -1].unsqueeze(0))
        #         print(self.data[0, -1, -1].unsqueeze(0).unsqueeze(1))
        train_num = int(train_split * len(self.data))
        if self.mode == 'train':
            self.data = self.data[:train_num, :, :]
        else:
            self.data = self.data[train_num:, :, :]
 
    def __len__(self):
        return len(self.data)
 
    def __getitem__(self, index):
        return self.data[index, :-self.pred_length, :], \
               self.data[index, (self.seq_length - self.label_length):, :], \
               self.data[index, (self.seq_length - self.label_length):, -1].unsqueeze(1)
 
    def get_data(self):
        data = pd.read_csv(self.data_path)
        data.index = pd.to_datetime(data['date'])
        data = data.drop('date', axis=1)
        data_max = data.max()
        data_min = data.min()
 
        data = (data - data_min) / (data_max - data_min)
        num_sample = len(data) - self.seq_length - self.pred_length + 1
        seq_data = torch.zeros(num_sample,
                               self.seq_length + self.pred_length,
                               len(self.features))
 
        #         print(data.iloc[0:0 + self.seq_length + 1, self.features].values)
 
        for i in range(num_sample):
            seq_data[i] = torch.tensor(data.iloc[i:i + self.seq_length + self.pred_length,
                                       self.features].values)
        #         print(data_max)
        #         print(data_min)
 
        return seq_data, data_max, data_min

4. 定义训练

def train(model, dataset, epochs, optim, loss_function, device, batch_size, shuffle=True):
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
 
    for epoch in range(epochs):
        train_loss = 0
        model.train()
        for x, y, label in data_loader:
            x, y, label = x.transpose(1, 0).to(device), 
                          y.transpose(1, 0).to(device), 
                          label.transpose(1, 0).to(device)
            #             print('x', x.shape)
            #             print('y', y.shape)
            pred = model(x, y, device)
            #             print('pred', pred)
            #             print(pred.shape)
 
            loss = loss_function(pred, label)
 
            optim.zero_grad()
            loss.backward()
            optim.step()
            train_loss += loss.item()
        train_loss /= len(data_loader)
        print('epoch / epochs : %d / %d, loss : %.6f' % (epoch, epochs, train_loss))

5. 定义测试

def test(model, dataset, device, batch_size, label_length, pred_length, root_path, shuffle=False):
 
    model.eval()
 
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
 
    # print('data_loader : ', len(data_loader))
    # print('dataset : ', len(dataset))
 
    preds, labels = np.zeros(len(dataset) * (label_length + pred_length)), \
                    np.zeros(len(dataset) * (label_length + pred_length))
    left, right = 0, 0
 
    for x, y, label in data_loader:
 
        left = right
 
        # if len(label) != 32:
        #     print('--')
 
        right += len(label) * (label_length + pred_length)
        x, y = x.transpose(1, 0).to(device), y.transpose(1, 0).to(device)
 
        pred = model(x, y, device).detach().cpu().numpy().flatten()
 
        # print('right:', right)
        # print('label : ', label.flatten())
        # print('pred : ', pred)
        # print(label.flatten().shape)
        # print(pred.shape)
        preds[left:right] = pred
        labels[left:right] = label.transpose(1, 0).detach().cpu().numpy().flatten()
 
    preds_ = preds * (dataset.data_max['OT'] - dataset.data_min['OT']) + dataset.data_min['OT']
    labels_ = labels * (dataset.data_max['OT'] - dataset.data_min['OT']) + dataset.data_min['OT']
 
    np.save(root_path + '_preds.npy', preds)
    np.save(root_path + '_labels.npy', labels)
    return preds_, labels_

6. 定义模型评估指标

def model_eva(pred, label):
    fig = px.line(title='transformer模型预测')
    fig.add_scatter(y=label, name='label')
    fig.add_scatter(y=pred, name='pred')
    fig.show()
 
    #     print(label)
    #     print(pred)
    #     label_nozero = labels[labels == 0] = 1e-3
    index = np.where(label > 0.01)
 
    mse = np.mean((label - pred) ** 2)
 
    r2 = 1 - np.sum((label - pred) ** 2) / np.sum((label - np.mean(label)) ** 2)
    mape = np.abs((pred[index] - label[index]) / label[index]).mean()
    mae = np.abs(label - pred).mean()
 
    print('MSE : %.6f' % (mse))
    print('R2 : %.6f' % (r2))
    print('MAPE : %.6f' % (mape))
    print('MAE : %.6f' % (mae))

7. main函数以及运行结果

seed = 0
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  
np.random.seed(seed)  
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
 
 
device = 'cuda' if torch.cuda.is_available() else 'cpu'
seq_length = 96
label_length = 48
pred_length = 24
features = [6]                 # [HUFL,HULL,MUFL,MULL,LUFL,LULL,OT]
input_size = len(features)
epochs = 100
lr = 0.005
batch_size = 32
train_split = 0.8
d_model = 128
nhead = 2
num_encoder_layers = 2
num_decoder_layers = 2
dim_feedforward = 128
dropout = 0.1
root_path = './' + 'seq_' + str(seq_length) + '_label_' + str(label_length) + '_pred_' + str(pred_length)
save_path = root_path + '_transformer.pth'
 
 
model = TransformerTimeSeriesModel(input_size, d_model, device, nhead,
                                   num_encoder_layers, num_decoder_layers,
                                   dim_feedforward, dropout=0.1).to(device)
    
optim = torch.optim.SGD(model.parameters(), lr=lr)
loss_function = nn.MSELoss()
 
dataset_train = get_dataset(data_path, seq_length, label_length, pred_length, features, train_split = train_split, mode = 'train')
dataset_test = get_dataset(data_path, seq_length, label_length, pred_length, features, train_split = train_split, mode = 'test')
 
 
train(model, dataset_train, epochs, optim, loss_function, device, batch_size, shuffle = True)
 
torch.save(model.state_dict(), save_path)
 
preds, labels = test(model, dataset_test, device, batch_size, label_length, pred_length, root_path, shuffle=False)
 
model_eva(preds, labels)