models/ITFormer.py

"""
ITFormer (Instruction-aware Time Series Transformer)
"""
import math
import torch
import torch.nn.functional as F
from torch import nn
from timm.layers import Mlp, DropPath
from timm.layers.helpers import to_2tuple
from transformers.modeling_outputs import CausalLMOutputWithPast
from utils.position_coding import LearnablePositionalEmbedding, SinusoidalPositionalEncoding,RotaryPositionalEncoding
from models.layers.attention import InstructTimeAttention
from utils.log_util import adaptive_print

class SeqCrossAttention(nn.Module):
    def __init__(
            self,
            dim,
            num_heads=8,
            qkv_bias=False,
            qk_norm=False,
            attn_drop=0.,
            proj_drop=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv_proj = nn.Linear(dim, dim * 2, bias=qkv_bias)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, query, key_value, attn_mask=None):
        B, N, C = query.shape
        _, V, L, _ = key_value.shape

        # Reshape Key and Value to focus only on L (time) dimension
        key_value = key_value.view(B * V, L, C)

        # Compute Query, Key, and Value projections
        q = self.q_proj(query).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        kv = self.kv_proj(key_value).reshape(B * V, L, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)

        # Adjust batch size for Key and Value
        k = k.view(B, V, self.num_heads, L, self.head_dim).permute(0, 2, 1, 3, 4).reshape(B * self.num_heads, V * L, self.head_dim)
        v = v.view(B, V, self.num_heads, L, self.head_dim).permute(0, 2, 1, 3, 4).reshape(B * self.num_heads, V * L, self.head_dim)

        # Apply normalization (if any)
        q = self.q_norm(q).reshape(B * self.num_heads, N, self.head_dim)
        k = self.k_norm(k)

        # Scaled Dot-Product Attention over L dimension
        x = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.attn_drop.p if self.training else 0.
        )

        # Reshape and project output
        x = x.view(B, self.num_heads, N, self.head_dim).permute(0, 2, 1, 3).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
    
class SeqAttBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            qkv_bias=False,
            qk_norm=False,
            proj_drop=0.,
            attn_drop=0.,
            init_values=None,
            drop_path=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn_seq = SeqCrossAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,
        )

        self.drop_path1 = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()
        self.proj = nn.Linear(dim, dim)

    def forward(self, x, key_value, attn_mask):
        x_input = x
        x = self.norm1(x)
        key_value = self.norm1(key_value)

        # key_value = torch.reshape(
        #     key_value, (-1, key_value.shape[-2], key_value.shape[-1]))
        x = self.attn_seq(x, key_value, attn_mask)
        x = x_input + self.drop_path1(x)
        return x

class VarCrossAttention(nn.Module):

    def __init__(
            self,
            dim,
            num_heads=8,
            qkv_bias=False,
            qk_norm=False,
            attn_drop=0.,
            proj_drop=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv_proj = nn.Linear(dim, dim * 2, bias=qkv_bias)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, query, key_value, attn_mask=None):
        B, N, C = query.shape
        _, V, L, _ = key_value.shape

        # Reshape Key and Value to focus only on V (variable) dimension
        key_value = key_value.view(B * L, V, C)

        # Compute Query, Key, and Value projections
        q = self.q_proj(query).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        kv = self.kv_proj(key_value).reshape(B * L, V, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)

        # Adjust batch size for Key and Value
        k = k.view(B, L, self.num_heads, V, self.head_dim).permute(0, 2, 1, 3, 4).reshape(B * self.num_heads, L * V, self.head_dim)
        v = v.view(B, L, self.num_heads, V, self.head_dim).permute(0, 2, 1, 3, 4).reshape(B * self.num_heads, L * V, self.head_dim)

        # Apply normalization (if any)
        q = self.q_norm(q).reshape(B * self.num_heads, N, self.head_dim)
        k = self.k_norm(k)

        # Scaled Dot-Product Attention over V dimension
        x = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.attn_drop.p if self.training else 0.
        )

        # Reshape and project output
        x = x.view(B, self.num_heads, N, self.head_dim).permute(0, 2, 1, 3).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class VarAttBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            qkv_bias=False,
            qk_norm=False,
            proj_drop=0.,
            attn_drop=0.,
            init_values=None,
            drop_path=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn_var = VarCrossAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,
        )

        self.drop_path1 = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()
        self.proj = nn.Linear(dim, dim)

    def forward(self, x, key_value, attn_mask):
        x_input = x
        x = self.norm1(x)
        key_value = self.norm1(key_value)

        x = self.attn_var(x, key_value, attn_mask)
        x = x_input + self.drop_path1(x)
        return x

