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