TrafficWheel/model/Informer/model.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from model.Informer.encoder import Encoder, EncoderLayer, ConvLayer, EncoderStack
from model.Informer.decoder import Decoder, DecoderLayer
from model.Informer.attn import FullAttention, ProbAttention, AttentionLayer
from model.Informer.embed import DataEmbedding
from model.Informer.masking import TriangularCausalMask, ProbMask

class Informer(nn.Module):
    def __init__(self, configs):
        super(Informer, self).__init__()
        # 从configs中提取参数
        self.enc_in = configs.get("enc_in", 1)
        self.dec_in = configs.get("dec_in", 1)
        self.c_out = configs.get("c_out", 1)
        self.seq_len = configs.get("seq_len", 96)
        self.label_len = configs.get("label_len", 48)
        self.out_len = configs.get("out_len", 24)
        self.factor = configs.get("factor", 5)
        self.d_model = configs.get("d_model", 512)
        self.n_heads = configs.get("n_heads", 8)
        self.e_layers = configs.get("e_layers", 3)
        self.d_layers = configs.get("d_layers", 2)
        self.d_ff = configs.get("d_ff", 512)
        self.dropout = configs.get("dropout", 0.0)
        self.attn = configs.get("attn", "prob")
        self.embed = configs.get("embed", "fixed")
        self.freq = configs.get("freq", "h")
        self.activation = configs.get("activation", "gelu")
        self.output_attention = configs.get("output_attention", False)
        self.distil = configs.get("distil", True)
        self.mix = configs.get("mix", True)
        self.device = configs.get("device", torch.device('cuda:0'))
        
        self.pred_len = self.out_len
        
        # 编码层
        self.enc_embedding = DataEmbedding(self.enc_in, self.d_model, self.embed, self.freq, self.dropout)
        self.dec_embedding = DataEmbedding(self.dec_in, self.d_model, self.embed, self.freq, self.dropout)
        
        # 注意力层
        Attn = ProbAttention if self.attn == 'prob' else FullAttention
        
        # 编码器
        self.encoder = Encoder(
            [
                EncoderLayer(
                    AttentionLayer(Attn(False, self.factor, attention_dropout=self.dropout, output_attention=self.output_attention), 
                                self.d_model, self.n_heads, mix=False),
                    self.d_model,
                    self.d_ff,
                    dropout=self.dropout,
                    activation=self.activation
                ) for l in range(self.e_layers)
            ],
            [
                ConvLayer(
                    self.d_model
                ) for l in range(self.e_layers - 1)
            ] if self.distil else None,
            norm_layer=torch.nn.LayerNorm(self.d_model)
        )
        
        # 解码器
        self.decoder = Decoder(
            [
                DecoderLayer(
                    AttentionLayer(Attn(True, self.factor, attention_dropout=self.dropout, output_attention=False), 
                                self.d_model, self.n_heads, mix=self.mix),
                    AttentionLayer(FullAttention(False, self.factor, attention_dropout=self.dropout, output_attention=False), 
                                self.d_model, self.n_heads, mix=False),
                    self.d_model,
                    self.d_ff,
                    dropout=self.dropout,
                    activation=self.activation,
                )
                for l in range(self.d_layers)
            ],
            norm_layer=torch.nn.LayerNorm(self.d_model)
        )
        
        # 投影层
        self.projection = nn.Linear(self.d_model, self.c_out, bias=True)
        
    def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None,
                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
        # 如果没有提供x_dec和x_mark_dec，则根据x_enc和label_len生成
        if x_dec is None:
            x_dec = torch.cat([x_enc[:, -self.label_len:, :], torch.zeros_like(x_enc[:, :self.pred_len, :])], dim=1)
        if x_mark_dec is None and x_mark_enc is not None:
            x_mark_dec = torch.cat([x_mark_enc[:, -self.label_len:, :], torch.zeros_like(x_mark_enc[:, :self.pred_len, :])], dim=1)
        
        # 编码
        enc_out = self.enc_embedding(x_enc, x_mark_enc)
        enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
        
        # 解码
        dec_out = self.dec_embedding(x_dec, x_mark_dec)
        dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
        dec_out = self.projection(dec_out)
        
        if self.output_attention:
            return dec_out[:, -self.pred_len:, :], attns
        else:
            return dec_out[:, -self.pred_len:, :] # [B, L, D]