class SeqAttention(nn.Module):

    def __init__(
            self,
            dim,
            num_heads=8,
            qkv_bias=False,
            qk_norm=False,
            attn_drop=0.,
            proj_drop=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, attn_mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads,
                                  self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        q, k = self.q_norm(q), self.k_norm(k)
        x = F.scaled_dot_product_attention(
            q, k, v,  attn_mask=attn_mask,
            dropout_p=self.attn_drop.p if self.training else 0.,
        )

        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class SelfAttBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            qkv_bias=False,
            qk_norm=False,
            proj_drop=0.,
            attn_drop=0.,
            init_values=None,
            drop_path=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn_seq = SeqAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,
        )

        self.drop_path1 = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x, attn_mask=None):
        x_input = x
        x = self.norm1(x)
        x = self.attn_seq(x, attn_mask)
        x = x_input + self.drop_path1(x)
        return x


class ITAttBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            qkv_bias=False,
            qk_norm=False,
            proj_drop=0.,
            attn_drop=0.,
            init_values=None,
            drop_path=0.,
            norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn_it = InstructTimeAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,

        )
        self.drop_path1 = DropPath(
            drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x,memory, attn_mask=None):
        x_input = x
        x = self.attn_it(x, memory,attn_mask)
        x = x_input + self.norm1(self.drop_path1(x))
        return x

class DecoderBasicBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            mlp_ratio=4.0,
            qkv_bias=False,
            qk_norm=False,
            proj_drop=0.,
            attn_drop=0.,
            drop_path=0.,
            act_layer=nn.GELU,
            norm_layer=nn.LayerNorm,
            prefix_num=10,
            legacy_double_residual=True,
    ):
        super().__init__()

        self.prefix_num = prefix_num
        self.legacy_double_residual = legacy_double_residual

        self.self_attn = SelfAttBlock(
              dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
            attn_drop=attn_drop, proj_drop=proj_drop, drop_path=drop_path, norm_layer=norm_layer
        ) 

        self.it_attn = ITAttBlock(
            dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
            attn_drop=attn_drop, proj_drop=proj_drop, drop_path=drop_path, norm_layer=norm_layer
        )

        
        self.feed_forward_prefix = nn.Sequential(
            norm_layer(dim),
            nn.Linear(dim, int(dim * mlp_ratio)),
            act_layer(),
            nn.Dropout(proj_drop),
            nn.Linear(int(dim * mlp_ratio), dim),
        )

        self.feed_forward_instruct = nn.Sequential(
            norm_layer(dim),
            nn.Linear(dim, int(dim * mlp_ratio)),
            act_layer(),
            nn.Dropout(proj_drop),
            nn.Linear(int(dim * mlp_ratio), dim),
        )


        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x, memory, attn_mask=None):

        # Self-attention block
        self_attended = self.self_attn(x, attn_mask)
        x = x + self_attended if self.legacy_double_residual else self_attended
        prefix  =x[:, :self.prefix_num, :]
        x = self.feed_forward_instruct(x)+x
        # Cross-attention block 10token vs b,n,c,d中的n
        time_attended = self.it_attn(prefix, memory, attn_mask)
        prefix = (
            prefix + time_attended
            if self.legacy_double_residual
            else time_attended
        )
        # 10token vs b,n,c,d中的c
        # Feed forward block
        prefix = prefix + self.feed_forward_prefix(prefix)
        # Concatenate prefix and x
        x = torch.cat([prefix, x[:, self.prefix_num:, :]], dim=1)
        return x


