Args: x: tensor of time series data. patch_size: Size of each patch. patch_stride: Stride of the patches. context_ratio_range: Tuple indicating the min and max ratio of context patches. target_ratio_range: Tuple indicating the min and max ratio of target patches. min_block_size: Minimum size of a block. max_block_size: Maximum size of a block (defaults to 25% of num_patches). return_nested: Whether to return nested tensors or 0 padded tensors. Returns: masks (target_masks): Tensor of target indices for each sequence in the batch. non_masks (context_masks): Tensor of context indices for each sequence in the batch.
Creates different masks for each batch item, with consistent sizes.
Args: bs: Batch size num_patches: Number of patches in sequence max_context_distance: Maximum distance between a target and its context patches target_ratio_range: Range for ratio of patches to use as targets context_ratio_range: Range for ratio of context to select from available context
Returns: mask_indices: Target patch indices [bs, num_targets] non_mask_indices: Context patch indices [bs, num_context]
Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
self.conv1 = nn.Conv2d(1, 20, 5)
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.
.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.
:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool
enc_model = Encoder( c_in=7, num_patches=500, patch_size=10, patch_stride=10, d_model=128, nhead=2, num_layers=1, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, norm_layer=nn.LayerNorm, jepa=True, embed_activation=nn.GELU(), tokenizer_type='simple', pe_type ='tAPE', use_tst_block=True, shared_embedding=False, )batch_size =2n_vars =7max_len =5000x = torch.randn(batch_size, n_vars, max_len)seq_lens = torch.randint(500, max_len, (batch_size,))# Create input tensors with different sequence lengthsx_list = [torch.randn(n_vars, length) for length in seq_lens]x_nested = torch.nested.as_nested_tensor(x_list, layout=torch.jagged)masks, non_masks = create_masks(x_nested, patch_size=10, patch_stride=10, context_mask_range=[0.1,0.1], target_mask_range=[0.1,0.1], melt_channels_to_batch=True, return_nested=True)#generate_block_masks(x=x, patch_size=self.patch_size, patch_stride=self.patch_stride, context_ratio_range=self.context_mask_range, target_ratio_range=self.target_mask_range, min_block_size=self.mask_block_range[0], max_block_size=self.mask_block_range[1], melt_channels_to_batch=self.melt_channels_to_batch, return_nested=x.is_nested)#print(masks.shape, non_masks.shape)#print([i.shape for i in x_nested.unbind()])#print([i for i in masks.unbind()], [i.shape for i in non_masks.unbind()])enc_model(x_nested, mask=non_masks)
/opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/torch/nested/__init__.py:109: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/NestedTensorImpl.cpp:182.)
return torch._nested_tensor_from_tensor_list(ts, dtype, None, device, None)
predictor = Predictor( num_patches=500, encoder_embed_dim=128, predictor_embed_dim=128, nhead=2, num_layers=1, pe_type='tAPE', c_in_mask_tokens=7,)batch_size =2n_vars =7max_len =5000x = torch.randn(batch_size, n_vars, max_len)seq_lens = torch.randint(4000, max_len, (batch_size,))# Create input tensors with different sequence lengthsx_list = [torch.randn(n_vars, length) for length in seq_lens]x_nested = torch.nested.as_nested_tensor(x_list, layout=torch.jagged)masks, non_masks = create_masks(x_nested, patch_size=10, patch_stride=10, context_mask_range=[0.1,0.1], target_mask_range=[0.1,0.1], melt_channels_to_batch=True, return_nested=True)#generate_block_masks(x=x, patch_size=self.patch_size, patch_stride=self.patch_stride, context_ratio_range=self.context_mask_range, target_ratio_range=self.target_mask_range, min_block_size=self.mask_block_range[0], max_block_size=self.mask_block_range[1], melt_channels_to_batch=self.melt_channels_to_batch, return_nested=x.is_nested)encoded_vals = enc_model(x, mask=non_masks)print(encoded_vals.shape)r = predictor(encoded_vals, mask=masks, non_masks=non_masks)print([i.shape for i in r], [i.shape for i in masks])
torch.Size([2, j141, 128])
---------------------------------------------------------------------------IndexError Traceback (most recent call last)
Cell In[39], line 27 25 encoded_vals = enc_model(x, mask=non_masks)
26print(encoded_vals.shape)
---> 27 r = predictor(encoded_vals, mask=masks, non_masks=non_masks)
28print([i.shape for i in r], [i.shape for i in masks])
File /opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs) 1737returnself._compiled_call_impl(*args, **kwargs) # type: ignore[misc] 1738else:
-> 1739returnself._call_impl(*args, **kwargs)
File /opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs) 1745# If we don't have any hooks, we want to skip the rest of the logic in 1746# this function, and just call forward. 1747ifnot (self._backward_hooks orself._backward_pre_hooks orself._forward_hooks orself._forward_pre_hooks
1748or _global_backward_pre_hooks or _global_backward_hooks
1749or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750return forward_call(*args, **kwargs)
1752 result =None 1753 called_always_called_hooks =set()
Cell In[37], line 151, in Predictor.forward(self, encoded_vals, mask, non_masks) 149 pred_tokens = []
150for i, target_pos_enc_i inenumerate(target_pos_enc):
--> 151 t = bs_pred_tokens[i].repeat(target_pos_enc_i.size(0), 1)
152 t = t + target_pos_enc_i
153 pred_tokens.append(t)
IndexError: index 0 is out of bounds for dimension 0 with size 0
def JEPASimpleLightning( learning_rate, train_size, batch_size, n_gpus, patchtsjepa_encoder_kwargs, patchtsjepa_predictor_kwargs, num_nodes:int=1, weight_decay:float=0.04, use_weight_decay_scheduler:bool=False, final_weight_decay:float=0.4, epochs:int=100, optimizer_type:str='adamw', scheduler_type:str='OneCycle', target_mask_range:tuple=(0.05, 0.3), # the target can be up to 50% of the original x context_mask_range:tuple=(0.5, 1.0), # the context can be up to 80% of masked out target (1-target_mask_ratio) mask_block_range:tuple=(1, 30), ema_decay:float=0.996, scheduler_kwargs:dict={}, transforms:NoneType=None, loss_fn:function=loss_pred, linear_probe:bool=False):
/opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/lightning/pytorch/core/module.py:441: You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`
/opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/lightning/pytorch/core/module.py:441: You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`
grid_size_h: int of the grid height grid_size_w: int of the grid width return: pos_embed: [grid_size_h*grid_size_w, embed_dim] or [1+grid_size_h*grid_size_w, embed_dim] (w/ or w/o cls_token)
def PatchTSJEPAPredictor( c_in, # number of input channels num_patches, # number of patches from encoder d_model:int=512, # encoder embedding dimension predictor_dim:int=384, # predictor embedding dimension (typically smaller) n_heads:int=4, n_layers:int=2, d_ff:int=1024, pos_encoding_type:str='learned', # 'learned' or 'tAPE' dropout:float=0.1, attn_dropout:float=0.0, act:str='gelu', pre_norm:bool=False, init_std:float=0.02):
Predictor network for PatchTSJEPA encoder for time series adaptation
model = PatchTSJEPAPredictor(c_in=7, num_patches=480, d_model=512, predictor_dim=384, n_heads=8, n_layers=4, d_ff=2048, dropout=0.1, attn_dropout=0., pos_encoding_type='learned', act="gelu", pre_norm=False)x = torch.randn(1, 7, 512, 480)z = model(x)batch_size =2n_vars =7max_len =480d_model =512# Create sequences of different lengthsseq_lens = torch.randint(50, max_len, (batch_size,))# Create input tensors with different sequence lengthsx_list = [torch.randn(n_vars, d_model, length) for length in seq_lens]x_nested = torch.nested.as_nested_tensor(x_list, layout=torch.jagged)z = model(x_nested)
/opt/miniconda3/envs/timeflies/lib/python3.12/site-packages/torch/nested/__init__.py:109: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/NestedTensorImpl.cpp:182.)
return torch._nested_tensor_from_tensor_list(ts, dtype, None, device, None)
def PatchTSJEPAEncoder( c_in:int, # the number of input channels win_length, # the length of the patch of time/interval or short time ft windown length (when time_domain=False) hop_length, # the length of the distance between each patch/fft max_seq_len, # maximum sequence len pos_encoding_type:str='learned', # 'learned' or 'tAPE' patch_encoder_type:str='linear', # 'linear' or 'conv' use_revin:bool=True, # if time_domain is true, whether or not to instance normalize time data affine:bool=True, # if time_domain is true, whether or not to learn revin normalization parameters n_layers:int=4, # the number of transformer encoder layers to use d_model:int=512, # the dimension of the input to the transofmrer encoder n_heads:int=8, # the number of heads in each layer shared_embedding:bool=False, # indicator for whether or not each channel should be projected with its own set of linear weights to the encoder dimension d_ff:int=2048, # the feedforward layer size in the transformer attn_dropout:float=0.0, # dropout in attention dropout:float=0.1, # dropout for linear layers act:str='gelu', # activation function pre_norm:bool=False, # indicator to pre batch or layer norm init_std:float=0.02):
Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
self.conv1 = nn.Conv2d(1, 20, 5)
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.
.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.
:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool
model = PatchTSJEPAEncoder(c_in=7, win_length=750, hop_length=750, max_seq_len=1*3600*100, use_revin=True, affine=True, n_layers=4, patch_encoder_type='linear', d_model=256, n_heads=8, shared_embedding=False, d_ff=2048, attn_dropout=0., dropout=0.1, act="gelu", pre_norm=True)x = torch.randn(2, 7, 1*3600*100)# z = model(x)# print(z.shape)batch_size =2n_vars =7max_len =1*3600*100# Create sequences of different lengthsseq_lens = torch.randint(50, max_len, (batch_size,))#seq_lens = [max_len-10000, max_len-10000]# Create input tensors with different sequence lengthsx_list = [torch.randn(length, n_vars) for length in seq_lens]x_nested = torch.nested.as_nested_tensor(x_list, layout=torch.jagged)x_nested = x_nested.transpose(1,2)z = model(x_nested)z.shape
def PatchTSJEPA( encoder_kwargs:dict, predictor_kwargs:dict, pretrain:bool=True, target_mask_range:tuple=(0.1, 0.5), # the target can be up to 50% of the original x context_mask_range:tuple=(0.2, 0.8), # the context can be up to 80% of masked out target (1-target_mask_ratio)):
Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
self.conv1 = nn.Conv2d(1, 20, 5)
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.
.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.
:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool
def PatchTSJEPALightning( learning_rate, train_size, batch_size, channels, patchtsjepa_encoder_kwargs, patchtsjepa_predictor_kwargs, loss_func, max_lr:float=0.01, weight_decay:float=0.0, epochs:int=100, optimizer_type:str='adamw', scheduler_type:str='OneCycle', target_mask_range:tuple=(0.1, 0.5), # the target can be up to 50% of the original x context_mask_range:tuple=(0.2, 0.8), # the context can be up to 80% of masked out target (1-target_mask_ratio) pretrain:bool=True, ema_decay:float=0.996):
