From Known to Unknown: Knowledge-guided Transformer for Time-Series Sales Forecasting Aliformer模型

Aliformer

\\

paper：From Known to Unknown: Knowledge-guided Transformer for Time-Series Sales Forecasting in Alibaba

\\

理论讲解部分

paper任务是基于历史的数据销售数据，对未来的销售数据进行预测，paper发现因为原来的模型都是基于历史的信息数据对未来进行预测，但是例如产品销量的预测，未来的一些促销信息会对当前的销量产生一个很大的影响，例如在促销之前的一小段时间，产品销量会有一个较大幅度的下跌，而促销开始，产品的销量会有一个较大的涨幅，这些原来的模型都无法很好的预测到，但是例如促销等信息都是可以提前知道的，因此提出了Aliformer，会基于未来的信息和历史的数据和信息对当前的进行预测。

Aliformer的网络流程：

Aliformer模型的主要部分在于AliAttention机制。

AliAttention

下图为AliAttention流程图

其中
$$x^{h}and\overline{x},x表示的是统计数据和知识信息，\overline{x}表示的是知识信息$$
shape 都是(batch_size,seq_len+pred_len,embeded_dim)

xh->v,k,q的矩阵的shape 为（embeded_dim,d_ff）,d_ff为一个更高的维度，或者是同维度变换也可以

x的计算公式如下，并且T表示的是已知的时间序列的长度，L表述需要预测的时间序列的长度，
$$\begin{gathered} S_t^{(n)}表示的是统计信息，k_t^{(n)}表示的是知识信息，u_t^{(n)}表示的是未来的统计信息, \\ 因为未来的统计信息是未知的，因此使用一个token来代表，这是一个可学习的参 \\ 或者是一个默认的值 \end{gathered}$$

通过三个全连接层得到了
$$注意力机制的Q,K,V,\overline{K},\overline{Q}$$
并且按照流程图送入多头的注意力机制中，分别得到
$$\begin{gathered} Attand\overline{Att}其中Att表的是统计信息和知识信息的注意力 \\ \overline{Att}表示的是纯知识信息的注意力 \end{gathered}$$
下图为具体的注意力机制的计算过程，其中d表示的就是
$$x_i^h\ast W_Q^x的结果的最后一个维度，d是一个规模因子$$

span masking

但是由于我们是使用一个token去替代未来的统计数据，因此这本身就是对未来数据的一个有偏差的估计，因此模型在训练的时候会更加依赖历史的知识数据和统计数据，但是这不是我们想要的，因此提出了一个跨度掩码的策略来使得模型更加的强调未来的统计信息和数据。

跨度掩码（span masking）在训练中会随机掩盖掉时间序列的中间一段序列，然后根据被mask掉的数据两边的数据进行预测被mask掉的数据。

代码讲解部分

任务简介

该任务是一个预测电力变压器的油温等的任务，因此知识信息就是时间信息

输入数据维度为 (batch_size,seq_len+pred_len,d_feature)，其中batch_size为批量大小

seq_len为已知时间序列的长度,pred_len为需要预测的时间长度，d_feature为特征维度

代码的流程

①先进入Aliformer的forward部分，然后处理数据

②进入AliAttention部分，然后在进入多头注意力机制的计算部分

③然后在返回到Aliformer继续进入AliAttention（堆叠了12层的AliAttention部分）

代码中的名称

$$\begin{gathered} x-fei表示的是\overline{x},k-fei表示的是\overline{k},q-fei表示的是\overline{q} \\ x-knowledge表示的也是\overline{x} \end{gathered}$$

初始化Aliformer

args参数是一个分装好的参数，里面包含了已知的时间序列的长度，需要预测的时间序列的长度，输入模型特征维度，embedding后特征维度，时间维度，升维后的维度d_ff，这一个维度可以自行设定

初始化Aliformer的参数