class InformerStack(nn.Module):
    def __init__(self, configs):
        super(InformerStack, self).__init__()
        # 从configs中提取参数
        self.enc_in = configs.get("enc_in", 1)
        self.dec_in = configs.get("dec_in", 1)
        self.c_out = configs.get("c_out", 1)
        self.seq_len = configs.get("seq_len", 96)
        self.label_len = configs.get("label_len", 48)
        self.out_len = configs.get("out_len", 24)
        self.factor = configs.get("factor", 5)
        self.d_model = configs.get("d_model", 512)
        self.n_heads = configs.get("n_heads", 8)
        self.e_layers = configs.get("e_layers", [3, 2, 1])
        self.d_layers = configs.get("d_layers", 2)
        self.d_ff = configs.get("d_ff", 512)
        self.dropout = configs.get("dropout", 0.0)
        self.attn = configs.get("attn", "prob")
        self.embed = configs.get("embed", "fixed")
        self.freq = configs.get("freq", "h")
        self.activation = configs.get("activation", "gelu")
        self.output_attention = configs.get("output_attention", False)
        self.distil = configs.get("distil", True)
        self.mix = configs.get("mix", True)
        self.device = configs.get("device", torch.device('cuda:0'))
        
        self.pred_len = self.out_len
        
        # 编码层
        self.enc_embedding = DataEmbedding(self.enc_in, self.d_model, self.embed, self.freq, self.dropout)
        self.dec_embedding = DataEmbedding(self.dec_in, self.d_model, self.embed, self.freq, self.dropout)
        
        # 注意力层
        Attn = ProbAttention if self.attn == 'prob' else FullAttention
        
        # 编码器栈
        inp_lens = list(range(len(self.e_layers))) # [0,1,2,...] you can customize here
        encoders = [
            Encoder(
                [
                    EncoderLayer(
                        AttentionLayer(Attn(False, self.factor, attention_dropout=self.dropout, output_attention=self.output_attention), 
                                    self.d_model, self.n_heads, mix=False),
                        self.d_model,
                        self.d_ff,
                        dropout=self.dropout,
                        activation=self.activation
                    ) for l in range(el)
                ],
                [
                    ConvLayer(
                        self.d_model
                    ) for l in range(el-1)
                ] if self.distil else None,
                norm_layer=torch.nn.LayerNorm(self.d_model)
            ) for el in self.e_layers]
        self.encoder = EncoderStack(encoders, inp_lens)
        
        # 解码器
        self.decoder = Decoder(
            [
                DecoderLayer(
                    AttentionLayer(Attn(True, self.factor, attention_dropout=self.dropout, output_attention=False), 
                                self.d_model, self.n_heads, mix=self.mix),
                    AttentionLayer(FullAttention(False, self.factor, attention_dropout=self.dropout, output_attention=False), 
                                self.d_model, self.n_heads, mix=False),
                    self.d_model,
                    self.d_ff,
                    dropout=self.dropout,
                    activation=self.activation,
                )
                for l in range(self.d_layers)
            ],
            norm_layer=torch.nn.LayerNorm(self.d_model)
        )
        
        # 投影层
        self.projection = nn.Linear(self.d_model, self.c_out, bias=True)
        
    def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None,
                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
        # 如果没有提供x_dec和x_mark_dec，则根据x_enc和label_len生成
        if x_dec is None:
            x_dec = torch.cat([x_enc[:, -self.label_len:, :], torch.zeros_like(x_enc[:, :self.pred_len, :])], dim=1)
        if x_mark_dec is None and x_mark_enc is not None:
            x_mark_dec = torch.cat([x_mark_enc[:, -self.label_len:, :], torch.zeros_like(x_mark_enc[:, :self.pred_len, :])], dim=1)
        
        # 编码
        enc_out = self.enc_embedding(x_enc, x_mark_enc)
        enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
        
        # 解码
        dec_out = self.dec_embedding(x_dec, x_mark_dec)
        dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
        dec_out = self.projection(dec_out)
        
        if self.output_attention:
            return dec_out[:, -self.pred_len:, :], attns
        else:
            return dec_out[:, -self.pred_len:, :] # [B, L, D]