transformer

import torch
import math
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
from tqdm import tqdm
import torch.nn as nn
import numpy as np

max_len=96
input_size = 5  # 特征数
# PositionalEncoding保持不变
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=96):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        # 初始化 dropout 层，用于在前向传播时应用 dropout
        position = np.arange(max_len).reshape(-1, 1)
        # 生成用于位置编码的除数，注意这里的计算对象应是模型维度的一半
        div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
        # 初始化位置编码矩阵为零矩阵
        pe = np.zeros((max_len, d_model))
        # 应用正弦函数于偶数索引处的维度
        pe[:, 0::2] = np.sin(position * div_term)
        # 应用余弦函数于奇数索引处的维度
        pe[:, 1::2] = np.cos(position * div_term)
        pe = torch.from_numpy(pe).float()  # 将 NumPy 数组转换为 PyTorch 张量，并转换数据类型为 float
        self.register_buffer('pe', pe)  # 将位置编码矩阵注册为模块的缓冲区
    def forward(self, x):
        x = x + self.pe[:x.size(0), :].unsqueeze(0).transpose(0, 1)  # 将位置编码矩阵与输入张量相加，根据输入张量的长度截取相应长度的位置编码
        return self.dropout(x)  # 对添加位置编码后的张量应用 dropout 后返回结果

class TransformerModel(nn.Module):
    def __init__(self, input_size, d_model, nhead, num_encoder_layers, dim_feedforward, max_len, pred_steps,
                 dropout=0.1):
        super(TransformerModel, self).__init__()
        self.pred_steps = pred_steps
        self.pos_encoder = PositionalEncoding(d_model, dropout, max_len)
        encoder_layers = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_encoder_layers)
        self.encoder = nn.Linear(input_size, d_model)
        self.d_model = d_model
        self.decoder = nn.Linear(d_model, input_size)  # 单时间步的输出维度
        # 添加额外的线性层用于处理拼接后的输出
        self.final_linear = nn.Linear(max_len*pred_steps*input_size, input_size)

    def forward(self, src):
        src = self.encoder(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)

        # 解码每个时间步的输出
        output = self.decoder(output)

        # 拼接解码后的所有时间步向量 (batch_size, sequence_length, features_per_step)
        output = output.contiguous().view(output.size(0), -1)  # (batch_size, sequence_length*features_per_step)

        # 通过最终的线性层调整维度
        output = self.final_linear(output)

        # 重新调整形状以匹配预期的输出形状 (batch_size, pred_steps, input_size)
        output = output.view(-1, self.pred_steps, input_size)

        return output