class Aliformer(nn.Module):
    def __init__(self,args,h=8,p2=0.5,loss_fn=torch.nn.MSELoss()):
        '''
        Args:
        :param h:头数
        :param d_feature:输入数据的最后一个维度dim
        :param d_mark:知识信息的维度（本数据集为时间信息的维度，4）
        :param d_model:embedding后的维度
        :param d_ff:升维后的维度
        d_ff % h==0
        :param p2: 在训练的时候使用跨度掩码的概率为p2
        :param loss_fn 使用的loss函数
        '''
        super().__init__()
        self.h=h
        self.d_feature=args.d_feature
        self.d_mark=args.d_mark
        self.d_model=args.d_model
        self.d_ff=args.d_ff
        self.p2=p2
        self.pred_len=args.pred_len
        self.label_len=args.label_len
        self.seq_len=args.seq_len

初始化embedding层

time_embeded是对时间维度进行embedding，把时间维度从d_mark–>d_model

embeded是把时间信息和数据一起进行embedding，输出的维度为d_model

  # embedding层
        self.time_embeded = TimeEmbedding(d_mark=self.d_mark, d_model=self.d_model)  # 对时间维度进行embedding
        self.embeded =  DataEmbedding_time_token(d_feature=self.d_feature, d_mark=self.d_mark, d_model=self.d_model)  # 把原始数据和时间维度embedding并且加在一起

初始化AliAttention层

self.aliattention=AliAttention(h=self.h,d_feature=self.d_feature,
                                       d_mark=self.d_mark,d_model=self.d_model,d_ff=self.d_ff) # 初始化AliAttention层

pred_len 就是预测的时间序列的长度，seq_len就是已知的时间序列的长度，label_len其实只是为了在后面 Aliformer的forward的变量y中把pred_len部分的数据切分出来，因为y是label_len+pred_len,label_len就是取seq_len的后面的部分长度为label_len的部分的数据

        self.pred_len=args.pred_len
        self.label_len=args.label_len
        self.seq_len=args.seq_len

初始化FC层，为了将预测的数据在特征维度和时间维度上进行降维

    self.out_dim=nn.Linear(self.d_model,self.d_feature) # 对特征维度进行降纬
    self.out_time=nn.Linear((self.seq_len+self.pred_len),self.pred_len) # 对时间维度进行降维

这里是为了span masking而设置的，因为有p2的概率是使用span masking方式的，因此先初始化一个有100个零的列表，然后把其中100*p2个设置为1，然后在这个列表中随机抽，如果抽中为1，那么就使用span masking ，否则就不使用span masking

       self.choice=torch.zeros((100))
        # 为了后面的选择是否使用span masking 策略,如果选择为1，那么就使用span mask策略
        self.choice[:int(p2*100)]=1 # 把p2*100的数字变成1，剩下的是0，然后在这一个列表中随机抽样
        self.loss_fn=loss_fn

初始化AliAttention的网络结构

初始化AliAttention都是按照AliAttention的流程图中的结构来初始化的

初始AliAttention中使用的参数

class AliAttention(nn.Module):
    def __init__(self,h=8,d_feature=7,d_mark=4,d_model=64,d_ff=128):
        '''
        Args:
        :param h:头数
        :param d_feature:输入数据的最后一个维度dim
        :param d_mark:知识信息的维度（本数据集为时间信息的维度，4）
        :param d_model:embedding后的维度
        :param d_ff:升为后的维度
        d_ff % h==0
        '''
        assert d_ff %h==0
        super().__init__()
        self.h=h
        self.d_feature=d_feature
        self.d_mark=d_mark
        self.d_model=d_model
        self.d_ff=d_ff
        self.dk=self.d_ff//self.h # 多头注意力机制中每一个头的维度

# x -->q,k,v
# 输入为以下三个输入输出的维度都是相同的（batch_size,seq_len+pred_len,d_model)-->输出为（batch_size,seq_len+pred_len,d_ff)
self.fc_x_to_v = nn.Linear(self.d_model, self.d_ff)
self.fc_x_to_k = nn.Linear(self.d_model, self.d_ff)
self.fc_x_to_q=nn.Linear(self.d_model,self.d_ff)

# 知识信息x的非-->k的非，Q的非
# 以下两个的输入输出的维度都是相同的（batch_size,seq_len+pred_len,d_model)-->输出为（batch_size,seq_len+pred_len,d_ff)
self.fc_knowledge_x_to_k=nn.Linear(self.d_model,self.d_ff)
self.fc_knowledge_x_to_q = nn.Linear(self.d_model, self.d_ff)

