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,
control_output: DistributionOutput,
dropout_rate: float,
cardinality: List[int],
embedding_dimension: List[int],
lags_seq: List[int],
scaling: bool = True,
dtype: np.dtype = np.float32,
) -> None:
super().__init__()
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.scaling = scaling
self.dtype = dtype
self.lags_seq = lags_seq
self.distr_output = distr_output
self.control_output = control_output
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_distr_args = distr_output.get_args_proj(num_cells + 1)
self.proj_control_args = control_output.get_args_proj(num_cells)
self.embedder = FeatureEmbedder(cardinalities=cardinality,
embedding_dims=embedding_dimension)
if scaling:
self.scaler = MeanScaler(keepdim=True)
self.control_scaler = NOPScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
self.control_scaler = NOPScaler(keepdim=True)
def test_feature_embedder(config):
out_shape = config["shape"][:-1] + (sum(config["kwargs"]["embedding_dims"]),)
embed_feature = FeatureEmbedder(**config["kwargs"])
for embed in embed_feature._FeatureEmbedder__embedders:
nn.init.constant_(embed.weight, 1.0)
def test_parameters_length():
exp_params_len = len([p for p in embed_feature.parameters()])
act_params_len = len(config["kwargs"]["embedding_dims"])
assert exp_params_len == act_params_len
def test_forward_pass():
act_output = embed_feature(torch.ones(config["shape"]).to(torch.long))
exp_output = torch.ones(out_shape)
assert act_output.shape == exp_output.shape
assert torch.abs(torch.sum(act_output - exp_output)) < 1e-20
test_parameters_length()
test_forward_pass()
def __init__(
self,
input_size: int,
decoder_size: int,
num_layers: int,
num_cells: int,
cell_type: str,
short_cycle: int,
history_length: int,
context_length: int, # should be equal to the length of long period
prediction_length: int, # should be integer multiples of small period
distr_output: DistributionOutput,
dropout_rate: float,
cardinality: List[int],
embedding_dimension: List[int],
lags_seq: List[int],
scaling: bool = True,
dtype: np.dtype = np.float32,
) -> None:
super().__init__()
self.num_layers = num_layers
self.num_cells = num_cells
self.cell_type = cell_type
self.short_cycle = short_cycle
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.scaling = scaling
self.dtype = dtype
# self.lags_seq = lags_seq + [0] # 0 is the current value
self.lags_seq = [l - 1 for l in lags_seq] # to distinguish with deepAR
self.distr_output = distr_output
rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type]
self.long_cell = nn.GRUCell(input_size=input_size, hidden_size=num_cells)
self.encoder = nn.ModuleList(
rnn(
input_size=num_cells,
hidden_size=num_cells,
num_layers=num_layers,
dropout=dropout_rate,
batch_first=True,
)
for _ in range(self.context_length // self.short_cycle)
)
# self.encoder = rnn(
# input_size=num_cells,
# hidden_size=num_cells,
# num_layers=num_layers,
# dropout=dropout_rate,
# batch_first=True,
# )
self.decoder = rnn(
input_size=decoder_size,
hidden_size=num_cells,
num_layers=num_layers,
dropout=dropout_rate,
batch_first=True,
)
self.decoder_dnn = nn.ModuleList(nn.Sequential(
nn.Linear(decoder_size, 1),
# nn.ReLU()
) for _ in range(self.prediction_length))
self.out = nn.Linear(num_cells, 1)
self.criterion = nn.MSELoss() # l2, l1
self.target_shape = distr_output.event_shape
self.proj_distr_args = distr_output.get_args_proj(num_cells)
self.embedder = FeatureEmbedder(
cardinalities=cardinality, embedding_dims=embedding_dimension
)
if scaling:
self.scaler = MeanScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
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 __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)
)
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()