informer代码

Informer是一种基于Transformer的序列模型，主要用于时间序列预测任务。其最大的特点是引入了全局和局部注意力机制，使得模型更加准确和稳定。以下是Informer的PyTorch实现代码示例：

import torch
import torch.nn as nn
import numpy as np

class InformerEncoderLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, dim_feedforward, dropout_rate=0.0):
        super(InformerEncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout_rate)
        self.linear1 = nn.Linear(embed_dim, dim_feedforward)
        self.dropout = nn.Dropout(dropout_rate)
        self.linear2 = nn.Linear(dim_feedforward, embed_dim)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout_rate)
        self.dropout2 = nn.Dropout(dropout_rate)

    def forward(self, x):
        # self-attention
        res, _ = self.self_attn(x, x, x)
        x = x + self.dropout1(res)
        x = self.norm1(x)

        # feedforward
        res = self.linear2(self.dropout(torch.relu(self.linear1(x))))
        x = x + self.dropout2(res)
        x = self.norm2(x)

        return x

class InformerEncoder(nn.Module):
    def __init__(self, input_size, input_dim, embed_dim, num_heads, num_layers):
        super(InformerEncoder, self).__init__()
        self.input_fc = nn.Linear(input_size * input_dim, embed_dim)
        self.pos_encoding = nn.Parameter(torch.zeros(1, input_size, embed_dim))
        self.layers = nn.ModuleList([InformerEncoderLayer(embed_dim, num_heads, dim_feedforward=2048) for _ in range(num_layers)])

    def forward(self, x):
        # flatten input
        x = x.reshape(x.shape[0], -1)

        # input projection
        x = self.input_fc(x)

        # add position encoding
        x = x.unsqueeze(1) + self.pos_encoding

        # pass through encoder layers
        for layer in self.layers:
            x = layer(x)

        return x

class InformerDecoderLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, dim_feedforward, dropout_rate=0.0):
        super(InformerDecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout_rate)
        self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout_rate)
        self.linear1 = nn.Linear(embed_dim, dim_feedforward)
        self.dropout = nn.Dropout(dropout_rate)
        self.linear2 = nn.Linear(dim_feedforward, embed_dim)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.norm3 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout_rate)
        self.dropout2 = nn.Dropout(dropout_rate)
        self.dropout3 = nn.Dropout(dropout_rate)

    def forward(self, x, encoder_out):
        # self-attention
        res, _ = self.self_attn(x, x, x)
        x = x + self.dropout1(res)
        x = self.norm1(x)

        # encoder-decoder attention
        res, _ = self.multihead_attn(x, encoder_out, encoder_out)
        x = x + self.dropout2(res)
        x = self.norm2(x)

        # feedforward
        res = self.linear2(self.dropout(torch.relu(self.linear1(x))))
        x = x + self.dropout3(res)
        x = self.norm3(x)

        return x

class InformerDecoder(nn.Module):
    def __init__(self, output_size, output_dim, embed_dim, num_heads, num_layers):
        super(InformerDecoder, self).__init__()
        self.output_fc = nn.Linear(output_dim, embed_dim)
        self.pos_encoding = nn.Parameter(torch.zeros(1, output_size, embed_dim))
        self.layers = nn.ModuleList([InformerDecoderLayer(embed_dim, num_heads, dim_feedforward=2048) for _ in range(num_layers)])
        self.output_proj = nn.Linear(embed_dim, output_dim)

    def forward(self, x, encoder_out):
        # output projection
        x = self.output_fc(x)

        # add position encoding
        x = x.unsqueeze(1) + self.pos_encoding

        # pass through decoder layers
        for layer in self.layers:
            x = layer(x, encoder_out)

        # output projection
        x = self.output_proj(x)

        return x

class Informer(nn.Module):
    def __init__(self, input_size, input_dim, output_size, output_dim, embed_dim=64, num_heads=4, enc_layers=2, dec_layers=1):
        super(Informer, self).__init__()
        self.encoder = InformerEncoder(input_size, input_dim, embed_dim, num_heads, enc_layers)
        self.decoder = InformerDecoder(output_size, output_dim, embed_dim, num_heads, dec_layers)

    def forward(self, x):
        encoder_out = self.encoder(x[:, :-1])
        decoder_out = self.decoder(x[:, -1], encoder_out)
        return decoder_out

# example usage
model = Informer(input_size=24, input_dim=1, output_size=24, output_dim=1)
x = torch.randn(16, 25, 1)
y = model(x)
print(y.shape)