# 多头注意力机制,同维度变换，然后再reshape为，h,dk,(d_ff=dk*h)，就相当于经过h个全连接层在拼接起来
# 以下四个的输入输出的shape都是相同的 batch_size,seq_len+pred_len,d_ff)-->输出为（batch_size,seq_len+pred_len,d_ff)
self.fc_q=nn.Linear(self.d_ff,self.d_ff)
self.fc_k=nn.Linear(self.d_ff,self.d_ff)
self.fc_k_fei=nn.Linear(self.d_ff,self.d_ff)
self.fc_q_fei=nn.Linear(self.d_ff,self.d_ff)

# 输入为（batch_size,seq_len+pred_len,d_ff)-->输出为（batch_size,seq_len+pred_len,d_model)
self.fc_out=nn.Linear(self.d_ff,self.d_model) # 将数据的维度变回输入进attention的维度

AliAttention的forward

首先先通过x是统计数据和知识信息embedding后的数据，embedding部分在Aliformer中

def forward(self,x,x_knowledge):
    '''
    Args:
    :param x: shape是(batch_size,seq_len+pred_len,d_model)
    :param x_knowledge: shape是(batch_size,seq_len+pred_len,mark),知识信息，时间维度
    :return:out: shape(batch_size,seq_len+pred_len,d_model)
    '''
    # 得到q，k，v
    v=self.fc_x_to_v(x)
    k=self.fc_x_to_k(x)
    q=self.fc_x_to_q(x) # v,q,v的shape都是(batch_size,seq_len+pred_len,d_ff)
    # 得到知识信息的q，k，称为q_fei,k_fei
    # k_fei,q_fei shape(batch_size,seq_len+pred_len,d_ff)
    k_fei=self.fc_knowledge_x_to_k(x_knowledge)
    q_fei=self.fc_knowledge_x_to_q(x_knowledge)
    # 输入进attention中得到输出
    out=self.attention(q,k,v,q_fei,k_fei)
    return out

定义多头Attention函数

其中输入的q , k , v , 是通过x得到的，q_fei , k_fei是通过x的非（纯知识信息）得到的

 def attention(self, q, k, v, q_fei, k_fei):
        '''
        Args:
        :param q: 综合信息的q,shape(batch_size,seq_len+pred_len,d_ff)
        :param k: 综合信息的k,shape(batch_size,seq_len+pred_len,d_ff)
        :param v: 综合信息的v,shape(batch_size,seq_len+pred_len,d_ff)
        :param q_fei: 知识信息的q,shape(batch_size,seq_len+pred_len,d_ff)
        :param v_fei: 知识信息的v,shape(batch_size,seq_len+pred_len,d_ff)
        :return: out,shape(batch_size,seq_len+pred_len,d_model)
        '''

首先q，k，q_fei , k_fei，先分别经过一个全连接然后在进行一下reshape变为（batch_size,h,seq_len+pred_len,dk)，因为经过一个全连接后在reshape，等同于经过多个全连接后在进行拼接

        # 多头注意力q,k,q_fei,k_fei的shape变为（batch_size,h,seq_len+pred_len,dk) ,dk*h=d_ff
        q=self.fc_q(q).reshape(q.shape[0],self.h,q.shape[1],-1)
        k = self.fc_k(k).reshape(k.shape[0],self.h,k.shape[1],-1)
        q_fei = self.fc_q_fei(q_fei).reshape(q_fei.shape[0],self.h ,q_fei.shape[1],-1)
        k_fei = self.fc_k_fei(k_fei).reshape(k_fei.shape[0], self.h, k_fei.shape[1],-1)

得到比例因子，为了后面计算Attention值的时候使用

    d= q.shape[-1]
    d_fei = q_fei.shape[-1] # 为了计算attention时候作为比例因子,就为dk

流程图中的
$$x_i^h\ast W_Q^x在代码中就是q，x_j^h\ast W_K^x在代码中就是K，剩下的也同理$$

      # K的最后两个维度进行一个转置，从(batch_size,h,seq_len+pred_len,dk)-->((batch_size,h,dk,seq_len+pred_len),为了能够进行矩阵的乘法
        # att and att_fei shape(batch_size,h,seq_len+pred_len,seq_len+pred_len)，表示的就是每一时间点之间的注意力关系
        att = torch.matmul(q, k.transpose(-1, -2)) / (math.sqrt(2 * d))
        att_fei = torch.matmul(q_fei, k_fei.transpose(-1, -2)) / (math.sqrt(2 * d_fei))
        att_final = att + att_fei

