Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 IBM & Hugging Face. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch PatchTST model."""16 17import math18from dataclasses import dataclass19from typing import Callable, Optional, Union20 21import torch22from torch import nn23 24from ...activations import ACT2CLS25from ...modeling_flash_attention_utils import FlashAttentionKwargs26from ...modeling_outputs import BaseModelOutput27from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel28from ...processing_utils import Unpack29from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput30from ...utils import ModelOutput, auto_docstring, logging31from .configuration_patchtst import PatchTSTConfig32 33 34logger = logging.get_logger(__name__)35 36 37# Copied from transformers.models.bart.modeling_bart.eager_attention_forward38def eager_attention_forward(39 module: nn.Module,40 query: torch.Tensor,41 key: torch.Tensor,42 value: torch.Tensor,43 attention_mask: Optional[torch.Tensor],44 scaling: Optional[float] = None,45 dropout: float = 0.0,46 head_mask: Optional[torch.Tensor] = None,47 **kwargs,48):49 if scaling is None:50 scaling = query.size(-1) ** -0.551 52 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling53 if attention_mask is not None:54 attn_weights = attn_weights + attention_mask55 56 attn_weights = nn.functional.softmax(attn_weights, dim=-1)57 58 if head_mask is not None:59 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)60 61 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)62 attn_output = torch.matmul(attn_weights, value)63 attn_output = attn_output.transpose(1, 2).contiguous()64 65 return attn_output, attn_weights66 67 68# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Attention with Wav2Vec2->PatchTST69class PatchTSTAttention(nn.Module):70 """Multi-headed attention from 'Attention Is All You Need' paper"""71 72 def __init__(73 self,74 embed_dim: int,75 num_heads: int,76 dropout: float = 0.0,77 is_decoder: bool = False,78 bias: bool = True,79 is_causal: bool = False,80 config: Optional[PatchTSTConfig] = None,81 ):82 super().__init__()83 self.embed_dim = embed_dim84 self.num_heads = num_heads85 self.dropout = dropout86 self.head_dim = embed_dim // num_heads87 self.config = config88 89 if (self.head_dim * num_heads) != self.embed_dim:90 raise ValueError(91 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"92 f" and `num_heads`: {num_heads})."93 )94 self.scaling = self.head_dim**-0.595 self.is_decoder = is_decoder96 self.is_causal = is_causal97 98 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)99 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)100 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)101 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)102 103 def forward(104 self,105 hidden_states: torch.Tensor,106 key_value_states: Optional[torch.Tensor] = None,107 attention_mask: Optional[torch.Tensor] = None,108 layer_head_mask: Optional[torch.Tensor] = None,109 output_attentions: Optional[bool] = False,110 # TODO: we need a refactor so that the different attention modules can get their specific kwargs111 # ATM, we have mixed things encoder, decoder, and encoder-decoder attn112 **kwargs: Unpack[FlashAttentionKwargs],113 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:114 """Input shape: Batch x Time x Channel"""115 116 # if key_value_states are provided this layer is used as a cross-attention layer117 # for the decoder118 is_cross_attention = key_value_states is not None119 120 # determine input shapes121 bsz, tgt_len = hidden_states.shape[:-1]122 src_len = key_value_states.shape[1] if is_cross_attention else tgt_len123 124 q_input_shape = (bsz, tgt_len, -1, self.head_dim)125 kv_input_shape = (bsz, src_len, -1, self.head_dim)126 127 # get query proj128 query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)129 130 current_states = key_value_states if is_cross_attention else hidden_states131 key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2)132 value_states = self.v_proj(current_states).view(*kv_input_shape).transpose(1, 2)133 134 attention_interface: Callable = eager_attention_forward135 if self.config._attn_implementation != "eager":136 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]137 138 attn_output, attn_weights = attention_interface(139 self,140 query_states,141 key_states,142 value_states,143 attention_mask,144 dropout=0.0 if not self.training else self.dropout,145 scaling=self.scaling,146 output_attentions=output_attentions,147 head_mask=layer_head_mask,148 **kwargs,149 )150 151 attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()152 attn_output = self.out_proj(attn_output)153 154 return attn_output, attn_weights, None155 156 157class PatchTSTBatchNorm(nn.Module):158 """159 Compute batch normalization over the sequence length (time) dimension.160 """161 162 def __init__(self, config: PatchTSTConfig):163 super().__init__()164 self.batchnorm = nn.BatchNorm1d(config.d_model, eps=config.norm_eps)165 166 def forward(self, inputs: torch.Tensor):167 """168 Parameters:169 inputs (`torch.Tensor` of shape `(batch_size, sequence_length, d_model)`):170 input for Batch norm calculation171 Returns:172 `torch.Tensor` of shape `(batch_size, sequence_length, d_model)`173 """174 output = inputs.transpose(1, 2) # output: (batch_size, d_model, sequence_length)175 output = self.batchnorm(output)176 return output.transpose(1, 2)177 178 179def random_masking(180 inputs: torch.Tensor,181 mask_ratio: float,182 unmasked_channel_indices: Optional[list] = None,183 channel_consistent_masking: bool = False,184 mask_value: int = 0,185):186 """random_masking: Mask the input considering the control variables.187 188 Args:189 inputs (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, num_features)`):190 The input tensor to mask.191 mask_ratio (`float`):192 Masking ratio applied to mask the input data during random pretraining. It is the number between 0 and 1.193 unmasked_channel_indices (list, *optional*):194 Indices of channels that will not be masked.195 channel_consistent_masking (bool, *optional*, defaults to `False`):196 When true, masking will be same across all channels of a timeseries. Otherwise, masking positions will vary197 across channels.198 mask_value (int, *optional*, defaults to 0):199 Define the value of masked patches for pretraining.200 201 Returns:202 `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as input Tensor and mask tensor of shape [bs x c x203 n]204 """205 if mask_ratio < 0 or mask_ratio >= 1:206 raise ValueError(f"Mask ratio {mask_ratio} has to be between 0 and 1.")207 208 batch_size, num_channels, sequence_length, num_features = inputs.shape209 device = inputs.device210 211 len_keep = int(sequence_length * (1 - mask_ratio))212 213 if channel_consistent_masking:214 noise = torch.rand(batch_size, 1, sequence_length, device=device) # noise in [0, 1], bs x 1 x L215 noise = noise.repeat(1, num_channels, 1) # bs x num_channels x time216 else:217 # noise in [0, 1], bs x num_channels x L218 noise = torch.rand(batch_size, num_channels, sequence_length, device=device)219 220 # mask: [bs x num_channels x num_patch]221 mask = torch.ones(batch_size, num_channels, sequence_length, device=device)222 mask[:, :, :len_keep] = 0223 224 # sort noise for each sample225 ids_shuffle = torch.argsort(noise, dim=-1) # ascend: small is keep, large is remove226 ids_restore = torch.argsort(ids_shuffle, dim=-1) # ids_restore: [bs x num_channels x L]227 228 mask = torch.gather(mask, dim=-1, index=ids_restore)229 mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features) # mask: [bs x num_channels x num_patches x patch_length]230 if unmasked_channel_indices is not None:231 mask[:, unmasked_channel_indices, :, :] = 0232 233 inputs_mask = inputs.masked_fill(mask.bool(), mask_value)234 return inputs_mask, mask[..., 0]235 236 237def forecast_masking(238 inputs: torch.Tensor,239 num_forecast_mask_patches: Union[list, int],240 unmasked_channel_indices: Optional[list] = None,241 mask_value: int = 0,242):243 """Forecast masking that masks the last K patches where K is from the num_forecast_mask_patches.244 If num_forecast_mask_patches is a list, samples in the batch will be randomly masked by numbers defined in the list.245 246 Parameters:247 inputs (`torch.Tensor`):248 Input of shape `(bs, num_channels, num_patch, patch_length)`249 num_forecast_mask_patches (`list`):250 Number of patches to be masked at the end of each batch sample. e.g. 4 or [3, 5].251 unmasked_channel_indices (`list`, *optional*):252 Indices of channels that are not masked.253 mask_value (`int`, *optional*, defaults to 0):254 Values in the masked patches will be filled by `mask_value`.255 256 Returns:257 `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as inputs Tensor and Mask tensor of shape `(bs,258 num_channels , num_patch)` or `(bs, tsg1, tsg2, num_channels, num_patch)`259 """260 261 if isinstance(num_forecast_mask_patches, int):262 num_forecast_mask_patches = [num_forecast_mask_patches]263 forecast_mask_ratios = [1 for _ in num_forecast_mask_patches]264 265 batch_size, num_channels, sequence_length, num_features = inputs.shape266 mask = torch.zeros(batch_size, num_channels, sequence_length, device=inputs.device)267 268 t_list = []269 total_length = 0270 total_ratio = sum(forecast_mask_ratios)271 272 for patch_length, ratio in zip(num_forecast_mask_patches, forecast_mask_ratios):273 if patch_length <= 0 or patch_length >= sequence_length:274 raise ValueError(275 f"num_forecast_mask_patches {patch_length} should be greater than 0 and less than total patches."276 )277 temp_len = int(batch_size * ratio / total_ratio)278 t_list.append([patch_length, ratio, temp_len])279 total_length += temp_len280 281 t_list = sorted(t_list, key=lambda x: x[2])282 283 if total_length < batch_size:284 t_list[0][2] = t_list[0][2] + (batch_size - total_length)285 elif total_length > batch_size:286 t_list[-1][2] = t_list[-1][2] + (total_length - batch_size)287 288 batch1 = 0289 for patch_len, _, temp_len in t_list:290 batch2 = batch1 + temp_len291 mask[batch1:batch2, :, -patch_len:] = 1292 batch1 = batch2293 294 perm = torch.randperm(mask.shape[0])295 mask = mask[perm]296 297 mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features) # mask: [bs x num_channels x num_patch x patch_len]298 if unmasked_channel_indices is not None:299 mask[:, unmasked_channel_indices, :, :] = 0300 301 inputs_mask = inputs.masked_fill(mask.bool(), mask_value)302 return inputs_mask, mask[..., 0]303 304 305class PatchTSTPatchify(nn.Module):306 """307 A class to patchify the time series sequence into different patches308 309 Returns:310 `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`311 """312 313 def __init__(self, config: PatchTSTConfig):314 super().__init__()315 316 self.sequence_length = config.context_length317 self.patch_length = config.patch_length318 self.patch_stride = config.patch_stride319 320 if self.sequence_length <= self.patch_length:321 raise ValueError(322 f"Sequence length ({self.sequence_length}) has to be greater than the patch length ({self.patch_length})"323 )324 325 # get the number of patches326 self.num_patches = (max(self.sequence_length, self.patch_length) - self.patch_length) // self.patch_stride + 1327 new_sequence_length = self.patch_length + self.patch_stride * (self.num_patches - 1)328 self.sequence_start = self.sequence_length - new_sequence_length329 330 def forward(self, past_values: torch.Tensor):331 """332 Parameters:333 past_values (`torch.Tensor` of shape `(batch_size, sequence_length, num_channels)`, *required*):334 Input for patchification335 336 Returns:337 `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`338 """339 sequence_length = past_values.shape[-2]340 if sequence_length != self.sequence_length:341 raise ValueError(342 f"Input sequence length ({sequence_length}) doesn't match model configuration ({self.sequence_length})."343 )344 # output: [bs x new_sequence_length x num_channels]345 output = past_values[:, self.sequence_start :, :]346 # output: [bs x num_patches x num_input_channels x patch_length]347 output = output.unfold(dimension=-2, size=self.patch_length, step=self.patch_stride)348 # output: [bs x num_input_channels x num_patches x patch_length]349 output = output.transpose(-2, -3).contiguous()350 return output351 352 353class PatchTSTMasking(nn.Module):354 """355 Class to perform random or forecast masking.356 357 Parameters:358 config (`PatchTSTConfig`): model config359 Returns:360 x_mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)361 Masked patched input362 mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)363 Bool tensor indicating True on masked points364 """365 366 def __init__(self, config: PatchTSTConfig):367 super().__init__()368 self.random_mask_ratio = config.random_mask_ratio369 self.channel_consistent_masking = config.channel_consistent_masking370 self.mask_type = config.mask_type371 self.num_forecast_mask_patches = config.num_forecast_mask_patches372 self.unmasked_channel_indices = config.unmasked_channel_indices373 self.mask_value = config.mask_value374 if self.unmasked_channel_indices is not None:375 self.unmasked_channel_indices = sorted(self.unmasked_channel_indices)376 377 def forward(self, patch_input: torch.Tensor):378 """379 Parameters:380 patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):381 Patch input382 383 Return:384 masked_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)385 Masked patched input386 mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)387 Bool tensor indicating True on masked points388 389 """390 if self.mask_type == "random":391 masked_input, mask = random_masking(392 inputs=patch_input,393 mask_ratio=self.random_mask_ratio,394 unmasked_channel_indices=self.unmasked_channel_indices,395 channel_consistent_masking=self.channel_consistent_masking,396 mask_value=self.mask_value,397 )398 elif self.mask_type == "forecast":399 masked_input, mask = forecast_masking(400 inputs=patch_input,401 num_forecast_mask_patches=self.num_forecast_mask_patches,402 unmasked_channel_indices=self.unmasked_channel_indices,403 mask_value=self.mask_value,404 )405 else:406 raise ValueError(f"Invalid mask type {self.mask_type}.")407 408 # mask: [bs x num_input_channels x num_patch]409 mask = mask.bool()410 return masked_input, mask411 412 413class PatchTSTEncoderLayer(nn.Module):414 """415 PatchTST encoder layer416 """417 418 def __init__(self, config: PatchTSTConfig):419 super().__init__()420 421 self.channel_attention = config.channel_attention422 # Multi-Head attention423 self.self_attn = PatchTSTAttention(424 embed_dim=config.d_model,425 num_heads=config.num_attention_heads,426 dropout=config.attention_dropout,427 config=config,428 )429 430 # Add & Norm of the sublayer 1431 self.dropout_path1 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()432 if config.norm_type == "batchnorm":433 self.norm_sublayer1 = PatchTSTBatchNorm(config)434 elif config.norm_type == "layernorm":435 self.norm_sublayer1 = nn.LayerNorm(config.d_model, eps=config.norm_eps)436 else:437 raise ValueError(f"{config.norm_type} is not a supported norm layer type.")438 439 # Add & Norm of the sublayer 2440 if self.channel_attention:441 self.dropout_path2 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()442 if config.norm_type == "batchnorm":443 self.norm_sublayer2 = PatchTSTBatchNorm(config)444 elif config.norm_type == "layernorm":445 self.norm_sublayer2 = nn.LayerNorm(config.d_model, eps=config.norm_eps)446 else:447 raise ValueError(f"{config.norm_type} is not a supported norm layer type.")448 449 # Position-wise Feed-Forward450 self.ff = nn.Sequential(451 nn.Linear(config.d_model, config.ffn_dim, bias=config.bias),452 ACT2CLS[config.activation_function](),453 nn.Dropout(config.ff_dropout) if config.ff_dropout > 0 else nn.Identity(),454 nn.Linear(config.ffn_dim, config.d_model, bias=config.bias),455 )456 457 # Add & Norm of sublayer 3458 self.dropout_path3 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()459 if config.norm_type == "batchnorm":460 self.norm_sublayer3 = PatchTSTBatchNorm(config)461 elif config.norm_type == "layernorm":462 self.norm_sublayer3 = nn.LayerNorm(config.d_model, eps=config.norm_eps)463 else:464 raise ValueError(f"{config.norm_type} is not a supported norm layer type.")465 466 self.pre_norm = config.pre_norm467 468 def forward(self, hidden_state: torch.Tensor, output_attentions: Optional[bool] = None):469 """470 Parameters:471 hidden_state (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`, *required*):472 Past values of the time series473 output_attentions (`bool`, *optional*):474 Whether or not to return the output attention of all layers475 Return:476 `torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`477 478 """479 batch_size, num_input_channels, sequence_length, d_model = hidden_state.shape480 481 # First sublayer: attention across time482 # hidden_states: [(bs*num_channels) x sequence_length x d_model]483 hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)484 485 if self.pre_norm:486 ## Norm and Multi-Head attention and Add residual connection487 attn_output, attn_weights, _ = self.self_attn(488 hidden_states=self.norm_sublayer1(hidden_state), output_attentions=output_attentions489 )490 # Add: residual connection with residual dropout491 hidden_state = hidden_state + self.dropout_path1(attn_output)492 else:493 ## Multi-Head attention and Add residual connection and Norm - Standard Transformer from BERT494 attn_output, attn_weights, _ = self.self_attn(495 hidden_states=hidden_state, output_attentions=output_attentions496 )497 # hidden_states: [(bs*num_channels) x sequence_length x d_model]498 hidden_state = self.norm_sublayer1(hidden_state + self.dropout_path1(attn_output))499 500 # hidden_state: [bs x num_channels x sequence_length x d_model]501 hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)502 503 # second sublayer: attention across variable at any given time504 if self.channel_attention:505 # hidden_state: [bs x sequence_length x num_channels x d_model]506 hidden_state = hidden_state.transpose(2, 1).contiguous()507 # hidden_state: [(bs*sequence_length) x num_channels x d_model]508 hidden_state = hidden_state.view(batch_size * sequence_length, num_input_channels, d_model)509 if self.pre_norm:510 ## Norm and Multi-Head attention and Add residual connection511 attn_output, channel_attn_weights, _ = self.self_attn(512 hidden_states=self.norm_sublayer2(hidden_state), output_attentions=output_attentions513 )514 # Add: residual connection with residual dropout515 hidden_state = hidden_state + self.dropout_path2(attn_output)516 else:517 ## Multi-Head attention and Add residual connection and Norm518 attn_output, channel_attn_weights, _ = self.self_attn(519 hidden_states=hidden_state, output_attentions=output_attentions520 )521 # hidden_states: [(bs*sequence_length) x num_channels x d_model]522 hidden_state = self.norm_sublayer2(hidden_state + self.dropout_path2(attn_output))523 524 # Reshape hidden state525 # hidden_state: [bs x sequence_length x num_channels x d_model]526 hidden_state = hidden_state.reshape(batch_size, sequence_length, num_input_channels, d_model)527 # hidden_state: [bs x num_channels x sequence_length x d_model]528 hidden_state = hidden_state.transpose(1, 2).contiguous()529 530 # Third sublayer: mixing across hidden531 # hidden_state: [(batch_size*num_channels) x sequence_length x d_model]532 hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)533 if self.pre_norm:534 ## Norm and Position-wise Feed-Forward and Add residual connection535 # Add: residual connection with residual dropout536 hidden_state = hidden_state + self.dropout_path3(self.ff(self.norm_sublayer3(hidden_state)))537 else:538 ## Position-wise Feed-Forward and Add residual connection and Norm539 # Add: residual connection with residual dropout540 hidden_state = self.norm_sublayer3(hidden_state + self.dropout_path3(self.ff(hidden_state)))541 542 # [bs x num_channels x sequence_length x d_model]543 hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)544 545 outputs = (hidden_state,)546 if output_attentions:547 outputs += (attn_weights, channel_attn_weights) if self.channel_attention else (attn_weights,)548 549 return outputs550 551 552@auto_docstring553class PatchTSTPreTrainedModel(PreTrainedModel):554 config: PatchTSTConfig555 base_model_prefix = "model"556 main_input_name = "past_values"557 supports_gradient_checkpointing = False558 559 def _init_weights(self, module: nn.Module):560 """561 Initialize weights562 """563 if isinstance(module, PatchTSTPositionalEncoding):564 # get the number of patches565 num_patches = (566 max(self.config.context_length, self.config.patch_length) - self.config.patch_length567 ) // self.config.patch_stride + 1568 # initialize cls_token569 if self.config.use_cls_token:570 nn.init.normal_(module.cls_token, std=0.02)571 num_patches += 1572 # initialize positional encoding573 module.position_enc = module._init_pe(self.config, num_patches)574 elif isinstance(module, nn.LayerNorm):575 module.bias.data.zero_()576 module.weight.data.fill_(1.0)577 elif isinstance(module, PatchTSTBatchNorm):578 module.batchnorm.bias.data.zero_()579 module.batchnorm.weight.data.fill_(1.0)580 elif isinstance(module, nn.Linear):581 module.weight.data.normal_(mean=0.0, std=self.config.init_std)582 if module.bias is not None:583 module.bias.data.zero_()584 585 def _set_gradient_checkpointing(self, module, value=False):586 if isinstance(module, (PatchTSTEncoder)):587 module.gradient_checkpointing = value588 589 590class PatchTSTEmbedding(nn.Module):591 def __init__(self, config: PatchTSTConfig):592 super().__init__()593 self.num_input_channels = config.num_input_channels594 self.share_embedding = config.share_embedding595 # Input encoding: projection of feature vectors onto a d-dim vector space596 if self.share_embedding:597 self.input_embedding = nn.Linear(config.patch_length, config.d_model)598 else:599 self.input_embedding = nn.ModuleList()600 for _ in range(config.num_input_channels):601 self.input_embedding.append(nn.Linear(config.patch_length, config.d_model))602 603 def forward(self, patch_input: torch.Tensor):604 """605 Parameters:606 patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):607 Patch input for embedding608 return:609 `torch.Tensor` of shape `(batch_size, num_channels, num_patches, d_model)`610 """611 # Input encoding612 num_input_channels = patch_input.shape[1]613 if num_input_channels != self.num_input_channels:614 raise ValueError(615 f"The defined number of input channels ({self.num_input_channels}) in the config "616 f"has to be the same as the number of channels in the batch input ({num_input_channels})"617 )618 if self.share_embedding:619 embeddings = self.input_embedding(patch_input) # x: [bs x num_channels x num_patches x d_model]620 else:621 embeddings = [self.input_embedding[i](patch_input[:, i, :, :]) for i in range(num_input_channels)]622 embeddings = torch.stack(embeddings, dim=1)623 return embeddings624 625 626class PatchTSTPositionalEncoding(nn.Module):627 """628 Class for positional encoding629 """630 631 def __init__(self, config: PatchTSTConfig, num_patches: int):632 super().__init__()633 self.use_cls_token = config.use_cls_token634 self.num_input_channels = config.num_input_channels635 if config.use_cls_token:636 # cls_token: [1 x num_input_channels x 1 x d_model]637 self.cls_token = nn.Parameter(torch.zeros(1, 1, 1, config.d_model))638 num_patches += 1639 # positional encoding: [num_patches x d_model]640 self.position_enc = self._init_pe(config, num_patches)641 # Positional dropout642 self.positional_dropout = (643 nn.Dropout(config.positional_dropout) if config.positional_dropout > 0 else nn.Identity()644 )645 646 @staticmethod647 def _init_pe(config: PatchTSTConfig, num_patches: int) -> nn.Parameter:648 # Positional encoding649 if config.positional_encoding_type == "random":650 position_enc = nn.Parameter(torch.randn(num_patches, config.d_model), requires_grad=True)651 elif config.positional_encoding_type == "sincos":652 position_enc = torch.zeros(num_patches, config.d_model)653 position = torch.arange(0, num_patches).unsqueeze(1)654 div_term = torch.exp(torch.arange(0, config.d_model, 2) * -(math.log(10000.0) / config.d_model))655 position_enc[:, 0::2] = torch.sin(position * div_term)656 position_enc[:, 1::2] = torch.cos(position * div_term)657 position_enc = position_enc - position_enc.mean()658 position_enc = position_enc / (position_enc.std() * 10)659 position_enc = nn.Parameter(position_enc, requires_grad=False)660 else:661 raise ValueError(662 f"{config.positional_encoding_type} is not a valid positional encoder. Available types are 'random' and 'sincos'."663 )664 return position_enc665 666 def forward(self, patch_input: torch.Tensor):667 if self.use_cls_token:668 # patch_input: [bs x num_channels x num_patches x d_model]669 patch_input = self.positional_dropout(patch_input + self.position_enc[1:, :])670 # append cls token where cls_token: [1 x num_channels x 1 x d_model]671 cls_token = self.cls_token + self.position_enc[:1, :]672 # get the same copy of cls_token for all the samples in batch: [bs x num_channels x 1 x d_model]673 cls_tokens = cls_token.expand(patch_input.shape[0], self.num_input_channels, -1, -1)674 # hidden_state: [bs x num_channels x (num_patches+1) x d_model]675 hidden_state = torch.cat((cls_tokens, patch_input), dim=2)676 else:677 # hidden_state: [bs x num_channels x num_patches x d_model]678 hidden_state = self.positional_dropout(patch_input + self.position_enc)679 return hidden_state680 681 682class PatchTSTEncoder(PatchTSTPreTrainedModel):683 """684 PatchTST Encoder685 """686 687 def __init__(self, config: PatchTSTConfig, num_patches: int):688 super().__init__(config)689 self.gradient_checkpointing = False690 691 # Input embedding: projection of feature vectors onto a d-dim vector space692 self.embedder = PatchTSTEmbedding(config)693 # Positional encoding694 self.positional_encoder = PatchTSTPositionalEncoding(config, num_patches)695 # Encoder696 self.layers = nn.ModuleList([PatchTSTEncoderLayer(config) for i in range(config.num_hidden_layers)])697 698 # Initialize weights and apply final processing699 self.post_init()700 701 def forward(702 self,703 patch_input: torch.Tensor,704 output_hidden_states: Optional[bool] = None,705 output_attentions: Optional[bool] = None,706 ) -> BaseModelOutput:707 """708 Parameters:709 patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):710 Past values of the time series711 output_hidden_states (bool, optional): Indicates if hidden states should be outputted.712 output_attentions (bool, optional): Indicates if attentions should be outputted.713 714 return:715 `BaseModelOutput`716 """717 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions718 output_hidden_states = (719 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states720 )721 722 # Input embedding723 patch_input = self.embedder(patch_input)724 # Positional encoding725 hidden_state = self.positional_encoder(patch_input)726 727 encoder_states = () if output_hidden_states else None728 all_attentions = () if output_attentions else None729 for encoder_layer in self.layers:730 if output_hidden_states:731 encoder_states = encoder_states + (hidden_state,)732 733 layer_outputs = encoder_layer(hidden_state=hidden_state, output_attentions=output_attentions)734 # get hidden state. hidden_state shape is [bs x num_channels x num_patches x d_model]735 # or [bs x num_channels x (num_patches+1) x d_model] if use cls_token736 hidden_state = layer_outputs[0]737 # append attention matrix at each layer738 if output_attentions:739 all_attentions = all_attentions + (layer_outputs[1],)740 # return past_values, hidden_states741 return BaseModelOutput(last_hidden_state=hidden_state, hidden_states=encoder_states, attentions=all_attentions)742 743 744@dataclass745@auto_docstring(746 custom_intro="""747 Base class for model's outputs, with potential hidden states.748 """749)750class PatchTSTModelOutput(ModelOutput):751 r"""752 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`):753 Sequence of hidden-states at the output of the last layer of the model.754 hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):755 Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +756 one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. Hidden-states of757 the model at the output of each layer plus the optional initial embedding outputs.758 mask (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`, *optional*):759 Bool masked tensor indicating which patches are masked760 loc (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*):761 Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length762 scale (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*):763 Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length764 patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`):765 Patched input to the Transformer766 """767 768 last_hidden_state: Optional[torch.FloatTensor] = None769 hidden_states: Optional[tuple[torch.FloatTensor]] = None770 attentions: Optional[tuple[torch.FloatTensor]] = None771 mask: Optional[torch.FloatTensor] = None772 loc: Optional[torch.FloatTensor] = None773 scale: Optional[torch.FloatTensor] = None774 patch_input: Optional[torch.FloatTensor] = None775 776 777@dataclass778@auto_docstring(779 custom_intro="""780 Output type of [`PatchTSTForPretraining`].781 """782)783class PatchTSTForPretrainingOutput(ModelOutput):784 r"""785 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):786 MSE loss.787 prediction_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):788 Prediction outputs of the time series modeling heads.789 """790 791 loss: Optional[torch.FloatTensor] = None792 prediction_output: Optional[torch.FloatTensor] = None793 hidden_states: Optional[tuple[torch.FloatTensor]] = None794 attentions: Optional[tuple[torch.FloatTensor]] = None795 796 797@dataclass798@auto_docstring(799 custom_intro="""800 Output type of [`PatchTSTForRegression`].801 """802)803class PatchTSTForRegressionOutput(ModelOutput):804 r"""805 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):806 MSE loss.807 regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`):808 Regression outputs of the time series modeling heads.809 """810 811 loss: Optional[torch.FloatTensor] = None812 regression_outputs: Optional[torch.FloatTensor] = None813 hidden_states: Optional[tuple[torch.FloatTensor]] = None814 attentions: Optional[tuple[torch.FloatTensor]] = None815 816 817@dataclass818@auto_docstring(819 custom_intro="""820 Output type of [`PatchTSTForPrediction`].821 """822)823class PatchTSTForPredictionOutput(ModelOutput):824 r"""825 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):826 MSE loss.827 prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, -1)`):828 Prediction outputs of the time series modeling heads.829 attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):830 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,831 sequence_length)`.832 833 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention834 heads.835 loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*)836 Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length837 scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*)838 Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length839 """840 841 loss: Optional[torch.FloatTensor] = None842 prediction_outputs: Optional[torch.FloatTensor] = None843 hidden_states: Optional[tuple[torch.FloatTensor]] = None844 attentions: Optional[tuple[torch.FloatTensor]] = None845 loc: Optional[torch.FloatTensor] = None846 scale: Optional[torch.FloatTensor] = None847 848 849@dataclass850@auto_docstring(851 custom_intro="""852 Output type of [`PatchTSTForClassification`].853 """854)855class PatchTSTForClassificationOutput(ModelOutput):856 r"""857 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):858 Total loss as the sum of the masked language modeling loss and the next sequence prediction859 (classification) loss.860 prediction_logits (`torch.FloatTensor` of shape `(batch_size, num_targets)`):861 Prediction scores of the PatchTST modeling head (scores before SoftMax).862 """863 864 loss: Optional[torch.FloatTensor] = None865 prediction_logits: Optional[torch.FloatTensor] = None866 hidden_states: Optional[tuple[torch.FloatTensor]] = None867 attentions: Optional[tuple[torch.FloatTensor]] = None868 869 870@dataclass871@auto_docstring(872 custom_intro="""873 Base class for time series model's predictions outputs that contains the sampled values from the chosen874 distribution.875 """876)877class SamplePatchTSTOutput(ModelOutput):878 r"""879 sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, num_targets)`):880 Sampled values from the chosen distribution.881 """882 883 sequences: Optional[torch.FloatTensor] = None884 885 886# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.nll887def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor:888 """889 Computes the negative log likelihood loss from input distribution with respect to target.890 """891 return -input.log_prob(target)892 893 894# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.weighted_average895def weighted_average(input_tensor: torch.Tensor, weights: Optional[torch.Tensor] = None, dim=None) -> torch.Tensor:896 """897 Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero,898 meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.899 900 Args:901 input_tensor (`torch.FloatTensor`):902 Input tensor, of which the average must be computed.903 weights (`torch.FloatTensor`, *optional*):904 Weights tensor, of the same shape as `input_tensor`.905 dim (`int`, *optional*):906 The dim along which to average `input_tensor`.907 908 Returns:909 `torch.FloatTensor`: The tensor with values averaged along the specified `dim`.910 """911 if weights is not None:912 weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))913 sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)914 return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights915 else:916 return input_tensor.mean(dim=dim)917 918 919# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesStdScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST920class PatchTSTStdScaler(nn.Module):921 """922 Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by923 subtracting from the mean and dividing by the standard deviation.924 """925 926 def __init__(self, config: PatchTSTConfig):927 super().__init__()928 self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1929 self.keepdim = config.keepdim if hasattr(config, "keepdim") else True930 self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5931 932 def forward(933 self, data: torch.Tensor, observed_indicator: torch.Tensor934 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:935 """936 Parameters:937 data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):938 input for Batch norm calculation939 observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):940 Calculating the scale on the observed indicator.941 Returns:942 tuple of `torch.Tensor` of shapes943 (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,944 `(batch_size, 1, num_input_channels)`)945 """946 denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)947 denominator = denominator.clamp_min(1.0)948 loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator949 950 variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator951 scale = torch.sqrt(variance + self.minimum_scale)952 return (data - loc) / scale, loc, scale953 954 955# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesMeanScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST956class PatchTSTMeanScaler(nn.Module):957 """958 Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data959 accordingly.960 """961 962 def __init__(self, config: PatchTSTConfig):963 super().__init__()964 self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1965 self.keepdim = config.keepdim if hasattr(config, "keepdim") else True966 self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10967 self.default_scale = config.default_scale if hasattr(config, "default_scale") else None968 969 def forward(970 self, data: torch.Tensor, observed_indicator: torch.Tensor971 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:972 """973 Parameters:974 data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):975 input for Batch norm calculation976 observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):977 Calculating the scale on the observed indicator.978 Returns:979 tuple of `torch.Tensor` of shapes980 (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,981 `(batch_size, 1, num_input_channels)`)982 """983 ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)984 num_observed = observed_indicator.sum(self.dim, keepdim=True)985 986 scale = ts_sum / torch.clamp(num_observed, min=1)987 988 # If `default_scale` is provided, we use it, otherwise we use the scale989 # of the batch.990 if self.default_scale is None:991 batch_sum = ts_sum.sum(dim=0)992 batch_observations = torch.clamp(num_observed.sum(0), min=1)993 default_scale = torch.squeeze(batch_sum / batch_observations)994 else:995 default_scale = self.default_scale * torch.ones_like(scale)996 997 # apply default scale where there are no observations998 scale = torch.where(num_observed > 0, scale, default_scale)999 1000 # ensure the scale is at least `self.minimum_scale`1001 scale = torch.clamp(scale, min=self.minimum_scale)1002 scaled_data = data / scale1003 1004 if not self.keepdim:1005 scale = scale.squeeze(dim=self.dim)1006 1007 return scaled_data, torch.zeros_like(scale), scale1008 1009 1010# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesNOPScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST1011class PatchTSTNOPScaler(nn.Module):1012 """1013 Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data.1014 """1015 1016 def __init__(self, config: PatchTSTConfig):1017 super().__init__()1018 self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 11019 self.keepdim = config.keepdim if hasattr(config, "keepdim") else True1020 1021 def forward(1022 self, data: torch.Tensor, observed_indicator: Optional[torch.Tensor] = None1023 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:1024 """1025 Parameters:1026 data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):1027 input for Batch norm calculation1028 Returns:1029 tuple of `torch.Tensor` of shapes1030 (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,1031 `(batch_size, 1, num_input_channels)`)1032 """1033 scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)1034 loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)1035 return data, loc, scale1036 1037 1038class PatchTSTScaler(nn.Module):1039 def __init__(self, config: PatchTSTConfig):1040 super().__init__()1041 if config.scaling == "mean" or config.scaling is True:1042 self.scaler = PatchTSTMeanScaler(config)1043 elif config.scaling == "std":1044 self.scaler = PatchTSTStdScaler(config)1045 else:1046 self.scaler = PatchTSTNOPScaler(config)1047 1048 def forward(1049 self, data: torch.Tensor, observed_indicator: torch.Tensor1050 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:1051 """1052 Parameters:1053 data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):1054 Input for scaler calculation1055 observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):1056 Calculating the scale on the observed indicator.1057 Returns:1058 tuple of `torch.Tensor` of shapes1059 (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,1060 `(batch_size, 1, um_input_channels)`)1061 """1062 data, loc, scale = self.scaler(data, observed_indicator)1063 return data, loc, scale1064 1065 1066@auto_docstring1067class PatchTSTModel(PatchTSTPreTrainedModel):1068 def __init__(self, config: PatchTSTConfig):1069 super().__init__(config)1070 1071 self.scaler = PatchTSTScaler(config)1072 self.patchifier = PatchTSTPatchify(config)1073 self.do_mask_input = config.do_mask_input1074 # get num_patches information from PatchTSTPatchify1075 num_patches = self.patchifier.num_patches1076 1077 if self.do_mask_input:1078 self.masking = PatchTSTMasking(config)1079 else:1080 self.masking = nn.Identity()1081 self.encoder = PatchTSTEncoder(config, num_patches=num_patches)1082 1083 # Initialize weights and apply final processing1084 self.post_init()1085 1086 def forward(1087 self,1088 past_values: torch.Tensor,1089 past_observed_mask: Optional[torch.Tensor] = None,1090 future_values: Optional[torch.Tensor] = None,1091 output_hidden_states: Optional[bool] = None,1092 output_attentions: Optional[bool] = None,1093 return_dict: Optional[bool] = None,1094 ) -> Union[tuple, PatchTSTModelOutput]:1095 r"""1096 Parameters:1097 past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):1098 Input sequence to the model1099 past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):1100 Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected1101 in `[0, 1]`:1102 1103 - 1 for values that are **observed**,1104 - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).1105 future_values (`torch.BoolTensor` of shape `(batch_size, prediction_length, num_input_channels)`, *optional*):1106 Future target values associated with the `past_values`1107 output_hidden_states (`bool`, *optional*):1108 Whether or not to return the hidden states of all layers1109 output_attentions (`bool`, *optional*):1110 Whether or not to return the output attention of all layers1111 return_dict (`bool`, *optional*):1112 Whether or not to return a `ModelOutput` instead of a plain tuple.1113 1114 Returns:1115 `PatchTSTModelOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False)1116 1117 Examples:1118 1119 ```python1120 >>> from huggingface_hub import hf_hub_download1121 >>> import torch1122 >>> from transformers import PatchTSTModel1123 1124 >>> file = hf_hub_download(1125 ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset"1126 ... )1127 >>> batch = torch.load(file)1128 1129 >>> model = PatchTSTModel.from_pretrained("namctin/patchtst_etth1_pretrain")1130 1131 >>> # during training, one provides both past and future values1132 >>> outputs = model(1133 ... past_values=batch["past_values"],1134 ... future_values=batch["future_values"],1135 ... )1136 1137 >>> last_hidden_state = outputs.last_hidden_state1138 ```"""1139 1140 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1141 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1142 output_hidden_states = (1143 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1144 )1145 1146 if past_observed_mask is None:1147 past_observed_mask = torch.ones_like(past_values)1148 1149 # x: tensor [bs x sequence_length x num_input_channels]1150 scaled_past_values, loc, scale = self.scaler(past_values, past_observed_mask)1151 1152 # patched_values: [bs x num_input_channels x num_patches x patch_length] for pretrain1153 patched_values = self.patchifier(scaled_past_values)1154 if self.do_mask_input:1155 masked_values, mask = self.masking(patched_values)1156 else:1157 masked_values, mask = self.masking(patched_values), None1158 1159 encoder_output = self.encoder(1160 patch_input=masked_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions1161 )1162 1163 if not return_dict:1164 outputs = (encoder_output.last_hidden_state, encoder_output.hidden_states, encoder_output.attentions)1165 outputs = outputs + (mask, loc, scale, patched_values)1166 return tuple(v for v in outputs if v is not None)1167 1168 return PatchTSTModelOutput(1169 last_hidden_state=encoder_output.last_hidden_state,1170 hidden_states=encoder_output.hidden_states,1171 attentions=encoder_output.attentions,1172 mask=mask,1173 loc=loc,1174 scale=scale,1175 patch_input=patched_values,1176 )1177 1178 1179class PatchTSTMaskPretrainHead(nn.Module):1180 """1181 Pretraining head for mask modelling1182 """1183 1184 def __init__(self, config: PatchTSTConfig):1185 super().__init__()1186 self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()1187 self.linear = nn.Linear(config.d_model, config.patch_length)1188 self.use_cls_token = config.use_cls_token1189 1190 def forward(self, embedding: torch.Tensor) -> torch.Tensor:1191 """1192 Parameters:1193 embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or1194 `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):1195 Embedding from the model1196 Returns:1197 `torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or1198 `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True1199 1200 """