def __init__(
self,
input_size: int,
num_layers: int,
num_cells: int,
cell_type: str,
history_length: int,
context_length: int,
prediction_length: int,
distr_output: DistributionOutput,
dropout_rate: float,
lags_seq: List[int],
target_dim: int,
cardinality: List[int],
embedding_dimension: List[int],
scaling: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.num_layers = num_layers
self.num_cells = num_cells
self.cell_type = cell_type
self.history_length = history_length
self.context_length = context_length
self.prediction_length = prediction_length
self.dropout_rate = dropout_rate
self.cardinality = cardinality
self.embedding_dimension = embedding_dimension
self.num_cat = len(cardinality)
self.target_dim = target_dim
self.scaling = scaling
self.target_dim_sample = target_dim
assert len(
set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!"
lags_seq.sort()
self.lags_seq = lags_seq
self.distr_output = distr_output
self.target_dim = target_dim
rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type]
self.rnn = rnn(
input_size=input_size,
hidden_size=num_cells,
num_layers=num_layers,
dropout=dropout_rate,
batch_first=True,
)
self.target_shape = distr_output.event_shape
self.proj_dist_args = distr_output.get_args_proj(num_cells)
self.embed = FeatureEmbedder(cardinalities=cardinality,
embedding_dims=embedding_dimension)
if scaling:
self.scaler = MeanScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
def test_feature_assembler(config):
# iterate over the power-set of all possible feature types, excluding the empty set
feature_types = {
"static_cat",
"static_real",
"dynamic_cat",
"dynamic_real",
}
feature_combs = chain.from_iterable(
combinations(feature_types, r) for r in range(1, len(feature_types) + 1)
)
# iterate over the power-set of all possible feature types, including the empty set
embedder_types = {"embed_static", "embed_dynamic"}
embedder_combs = chain.from_iterable(
combinations(embedder_types, r) for r in range(0, len(embedder_types) + 1)
)
for enabled_embedders in embedder_combs:
embed_static = (
FeatureEmbedder(**config["embed_static"])
if "embed_static" in enabled_embedders
else None
)
embed_dynamic = (
FeatureEmbedder(**config["embed_dynamic"])
if "embed_dynamic" in enabled_embedders
else None
)
for enabled_features in feature_combs:
assemble_feature = FeatureAssembler(
T=config["T"], embed_static=embed_static, embed_dynamic=embed_dynamic,
)
# assemble_feature.collect_params().initialize(mx.initializer.One())
def test_parameters_length():
exp_params_len = sum(
[
len(config[k]["embedding_dims"])
for k in ["embed_static", "embed_dynamic"]
if k in enabled_embedders
]
)
act_params_len = len([p for p in assemble_feature.parameters()])
assert exp_params_len == act_params_len
def test_forward_pass():
N, T = config["N"], config["T"]
inp_features = []
out_features = []
if "static_cat" not in enabled_features:
inp_features.append(torch.zeros((N, 1)))
out_features.append(torch.zeros((N, T, 1)))
elif embed_static: # and 'static_cat' in enabled_features
C = config["static_cat"]["C"]
inp_features.append(
torch.cat(
[
torch.randint(
0,
config["embed_static"]["cardinalities"][c],
(N, 1),
)
for c in range(C)
],
dim=1,
)
)
out_features.append(
torch.ones(
(N, T, sum(config["embed_static"]["embedding_dims"]),)
)
)
else: # not embed_static and 'static_cat' in enabled_features
C = config["static_cat"]["C"]
inp_features.append(
torch.cat(
[
torch.randint(
0,
config["embed_static"]["cardinalities"][c],
(N, 1),
)
for c in range(C)
],
dim=1,
)
)
out_features.append(
inp_features[-1].unsqueeze(1).expand(-1, T, -1).float()
)
if "static_real" not in enabled_features:
inp_features.append(torch.zeros((N, 1)))
out_features.append(torch.zeros((N, T, 1)))
else:
C = config["static_real"]["C"]
static_real = torch.empty((N, C)).uniform_(0, 100)
inp_features.append(static_real)
out_features.append(static_real.unsqueeze(-2).expand(-1, T, -1))
if "dynamic_cat" not in enabled_features:
inp_features.append(torch.zeros((N, T, 1)))
out_features.append(torch.zeros((N, T, 1)))
elif embed_dynamic: # and 'static_cat' in enabled_features
C = config["dynamic_cat"]["C"]
inp_features.append(
torch.cat(
[
torch.randint(
0,
config["embed_dynamic"]["cardinalities"][c],
(N, T, 1),
)
for c in range(C)
],
dim=2,
)
)
out_features.append(
torch.ones(
(N, T, sum(config["embed_dynamic"]["embedding_dims"]),)
)
)
else: # not embed_dynamic and 'dynamic_cat' in enabled_features
C = config["dynamic_cat"]["C"]
inp_features.append(
torch.cat(
[
torch.randint(
0,
config["embed_dynamic"]["cardinalities"][c],
(N, T, 1),
)
for c in range(C)
],
dim=2,
)
)
out_features.append(inp_features[-1].float())
if "dynamic_real" not in enabled_features:
inp_features.append(torch.zeros((N, T, 1)))
out_features.append(torch.zeros((N, T, 1)))
else:
C = config["dynamic_real"]["C"]
dynamic_real = torch.empty((N, T, C)).uniform_(0, 100)
inp_features.append(dynamic_real)
out_features.append(dynamic_real)
act_output = assemble_feature(*inp_features)
exp_output = torch.cat(out_features, dim=2)
assert exp_output.shape == act_output.shape
assert torch.sum(exp_output - act_output) < 1e-20
test_parameters_length()
test_forward_pass()
def __init__(
self,
input_size: int,
d_model: int,
num_heads: int,
act_type: str,
dropout_rate: float,
dim_feedforward_scale: int,
num_encoder_layers: int,
num_decoder_layers: int,
history_length: int,
context_length: int,
prediction_length: int,
distr_output: DistributionOutput,
cardinality: List[int],
embedding_dimension: List[int],
lags_seq: List[int],
scaling: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.history_length = history_length
self.context_length = context_length
self.prediction_length = prediction_length
self.scaling = scaling
self.cardinality = cardinality
self.embedding_dimension = embedding_dimension
self.distr_output = distr_output
assert len(
set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!"
lags_seq.sort()
self.lags_seq = lags_seq
self.target_shape = distr_output.event_shape
# [B, T, input_size] -> [B, T, d_model]
self.encoder_input = nn.Linear(input_size, d_model)
self.decoder_input = nn.Linear(input_size, d_model)
# [B, T, d_model] where d_model / num_heads is int
self.transformer = nn.Transformer(
d_model=d_model,
nhead=num_heads,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward_scale * d_model,
dropout=dropout_rate,
activation=act_type,
)
self.proj_dist_args = distr_output.get_args_proj(d_model)
self.embedder = FeatureEmbedder(
cardinalities=cardinality,
embedding_dims=embedding_dimension,
)
if scaling:
self.scaler = MeanScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
# mask
self.register_buffer(
"tgt_mask",
self.transformer.generate_square_subsequent_mask(
prediction_length))