流程图中的
$$\sum _{j}Softmax(Att(i,j{)}^{\ast })在代码中就是score，（x_i^h\ast W_V^x）在代码中就是v，W在代码中为fc-out$$

        score = torch.softmax(att_final, dim=-1) # score的shape(batch_size,h,seq_len+pred_len,seq_len+pred_len)
        # 为了使得V可以和score进行矩阵乘法，因此把V reshape成(batch_size,h,seq_len+pred_len,dk)
        v=v.reshape(v.shape[0],self.h,v.shape[1],-1) # v从(batch_size,seq_len+pred_len,d_ff)-->(batch_size,h,seq_len+pred_len,dk)
        out = torch.matmul(score, v) # out shape(batch_size,seq_len+pred_len,dk)
        out=out.reshape(out.shape[0],out.shape[2],-1) # reshape为(batch_size,seq_len+pred_len,d_ff)
        out=self.fc_out(out) # 对特征维度进行降维，输入为(batch_size,seq_len+pred_len,d_ff)-->（batch_size,seq_len+pred_len,d_model）

Aliformer的forward函数

模型的主体部分在此

def forward(self, enc_x, enc_mark, y, y_mark,mode):
    '''
    Args:
    :param enc_x: (batch_size,seq_len,dim)
    :param enc_mark: (batch_size,seq_len,4)
    :param y: (batch_size,label_len+pred_len,dim)
    :param y_mark: (batch_size,label_len+pred_len,4)
    :param  mode:判断是否是在训练
    :return:
    '''

使用全零进行初始化综合信息 x和知识信息x_knowledge

x = torch.zeros(enc_x.shape[0], enc_x.shape[1] + self.pred_len, enc_x.shape[2])
x_knowledge = torch.zeros(enc_mark.shape[0], enc_mark.shape[1] + self.pred_len, enc_mark.shape[2]) #初始化

训练的时候，由于有两个训练模式，模式1，普通的训练模式，模式2:，span masking训练模型

因此先将seq_len和pred_len部分对应的数据都拼接起来

x是将seq_len,pred_len的data数据拼接起来

x_knowledge是将seq_len,pred_len部分的时间数据拼接起来

if mode == 'train':
    x[:,:self.seq_len,:]=x[:,:self.seq_len,:]+enc_x
    x[:,self.seq_len:,:]=x[:,self.seq_len:,:]+y[:,self.label_len:,:] # 将pred_len的数据也拼接上去

    x_knowledge[:, :self.seq_len, :] = x_knowledge[:, :self.seq_len, :] + enc_mark
    x_knowledge[:, self.seq_len:, :] = x_knowledge[:, self.seq_len:, :] + y_mark[:, self.label_len:, :] # 将pred_len的数据也拼接上去

选择训练模式

在之前就已经按照概率将列表中放置了0,1的列表中随机抽取，如果抽到的是0，那么代表不是span masking，即就把pred_len部分设置为0即可，并且对应的label也是pred_len部分的数据

如果抽到是1，那么就是span masking ，就随机在0~seq_len中抽取起始点，从抽取到起始点往后pred_len长度就是本次预测的label，把该部分设置为0

    choice=random.choice(self.choice)
    if choice==0:
        label=x[:,self.pred_len:,:] # 把label分出来
        x[:,self.pred_len:,:]=0

    else:
        star=random.choice(range(self.seq_len)) # 随机选取初始点
        label=x[:,star:star+self.pred_len,:]
        x[:,star:star+self.pred_len,:]=0

测试

测试的时候就和训练的普通模式差不多

测试的时候就只把seq_len部分的数据放上，pred_len部分的数据直接使用零进行代替

  else: # 如果是在测试，那么直接把pred_len部分的数据mask掉，设置为0
        x[:, :self.seq_len, :] = x[:, :self.seq_len, :] + enc_x
        label = y[:, self.label_len:, :]
        x_knowledge[:, :self.seq_len, :] = x_knowledge[:, :self.seq_len, :] + enc_mark
        x_knowledge[:, self.seq_len:, :] = x_knowledge[:, self.seq_len:, :] + y_mark[:, self.label_len:,:] # 将pred_len的数据也拼接上,因为知识数据是可以预先知道的

