AI时序预测: iTransformer算法代码深度解析

在之前的文章中，我对iTransformer的Paper进行了详细解析，具体文章如下：

文章链接：深度解析iTransformer：维度倒置与高效注意力机制的结合

今天，我将对iTransformer代码进行解析。回顾Paper，我们知道iTransformer通过简单地将注意力机制和前馈网络应用于倒置的维度上。具体而言，单个序列的时间点被嵌入为变量令牌（variate tokens），并利用注意力机制捕捉变量间的相关性；同时，前馈网络被应用于每个变量令牌，以学习非线性表示。

iTransformer 整体架构如下图所示，采用了 Transformer（Vaswani et al., 2017）的仅编码器（encoder-only）架构，包括嵌入层、投影层和 Transformer 块。接下来，我们看看每个模块是如何通过代码实现的。

1. 嵌入层

import torch
import torch.nn as nn

class DataEmbedding_inverted(nn.Module):
    """
    该类用于数据嵌入（Embedding），适用于时间序列建模或其他需要将输入转换为高维表示的任务。
    它通过线性变换将输入数据映射到 `d_model` 维度，并可选地结合额外的时间信息进行处理。

    参数：
    - c_in (int): 输入特征的维度（即变量数）。
    - d_model (int): 目标嵌入维度（即转换后的特征维度）。
    - embed_type (str, 可选): 嵌入类型，当前代码未使用该参数，默认值为 'fixed'。
    - freq (str, 可选): 频率信息（如 'h' 代表小时级别），当前代码未使用该参数。
    - dropout (float, 可选): Dropout 比例，控制神经元随机失活的概率，以防止过拟合。
    """

    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
        super(DataEmbedding_inverted, self).__init__()

        # 线性映射层：将输入数据从 `c_in` 维度投影到 `d_model` 维度
        self.value_embedding = nn.Linear(c_in, d_model)

        # Dropout 层：用于在训练时随机丢弃部分神经元，以增强模型的泛化能力
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, x_mark):
        """
        前向传播函数，将输入 `x` 及可选的时间标记 `x_mark` 进行处理，返回嵌入后的表示。

        参数：
        - x (Tensor): 输入数据，形状为 [Batch, Time, Variate]，其中：
            - Batch: 批量大小
            - Time: 时间步数
            - Variate: 变量数（即 `c_in`）
        - x_mark (Tensor 或 None): 时间标记信息，形状为 [Batch, Time, Extra_Features]，
          如果为 None，则仅使用 `x` 进行嵌入。

        返回：
        - Tensor: 经过嵌入和 Dropout 处理后的数据，形状为 [Batch, Variate, d_model]。
        """

        # 交换 `Time` 和 `Variate` 维度，调整形状以适配后续处理
        x = x.permute(0, 2, 1)  # 变换后形状：[Batch, Variate, Time]

        if x_mark is None:
            # 仅使用输入数据 `x` 进行嵌入
            x = self.value_embedding(x)
        else:
            # 如果提供了 `x_mark`，先调整其形状，再与 `x` 拼接后进行嵌入
            x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], dim=1))

        # 经过 Dropout 处理后返回最终的嵌入表示
        return self.dropout(x)

完整文章链接：AI时序预测: iTransformer算法代码深度解析