mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-11 22:00:32 +08:00
Dr. Kashif Rasul
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class FeatureEmbedder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
cardinalities: List[int],
|
|
embedding_dims: List[int],
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.__num_features = len(cardinalities)
|
|
|
|
def create_embedding(c: int, d: int) -> nn.Embedding:
|
|
embedding = nn.Embedding(c, d)
|
|
return embedding
|
|
|
|
self.__embedders = nn.ModuleList([
|
|
create_embedding(c, d)
|
|
for c, d in zip(cardinalities, embedding_dims)
|
|
])
|
|
|
|
def forward(self, features: torch.Tensor) -> torch.Tensor:
|
|
if self.__num_features > 1:
|
|
# we slice the last dimension, giving an array of length
|
|
# self.__num_features with shape (N,T) or (N)
|
|
cat_feature_slices = torch.chunk(features,
|
|
self.__num_features,
|
|
dim=-1)
|
|
else:
|
|
cat_feature_slices = [features]
|
|
|
|
return torch.cat([
|
|
embed(cat_feature_slice.squeeze(-1)) for embed, cat_feature_slice
|
|
in zip(self.__embedders, cat_feature_slices)
|
|
], dim=-1)
|
|
|
|
|
|
class FeatureAssembler(nn.Module):
|
|
def __init__(T: int,
|
|
use_static_cat: bool = False,
|
|
use_static_real: bool = False,
|
|
use_dynamic_cat: bool = False,
|
|
use_dynamic_real: bool = False,
|
|
embed_static: Optional[FeatureEmbedder] = None,
|
|
embed_dynamic: Optional[FeatureEmbedder] = None,
|
|
dtype: torch.dtype = torch.float32) -> None:
|
|
super().__init__()
|
|
|
|
self.T = T
|
|
self.use_static_cat = use_static_cat
|
|
self.use_static_real = use_static_real
|
|
self.use_dynamic_cat = use_dynamic_cat
|
|
self.use_dynamic_real = use_dynamic_real
|
|
|
|
self.embed_static: Callable[[torch.Tensor], torch.
|
|
Tensor] = embed_static or (lambda x: x)
|
|
self.embed_dynamic: Callable[[torch.Tensor], torch.
|
|
Tensor] = embed_dynamic or (lambda x: x)
|
|
|
|
def forward(
|
|
self,
|
|
feat_static_cat: torch.Tensor,
|
|
feat_static_real: torch.Tensor,
|
|
feat_dynamic_cat: torch.Tensor,
|
|
feat_dynamic_real: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
processed_features = [
|
|
self.process_static_cat(feat_static_cat),
|
|
self.process_static_real(feat_static_real),
|
|
self.process_dynamic_cat(feat_dynamic_cat),
|
|
self.process_dynamic_real(feat_dynamic_real),
|
|
]
|
|
|
|
return torch.cat(processed_features, dim=-1)
|
|
|
|
def process_static_cat(self, feature: torch.Tensor) -> torch.Tensor:
|
|
feature = self.embed_static(feature.to(self.dtype))
|
|
return feature.unsqueeze(1).expand(-1, self.T, -1)
|
|
|
|
def process_dynamic_cat(self, feature: torch.Tensor) -> torch.Tensor:
|
|
return self.embed_dynamic(feature.to(self.dtype))
|
|
|
|
def process_static_real(self, feature: torch.Tensor) -> torch.Tensor:
|
|
return feature.unsqueeze(1).expand(-1, self.T, -1)
|
|
|
|
def process_dynamic_real(self, feature: torch.Tensor) -> torch.Tensor:
|
|
return feature
|