将数据进行embedding

x=self.embeded(x,x_knowledge)
x_knowledge=self.time_embeded(x_knowledge) # shape都为（batch_size,seq_len+pred_len,d_model）

堆叠AliAttention层

'''
Args:
:param x: shape是(batch_size,seq_len+pred_len,d_feature)
:param x_knowledge: shape是(batch_size,seq_len+pred_len,mark),知识信息，时间维度
'''
#  堆叠12层的AliAttention
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x,x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x,x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)
x = self.aliattention(x, x_knowledge)

将数据在特征维度和时间维度进行降维，然后和label进行loss计算，返回pred和loss值

# 输入为(batch_size,seq_len+pred_len,d_model)-->输出为(batch_size,seq_len+pred_len,d_feature)
x = self.out_dim(x)
# 输入为(batch_size,seq_len+pred_len,d_feature)-->输出为(batch_size,pred_len,d_feature)
y_hat=(self.out_time(x.permute(0,2,1))).permute(0,2,1)

loss=self.loss_fn(y_hat,label)
return y_hat,loss

完整的Aliformer代码

import torch
import torch.nn as nn
import math
import random
from layers.embeded import  DataEmbedding_time_token,TimeEmbedding


#----------------------------------------------------------------------------------------------
# 首先我们得到的输入数据为enc_x(batch_size,seq_len,dim)，y(batch_size,label_len+pred_len,dim)
# 其对应的时间数据为enc_mark(batch_size,seq_len,dim)，y_mark（batch_size,label_len+pred_len,dim)
# lable_len就是seq_len的后面长度为label_len部分的数据
# 因此我们通过拼接将数据enc_x和y变为x=（batch_size,seq_len+pred_len,dim)
# 并且并且对应的时间数据enc_mark和y_mark变为x_knowledge=（batch_size,seq_len+pred_len,dim)
# 然后我们在Aliformer中把x和x_knowledge经过embedding后变为（batch_size,seq_len+pred_len,d_model)
# 把x_knowledge经过dateembedding后变为（batch_size,seq_len+pred_len,d_model)
# 然后在送入AliAttention中，在AliAttention输出也为(batch_size,seq_len+pred_len,d_model)
#----------------------------------------------------------------------------------------------

