""" Time Series Encoder """ 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 def calculate_unfold_output_length(input_length, size, step): # Calculate the number of windows num_windows = (input_length - size) // step + 1 return num_windows class CrossAttention(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, var_num=None, ): 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 = nn.Linear(dim, dim, bias=qkv_bias) self.kv = 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) if var_num is not None: self.template = nn.Parameter( torch.zeros(var_num, dim), requires_grad=True) torch.nn.init.normal_(self.template, std=.02) self.var_num = var_num def forward(self, x, query=None): B, N, C = x.shape if query is not None: q = self.q(query).reshape( B, query.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3) q = self.q_norm(q) var_num = query.shape[1] else: q = self.q(self.template).reshape(1, self.var_num, self.num_heads, self.head_dim).permute(0, 2, 1, 3) q = self.q_norm(q) q = q.repeat(B, 1, 1, 1) var_num = self.var_num kv = self.kv(x).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) k, v = kv.unbind(0) k = self.k_norm(k) # Ensure memory continuity and use stable attention backend to prevent CUDA illegal memory access q, k, v = q.contiguous(), k.contiguous(), v.contiguous() with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): x = F.scaled_dot_product_attention( q, k, v, dropout_p=self.attn_drop.p if self.training else 0., ) x = x.transpose(1, 2).reshape(B, var_num, C) x = self.proj(x) x = self.proj_drop(x) return x class FeedFoward(nn.Module): def __init__( self, dim, hidden_features=None, out_features=None, act_layer=nn.GELU, norm_layer=None, bias=True, drop=0., prefix_token_length=None, group=1, ): super().__init__() dim = dim hidden_features = hidden_features or 4*dim out_features = out_features or dim bias = to_2tuple(bias) drop_probs = to_2tuple(drop) self.fc1 = nn.Linear(dim, hidden_features, bias=bias[0]) self.act = act_layer() self.drop1 = nn.Dropout(drop_probs[0]) self.norm = norm_layer( hidden_features) if norm_layer is not None else nn.Identity() self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1]) self.drop2 = nn.Dropout(drop_probs[1]) self.hidden_features = hidden_features self.prefix_token_length = prefix_token_length def forward(self, x): n, var, l, d = x.shape # x = x.view(-1, d) # (n*var, l, c) # x = x.transpose(-1, -2) # (n*var, c, l) x = self.fc1(x) x = self.act(x) x = self.drop1(x) x = self.fc2(x) x = self.norm(x) x = self.drop2(x)+x return x class LearnablePositionalEmbedding(nn.Module): def __init__(self, d_model, max_len=5000): super(LearnablePositionalEmbedding, self).__init__() # Compute the positional encodings once in log space. self.pe = nn.Parameter(torch.zeros( 1, 1, max_len, d_model), requires_grad=True) pe = torch.zeros(max_len, d_model).float() position = torch.arange(0, max_len).float().unsqueeze(1) div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp() pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0).unsqueeze(0) self.pe.data.copy_(pe.float()) del pe def forward(self, x, offset=0): return self.pe[:, :, offset:offset+x.size(2)] 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) # Ensure memory continuity and use stable attention backend to prevent CUDA illegal memory access q, k, v = q.contiguous(), k.contiguous(), v.contiguous() with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): 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 VarAttention(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): B, N, P, C = x.shape qkv = self.qkv(x).reshape(B, N, P, 3, self.num_heads, self.head_dim).permute(3, 0, 2, 4, 1, 5) q, k, v = qkv.unbind(0) q, k = self.q_norm(q), self.k_norm(k) q = q.mean(dim=1, keepdim=False) k = k.mean(dim=1, keepdim=False) v = v.permute(0, 2, 3, 4, 1).reshape(B, self.num_heads, N, -1) # Ensure memory continuity and use stable attention backend to prevent CUDA illegal memory access q, k, v = q.contiguous(), k.contiguous(), v.contiguous() with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): x = F.scaled_dot_product_attention( q, k, v, dropout_p=self.attn_drop.p if self.training else 0., ) # Robust reshape to handle dimension mismatches # Expected: x has shape [B, num_heads, N, P * head_dim] # Target: [B, N, P, C] where C = num_heads * head_dim expected_last_dim = P * self.head_dim actual_last_dim = x.shape[-1] if actual_last_dim == expected_last_dim: # Standard case: reshape back to separate P dimension x = x.view(B, self.num_heads, N, self.head_dim, P).permute(0, 2, 4, 1, 3).reshape(B, N, P, -1) else: # Handle dimension mismatch case # When attention produces unexpected dimensions, we adapt by using a simpler reshape # This preserves the total information while ensuring dimensional consistency x = x.transpose(1, 2) # [B, N, num_heads, actual_last_dim] # Flatten the last two dimensions and then reshape to target flattened_dim = self.num_heads * actual_last_dim x = x.reshape(B, N, flattened_dim) # [B, N, flattened_dim] # Now we reshape to [B, N, P, remaining_dim] where remaining_dim = flattened_dim / P if flattened_dim % P == 0: remaining_dim = flattened_dim // P x = x.reshape(B, N, P, remaining_dim) else: # If not evenly divisible, we keep P=1 and adjust accordingly # This is a fallback that preserves all information x = x.reshape(B, N, 1, flattened_dim) # Expand to P patches by repeating (not ideal but prevents crash) x = x.expand(B, N, P, flattened_dim) x = self.proj(x) x = self.proj_drop(x) return x class GateLayer(nn.Module): def __init__(self, dim, init_values=1e-5, inplace=False): super().__init__() self.inplace = inplace self.gate = nn.Linear(dim, 1) def forward(self, x): gate_value = self.gate(x) return gate_value.sigmoid() * 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 = 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() # self.proj = nn.Linear(dim, dim) def forward(self, x, attn_mask): x_input = x x = self.norm1(x) n_vars, n_seqs = x.shape[1], x.shape[2] x = torch.reshape( x, (-1, x.shape[-2], x.shape[-1])) x = self.attn_seq(x, attn_mask) x = torch.reshape( x, (-1, n_vars, n_seqs, x.shape[-1])) x = x_input + self.drop_path1(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 = VarAttention( 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.ls1 = GateLayer(dim, init_values=init_values) self.drop_path1 = DropPath( drop_path) if drop_path > 0. else nn.Identity() # self.proj = nn.Linear(dim, dim) def forward(self, x): x = x + self.drop_path1(self.attn_var(self.norm1(x))) return x class MLPBlock(nn.Module): def __init__( self, dim, mlp_ratio=4., proj_drop=0., init_values=None, drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, mlp_layer=None, prefix_token_length=0, ): super().__init__() self.norm2 = norm_layer(dim) if mlp_layer is FeedFoward: self.mlp = mlp_layer( in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=proj_drop, prefix_token_length=prefix_token_length, ) else: self.mlp = mlp_layer( in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=proj_drop, ) self.ls2 = GateLayer(dim, init_values=init_values) self.drop_path2 = DropPath( drop_path) if drop_path > 0. else nn.Identity() def forward(self, x, prefix_seq_len=None): if prefix_seq_len is not None: x = x + \ self.drop_path2( self.ls2(self.mlp(self.norm2(x), prefix_seq_len=prefix_seq_len))) else: x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x)))) return x class BasicBlock(nn.Module): def __init__( self, dim, num_heads, mlp_ratio=8., qkv_bias=False, qk_norm=False, proj_drop=0., attn_drop=0., init_values=None, drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, prefix_token_length=0, ): super().__init__() self.seq_att_block = SeqAttBlock(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, init_values=init_values, proj_drop=proj_drop, drop_path=drop_path, norm_layer=norm_layer) self.var_att_block = VarAttBlock(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, init_values=init_values, proj_drop=proj_drop, drop_path=drop_path, norm_layer=norm_layer) self.feed_forward = FeedFoward(dim=dim, hidden_features=dim*4, act_layer=act_layer, drop=proj_drop) def forward(self, x, prefix_seq_len, attn_mask): x = self.var_att_block(x) x = self.seq_att_block(x, attn_mask) x = self.feed_forward(x) return x class Patchfy(nn.Module): def __init__(self, patch_len, stride): super(Patchfy, self).__init__() self.patch_len = patch_len self.stride = stride assert self.patch_len == self.stride, "non-overlap" def forward(self, x): x = x.transpose(1, 2) x = x.unfold(dimension=-1, size=self.patch_len, step=self.stride) # x = x.transpose(1, 2) return x class Model(nn.Module): def __init__(self, args): super(Model, self).__init__() self.patchfy = Patchfy(args.patch_len, args.stride) self.layers = nn.ModuleList([ BasicBlock( dim=args.d_model, num_heads=args.n_heads, mlp_ratio=4., qkv_bias=True, qk_norm=False, proj_drop=args.dropout, attn_drop=args.dropout, init_values=None, drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, prefix_token_length=0 ) for _ in range(args.e_layers) ]) self.norm = nn.LayerNorm(args.d_model) self.patch_embedding = nn.Sequential( nn.Linear(args.patch_len, args.d_model, bias=False), nn.Dropout(args.dropout) ) self.pretrain = getattr(args, 'pretrain', False) self.min_mask_ratio = getattr(args, 'min_mask_ratio', 0.7) self.max_mask_ratio = getattr(args, 'max_mask_ratio', 0.8) self.proj_head = nn.Linear(args.d_model, args.patch_len) # self.times_project = nn.Linear(6*args.d_model, args.d_model) def choose_masking(self, x, min_mask_ratio, max_mask_ratio): # Generate a random number to decide which masking function to use # if torch.rand(1).item() > right_prob: # return self.random_masking(x, min_mask_ratio, max_mask_ratio) # else: # return self.right_masking(x, min_mask_ratio, max_mask_ratio) return self.random_masking(x,min_mask_ratio,max_mask_ratio) def random_masking(self, x, min_mask_ratio, max_mask_ratio): """ Perform random masking where a specified ratio of the total V*L blocks are masked. """ N, V, L, D = x.shape # batch, var, length, dim total_elements = V * L mask_ratio = (min_mask_ratio+max_mask_ratio)/2 # Calculate the number of elements to keep based on the mask ratio total_keeps = int((1 - mask_ratio) * total_elements) # Generate a random noise array for each sample in the batch noise = torch.rand(N, V, L, device=x.device) # noise in [0, 1] for V*L blocks # Flatten noise for easier processing noise_flat = noise.view(N, V * L) # Get indices to sort and restore noise ids_shuffle = torch.argsort(noise_flat, dim=1) ids_restore = torch.argsort(ids_shuffle, dim=1) # Create the binary mask: 0 is keep, 1 is remove # We create a range tensor and compare it with total_keeps to generate the mask range_tensor = torch.arange(V * L, device=x.device).repeat(N, 1) mask_flat = range_tensor >= total_keeps # Unshuffle to get the binary mask in original order mask_flat = mask_flat.gather(dim=1, index=ids_restore) # Reshape mask back to the original V, L dimensions mask = mask_flat.view(N, V, L) return mask def encode(self, x): B,n_vars,N,C = x.shape for layer in self.layers: x = layer(x, prefix_seq_len=None, attn_mask=None) x = self.norm(x) # x = x.view(B // n_vars, n_vars , N* C) # x = self.vars_project(x) return x def forward(self, ts_values): #X_enc [B,L,D] x = ts_values x = self.patchfy(x) if self.pretrain: orin_x = x mask = self.choose_masking(x,self.min_mask_ratio, self.max_mask_ratio) mask_repeat = mask.unsqueeze(dim=-1) #[B,D,N,1] mask_repeat = mask_repeat.repeat(1, 1, 1, x.shape[-1])#[B,D,N,d] #进行掩码 x = x.masked_fill(mask_repeat, 0) x = self.patch_embedding(x) x = self.encode(x) if self.pretrain: predict_x = self.proj_head(x) loss = F.mse_loss(predict_x, orin_x, reduction='mean') return CausalLMOutputWithPast(loss=loss, logits=x) return CausalLMOutputWithPast(logits=x,loss=None) def test_model(): import yaml import argparse #读取args parser = argparse.ArgumentParser(description='TsEncoder Pretrain') parser.add_argument('--fix_seed', type=int, default=None, help='seed') # model settings parser.add_argument('--d_model', type=int, default=512, help='dimension of model') parser.add_argument('--n_heads', type=int, default=8, help='num of heads') parser.add_argument('--e_layers', type=int, default=4, help='num of encoder layers') parser.add_argument("--patch_len", type=int, default=100) parser.add_argument("--stride", type=int, default=100) parser.add_argument("--input_len", type=int, default=600) parser.add_argument('--dropout', type=float, default=0.1, help='dropout') parser.add_argument('--pretrain', type=bool, default=True, help='pretrain mode') parser.add_argument('--min_mask_ratio', type=float, default=0.1, help='minimum mask ratio') parser.add_argument('--max_mask_ratio', type=float, default=0.3, help='maximum mask ratio') args = parser.parse_args() # 创建一个Model实例 model = Model(args) # 创建一些假的输入数据 x_enc = torch.randn(10, 600, 33) # 假设有10个样本，每个样本有50个时间步，每个时间步有100个特征 # 调用forward方法 dec_out = model.forward(x_enc) if 'loss' in dec_out: print(f"Loss: {dec_out['loss']}") print(f"Logits Shape: {dec_out['logits'].shape}") else: print(f"Logits Shape: {dec_out['logits'].shape}") # 在脚本的最后调用测试函数 if __name__ == '__main__': test_model()