| import torch |
| import torch.nn as nn |
| import math |
|
|
|
|
| class AstromerPosEmbedding(nn.Module): |
| def __init__(self, d_model, trainable=False, base=1000, **kwargs): |
| super(AstromerPosEmbedding, self).__init__() |
|
|
| self.d_model = d_model |
|
|
| initial_div_term = torch.exp( |
| torch.arange(0.0, d_model).float() * -(math.log(base) / d_model) |
| ) |
|
|
| if trainable: |
| self.w = nn.Parameter(initial_div_term) |
| else: |
| self.register_buffer("w", initial_div_term) |
|
|
| def combine_input_embeddings(self, m_emb, t_emb): |
| x = m_emb + t_emb |
| return x |
|
|
| def forward(self, emb_x, t): |
| """ |
| t: Tensor de series de tiempo con dimensiones [batch_size, seq_len, 1] |
| """ |
| w = self.w.unsqueeze(0).unsqueeze(1) |
| pe = t * w |
|
|
| |
| pe_sin = torch.sin(pe[:, :, 0::2]) |
| pe_cos = torch.cos(pe[:, :, 1::2]) |
|
|
| pe = torch.cat([pe_sin.unsqueeze(-1), pe_cos.unsqueeze(-1)], dim=-1).reshape( |
| pe.shape |
| ) |
|
|
| if emb_x == None: |
| return pe |
|
|
| return self.combine_input_embeddings(emb_x, pe) |
|
|