class AliAttention(nn.Module):
    def __init__(self,h=8,d_feature=7,d_mark=4,d_model=64,d_ff=128,drop_out=0.1):
        '''
        Args:
        :param h:头数
        :param d_feature:输入数据的最后一个维度dim
        :param d_mark:知识信息的维度（本数据集为时间信息的维度，4）
        :param d_model:embedding后的维度
        :param d_ff:升为后的维度
        d_ff % h==0
        :param drop_out：drop_out的概率
        '''
        assert d_ff %h==0
        super().__init__()
        self.h=h
        self.d_feature=d_feature
        self.d_mark=d_mark
        self.d_model=d_model
        self.d_ff=d_ff
        self.drop_out=nn.Dropout(p=drop_out)
        self.dk=self.d_ff//self.h # 多头注意力机制中每一个头的维度

        # x -->q,k,v
        # 输入为以下三个输入输出的维度都是相同的（batch_size,seq_len+pred_len,d_model)-->输出为（batch_size,seq_len+pred_len,d_ff)
        self.fc_x_to_v = nn.Linear(self.d_model, self.d_ff)
        self.fc_x_to_k = nn.Linear(self.d_model, self.d_ff)
        self.fc_x_to_q=nn.Linear(self.d_model,self.d_ff)

        # 知识信息x的非-->k的非，Q的非
        # 以下两个的输入输出的维度都是相同的（batch_size,seq_len+pred_len,d_model)-->输出为（batch_size,seq_len+pred_len,d_ff)
        self.fc_knowledge_x_to_k=nn.Linear(self.d_model,self.d_ff)
        self.fc_knowledge_x_to_q = nn.Linear(self.d_model, self.d_ff)

        # 多头注意力机制,同维度变换，然后再reshape为，h,dk,(d_ff=dk*h)，就相当于经过h个全连接层在拼接起来
        # 以下四个的输入输出的shape都是相同的 batch_size,seq_len+pred_len,d_ff)-->输出为（batch_size,seq_len+pred_len,d_ff)
        self.fc_q=nn.Linear(self.d_ff,self.d_ff)
        self.fc_k=nn.Linear(self.d_ff,self.d_ff)
        self.fc_k_fei=nn.Linear(self.d_ff,self.d_ff)
        self.fc_q_fei=nn.Linear(self.d_ff,self.d_ff)

        # 输入为（batch_size,seq_len+pred_len,d_ff)-->输出为（batch_size,seq_len+pred_len,d_model)
        self.fc_out=nn.Linear(self.d_ff,self.d_model) # 将数据的维度变回输入进attention的维度

    def attention(self, q, k, v, q_fei, k_fei):
        '''
        Args:
        :param q: 综合信息的q,shape(batch_size,seq_len+pred_len,d_ff)
        :param k: 综合信息的k,shape(batch_size,seq_len+pred_len,d_ff)
        :param v: 综合信息的v,shape(batch_size,seq_len+pred_len,d_ff)
        :param q_fei: 知识信息的q,shape(batch_size,seq_len+pred_len,d_ff)
        :param v_fei: 知识信息的v,shape(batch_size,seq_len+pred_len,d_ff)
        :return: out,shape(batch_size,seq_len+pred_len,d_model)
        '''

        # 多头注意力q,k,q_fei,k_fei的shape变为（batch_size,h,seq_len+pred_len,dk) ,dk*h=d_ff
        q=self.fc_q(q).reshape(q.shape[0],self.h,q.shape[1],-1)
        k = self.fc_k(k).reshape(k.shape[0],self.h,k.shape[1],-1)
        q_fei = self.fc_q_fei(q_fei).reshape(q_fei.shape[0],self.h ,q_fei.shape[1],-1)
        k_fei = self.fc_k_fei(k_fei).reshape(k_fei.shape[0], self.h, k_fei.shape[1],-1)

        d = q.shape[-1]
        d_fei = q_fei.shape[-1] # 为了计算attention时候作为比例因子,就为dk

        # K的最后两个维度进行一个转置，从(batch_size,h,seq_len+pred_len,dk)-->((batch_size,h,dk,seq_len+pred_len),为了能够进行矩阵的乘法
        # att and att_fei shape(batch_size,h,seq_len+pred_len,seq_len+pred_len)，表示的就是每一时间点之间的注意力关系
        att = torch.matmul(q, k.transpose(-1, -2)) / (math.sqrt(2 * d))
        att_fei = torch.matmul(q_fei, k_fei.transpose(-1, -2)) / (math.sqrt(2 * d_fei))
        att_final = att + att_fei
        score = torch.softmax(att_final, dim=-1) # score的shape(batch_size,h,seq_len+pred_len,seq_len+pred_len)
        score=self.drop_out(score)
        # 为了使得V可以和score进行矩阵乘法，因此把V reshape成(batch_size,h,seq_len+pred_len,dk)
        v=v.reshape(v.shape[0],self.h,v.shape[1],-1) # v从(batch_size,seq_len+pred_len,d_ff)变为(batch_size,h,seq_len+pred_len,dk)
        out = torch.matmul(score, v) # out shape(batch_size,seq_len+pred_len,dk)
        out=out.reshape(out.shape[0],out.shape[2],-1) # reshape为(batch_size,seq_len+pred_len,d_ff)
        out=self.fc_out(out) # 对特征维度进行降维，输入为(batch_size,seq_len+pred_len,d_ff)-->（batch_size,seq_len+pred_len,d_model）
        return out

    def forward(self,x,x_knowledge):
        '''
        Args:
        :param x: shape是(batch_size,seq_len+pred_len,d_model)
        :param x_knowledge: shape是(batch_size,seq_len+pred_len,mark),知识信息，时间维度
        :return:out: shape(batch_size,seq_len+pred_len,d_model)
        '''
        # 得到q，k，v
        # v,k,q都是(batch_size, seq_len + pred_len, d_model)-->（batch_size, seq_len + pred_len, d_ff)
        v=self.fc_x_to_v(x)
        k=self.fc_x_to_k(x)
        q=self.fc_x_to_q(x)
        # 得到知识信息的q，k，称为q_fei,k_fei
        # k_fei,q_fei（batch_size, seq_len + pred_len, d_model)-->输出为（batch_size, seq_len + pred_len, d_ff)
        k_fei=self.fc_knowledge_x_to_k(x_knowledge)
        q_fei=self.fc_knowledge_x_to_q(x_knowledge)

        # 输入进attention中得到输出out为（batch_size,seq_len+pred_len,d_model)
        out=self.attention(q,k,v,q_fei,k_fei)
        return out

    
    