class ITFormer(nn.Module):
    def __init__(self, args):
        super(ITFormer, self).__init__()
        self.layers = nn.ModuleList([
            DecoderBasicBlock(
                dim=args.it_d_model,
                num_heads=args.it_n_heads,
                mlp_ratio=4.,
                qkv_bias=True,
                qk_norm=False,
                proj_drop=args.it_dropout,
                attn_drop=args.it_dropout,
                drop_path=0.,
                act_layer=nn.GELU,
                norm_layer=nn.LayerNorm,
                prefix_num=args.prefix_num,
                legacy_double_residual=getattr(
                    args, "itformer_legacy_double_residual", True
                ),
            ) for _ in range(args.it_layers)
        ])
        self.norm = nn.LayerNorm(args.it_d_model)


        #time posi
        self.time_pos = SinusoidalPositionalEncoding(args.it_d_model)
        #variable posi
        self.var_pos = LearnablePositionalEmbedding(args.it_d_model)
        #instruction posi
        self.instruc_pos = SinusoidalPositionalEncoding(args.it_d_model)
        # cycle posi
        self.cycle_pos = RotaryPositionalEncoding(args.it_d_model)

        #prefix num
        self.prefix_num = args.prefix_num
        self.prefix_token = nn.Parameter(
            torch.empty(1, args.prefix_num, args.it_d_model)
        )
        nn.init.normal_(self.prefix_token, std=0.02)
    def forward(self, x, memory, stage=None,attn_mask=None):
        # Add prefix token to x 
        x = torch.cat([self.prefix_token.repeat(x.shape[0], 1, 1), x], dim=1)
        # Positional encoding
        # Apply positional encoding to x
        x = x + self.instruc_pos(x)

        # Stage 处理逻辑改进
        if torch.is_tensor(stage):
            stage_list = stage.tolist()
        else:
            stage_list = stage

        # 找出不同 stage 的索引
        cycle_index = [i for i, s in enumerate(stage_list) if s != 3 and s != 4]
        cross_cycle_index = [i for i, s in enumerate(stage_list) if s == 3 or s == 4]
        
        # 记录原始顺序以便恢复，确保与 x 对齐
        original_indices = cycle_index + cross_cycle_index
        reorder_map = {idx: i for i, idx in enumerate(original_indices)}
        reverse_indices = [reorder_map[i] for i in range(len(stage_list))]

        processed_memories = []

        if len(cycle_index) > 0:
            sub_memory = memory[cycle_index]
            b, l, v, d = sub_memory.shape
            sub_memory = sub_memory.view(b * l, v, d)
            sub_memory = sub_memory + self.time_pos(sub_memory)
            sub_memory = sub_memory.view(b, l, v, d)
            processed_memories.append((cycle_index, sub_memory))

        if len(cross_cycle_index) > 0:
            sub_memory = memory[cross_cycle_index]
            b, l, v, d = sub_memory.shape
            sub_memory = sub_memory.view(b * v, l, d)
            sub_memory = sub_memory + self.cycle_pos(sub_memory)
            sub_memory = sub_memory.view(b, l, v, d)
            processed_memories.append((cross_cycle_index, sub_memory))

        # 按照拼接后的顺序排列
        all_processed = torch.cat([m for _, m in processed_memories], dim=0)
        # 关键步骤：恢复原始 batch 顺序以匹配 x
        memory = all_processed[reverse_indices]

        # 再次处理变量维度
        b, l, v, d = memory.shape
        memory = memory.view(b * l, v, d)
        memory = memory + self.var_pos(memory)
        memory = memory.view(b, l, v, d)

        for i, layer in enumerate(self.layers):
            x = layer(x, memory, attn_mask)

        x = self.norm(x)
        return x[:, :self.prefix_num, :]

def count_parameters(model):
    """统计模型中可训练参数的总数"""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)



if __name__ == "__main__":

    # dim = 64
    # num_heads = 8
    # seq_len = 20
    # var_num = 5
    # memory_len = 30
    # batch_size = 2

    # x = torch.randn(batch_size, seq_len, dim)
    # memory = torch.randn(batch_size, var_num,memory_len, dim)
    # attn_mask = None

    # decoder_block = DecoderBasicBlock(
    #     dim=dim, num_heads=num_heads, qkv_bias=True, proj_drop=0.1, attn_drop=0.1
    # )
    # output = decoder_block(x, memory, attn_mask)
    # print("DecoderBasicBlock Output Shape:", output.shape)
    # class Args:
    #     def __init__(self):
    #         self.it_d_model = 64
    #         self.it_n_heads = 8
    #         self.it_layers = 6
    #         self.it_dropout = 0.1
    #         self.prefix_num = 10

    # args = Args()
    # model = ITformer(args)

    # x = torch.randn(batch_size, seq_len, dim)
    # memory = torch.randn(batch_size, var_num, memory_len, dim)
    # attn_mask = None
    # stage = [1,2]
    # output = model(x, memory,stage, attn_mask)
    # print("Model Output Shape:", output.shape)
    class Args:
        def __init__(self):
            self.it_d_model = 512
            self.it_n_heads = 8
            self.it_layers = 4
            self.it_dropout = 0.1
            self.prefix_num = 10

    args = Args()
    model = ITFormer(args)

    # 打印可训练参数量
    total_trainable_params = count_parameters(model)
    print(f"Total Trainable Parameters: {total_trainable_params:,}")

    # # 可选：打印每一层的参数量
    # print("\nLayer-wise Parameters:")
    # for name, param in model.named_parameters():
    #     if param.requires_grad:
    #         print(f"{name}: {param.numel():,}")