class Aliformer(nn.Module):
    def __init__(self,args,h=8,p2=0.5,loss_fn=torch.nn.MSELoss()):
        '''
        Args:
        :param h:头数 多头注意力网络中的头数
        :param d_feature:输入数据的最后一个维度dim（本数据集为 7）
        :param d_mark:知识信息的维度（本数据集为时间信息的维度 4）
        :param d_model:embedding后的维度
        :param d_ff:升为后的维度
        d_ff % h==0
        :param p2: 在训练的时候使用跨度掩码的概率为p2
        :param loss_fn 使用的loss函数
        '''
        super().__init__()
        self.h=h
        self.d_feature=args.d_feature
        self.d_mark=args.d_mark
        self.d_model=args.d_model
        self.d_ff=args.d_ff
        self.p2=p2

        # embedding层
        self.time_embeded = TimeEmbedding(d_mark=self.d_mark, d_model=self.d_model)  # 对时间维度进行embedding
        self.embeded = DataEmbedding(d_feature=self.d_feature, d_mark=self.d_mark,
                                     d_model=self.d_model, dropout=0, pos=False)  # 把原始数据和时间维度embedding并且加在一起

        self.aliattention=AliAttention(h=self.h,d_feature=self.d_feature,
                                       d_mark=self.d_mark,d_model=self.d_model,d_ff=self.d_ff) # 初始化AliAttention层

        self.pred_len=args.pred_len
        self.label_len=args.label_len
        self.seq_len=args.seq_len

        self.out_dim=nn.Linear(self.d_model,self.d_feature) # 对特征维度进行降纬
        self.out_time=nn.Linear((self.seq_len+self.pred_len),self.pred_len) # 对时间维度进行降维

        self.choice=torch.zeros((100))
        # 为了后面的选择是否使用span masking 策略,如果选择为1，那么就使用span mask策略
        self.choice[:int(p2*100)]=1 # 把p2*100的数字变成1，剩下的是0，然后在这一个列表中随机抽样
        self.loss_fn=loss_fn

    def forward(self, enc_x, enc_mark, y, y_mark,mode):
        '''
        Args:
        :param enc_x: (batch_size,seq_len,dim)
        :param enc_mark: (batch_size,seq_len,d_mark)
        :param y: (batch_size,label_len+pred_len,dim)
        :param y_mark: (batch_size,label_len+pred_len,d_mark)
        :param  mode:判断是否是在训练
        :return:
        '''
        # 初始化x装的是data，x_knowledge装的是时间信息
        x = torch.zeros(enc_x.shape[0], enc_x.shape[1] + self.pred_len, enc_x.shape[2],device=enc_x.device) # x（batch_size,seq_len+pred_len,dim)
        x_knowledge = torch.zeros(enc_mark.shape[0], enc_mark.shape[1] + self.pred_len, enc_mark.shape[2],device=enc_x.device) # x_knowledge（batch_size,seq_len+pred_len,d_mark)

        if mode == 'train':
            # 将seq_len+pred_len的时间数据都拼到x上
            x[:,:self.seq_len,:]=x[:,:self.seq_len,:]+enc_x
            x[:,self.seq_len:,:]=x[:,self.seq_len:,:]+y[:,self.label_len:,:] # 将pred_len的数据拼接上去

            # 将seq_len+pred_len的时间数据都拼接到x_knowledge上
            x_knowledge[:, :self.seq_len, :] = x_knowledge[:, :self.seq_len, :] + enc_mark
            x_knowledge[:, self.seq_len:, :] = x_knowledge[:, self.seq_len:, :] + y_mark[:, self.label_len:, :] # 将pred_len的数据拼接上去

            choice=random.choice(self.choice) #self.choice中有100*int(p2)个1，剩下的都是0,在其中随机抽取来代表是使用span masking还是正常的预测
            if choice==0: #正常的预测
                label=x[:,self.pred_len:,:].clone() # 把label分出来
                x[:,self.pred_len:,:]=0 # 把label对应的部分设置为0
                star=self.seq_len # 需要预测的部分对应的起始index

            else: # span masking
                star=random.choice(range(self.seq_len)) # 在0~seq_len中随机选取初始点
                label=x[:,star:star+self.pred_len,:].clone() # label就是strar-star+pred_len部分
                x[:,star:star+self.pred_len,:]=0 # 把label对应的部分设置为0

        else: # 如果是在测试，那么直接把pred_len部分的数据mask掉，设置为0
        # 测试下的x shape也是(batch_size,seq_len+pred_len,dim),x_knowledge (batch_size,seq_len+pred_len,d_mark),
        # 只不过pred_len部分全部为0
            x[:, :self.seq_len, :] = x[:, :self.seq_len, :] + enc_x # 将seq_len部分的data赋值给x
            label = y[:, self.label_len:, :] # 得到label 是pred_len 部分对应的data
            x_knowledge[:, :self.seq_len, :] = x_knowledge[:, :self.seq_len, :] + enc_mark # 将seq_len部分的时间数据赋值给x_knowledge
            x_knowledge[:, self.seq_len:, :] = x_knowledge[:, self.seq_len:, :] + y_mark[:, self.label_len:,:] # 将pred_len的数据也赋值给x_knowledge,因为知识数据是可以预先知道的
            star=self.seq_len # 需要预测的部分对应的index
        '''
        Args:
        :param x: shape是(batch_size,seq_len+pred_len,d_feature)
        :param x_knowledge: shape是(batch_size,seq_len+pred_len,d_mark),知识信息，时间维度
        '''

        x=self.embeded(x,x_knowledge) # 输出x的shape为 (batch_size,seq_len+pred_len，d_model）
        x_knowledge=self.time_embeded(x_knowledge) # x_knowledge shape为（batch_size,seq_len+pred_len,d_model）

        x=self.aliattention(x,x_knowledge)
        x_new=x.clone()

        pred=self.out_time(x.permute(0,2,1)).permute(0,2,1)
        x_new[:,star:star+self.pred_len,:]=pred # 把使用AliAttention预测的值拼接回去

        x=self.embeded(x,x_knowledge)
        x_knowledge=self.time_embeded(x_knowledge) # shape都为（batch_size,seq_len+pred_len,d_model）

        #  堆叠12层的AliAttention
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        x = self.aliattention(x, x_knowledge)
        # 输入为(batch_size,seq_len+pred_len,d_model)-->输出为(batch_size,seq_len+pred_len,d_feature)
        x = self.out_dim(x)
        # 输入为(batch_size,seq_len+pred_len,d_feature)-->输出为(batch_size,pred_len,d_feature)
        y_hat=(self.out_time(x.permute(0,2,1))).permute(0,2,1)
        loss=self.loss_fn(y_hat,label) # 计算损失loss
        return y_hat,loss

embedding层代码

class TokenEmbedding_Aliformer(nn.Module):
    def __init__(self, d_feature, d_model):
        super(TokenEmbedding_Aliformer, self).__init__()
        self.embed = nn.Linear(d_feature, d_model, bias=False)

    def forward(self, x):
        return self.embed(x)


class DataEmbedding_time_token(nn.Module):
    def __init__(self, d_feature, d_mark, d_model):
        super(DataEmbedding_time_token, self).__init__()

        self.value_embedding = TokenEmbedding_Aliformer(d_feature=d_feature, d_model=d_model)
        self.time_embedding = TimeEmbedding(d_mark=d_mark, d_model=d_model)

    def forward(self, x, x_mark):

        x = self.value_embedding(x) + self.time_embedding(x_mark)
        return x
    

class TimeEmbedding(nn.Module):
    def __init__(self, d_mark, d_model):
        super(TimeEmbedding, self).__init__()
        self.embed = nn.Linear(d_mark, d_model, bias=False)

    def forward(self, x):
        return self.embed(x)