Breaking: change num_output with quantiles in TFT (mxnet) (#2879)

src/gluonts/mx/model/tft/_estimator.py

@@ -66,6 +66,53 @@ def _default_feat_args(dims_or_cardinalities: List[int]):
 
 
 class TemporalFusionTransformerEstimator(GluonEstimator):
+ """
+ Parameters
+ ----------
+ freq
+ Frequency of the data to train on and predict.
+ prediction_length
+ Length of the prediction horizon.
+ context_length
+ Number of previous time series values provided as input to the encoder.
+ (default: None).
+ trainer
+ Trainer object to be used (default: Trainer())
+ hidden_dim
+ Size of the LSTM & transformer hidden states.
+ variable_dim
+ Size of the feature embeddings.
+ num_heads
+ Number of attention heads in self-attention layer in the decoder.
+ quantiles
+ List of quantiles that the model will learn to predict.
+ Defaults to [0.1, 0.5, 0.9]
+ num_instances_per_series
+ Number of samples to generate for each time series when training.
+ dropout_rate
+ Dropout regularization parameter (default: 0.1).
+ time_features
+ List of time features, from :py:mod:`gluonts.time_feature`, to use as
+ dynamic real features in addition to the provided data (default: None,
+ in which case these are automatically determined based on freq).
+ static_cardinalities
+ Cardinalities of the categorical static features.
+ dynamic_cardinalities
+ Cardinalities of the categorical dynamic features that are known in the future.
+ static_feature_dims
+ Sizes of the real-valued static features.
+ dynamic_dims
+ Sizes of the real-valued dynamic features that are known in the future.
+ past_dynamic_features
+ List of names of the real-valued dynamic features that are only known in the past.
+ train_sampler
+ Controls the sampling of windows during training.
+ validation_sampler
+ Controls the sampling of windows during validation.
+ batch_size
+ The size of the batches to be used training and prediction.
+ """
+
  @validated()
  def __init__(
  self,
@@ -76,7 +123,7 @@ def __init__(
  hidden_dim: int = 32,
  variable_dim: Optional[int] = None,
  num_heads: int = 4,
- num_outputs: int = 3,
+ quantiles: List[float] = [0.1, 0.5, 0.9],
  num_instance_per_series: int = 100,
  dropout_rate: float = 0.1,
  time_features: List[TimeFeature] = [],
@@ -104,7 +151,7 @@ def __init__(
  self.hidden_dim = hidden_dim
  self.variable_dim = variable_dim or hidden_dim
  self.num_heads = num_heads
- self.num_outputs = num_outputs
+ self.quantiles = quantiles
  self.num_instance_per_series = num_instance_per_series
 
  if not time_features:
@@ -372,7 +419,7 @@ def create_training_network(
  d_var=self.variable_dim,
  d_hidden=self.hidden_dim,
  n_head=self.num_heads,
- n_output=self.num_outputs,
+ quantiles=self.quantiles,
  d_past_feat_dynamic_real=_default_feat_args(
  list(self.past_dynamic_feature_dims.values())
  ),
@@ -406,7 +453,7 @@ def create_predictor(
  d_var=self.variable_dim,
  d_hidden=self.hidden_dim,
  n_head=self.num_heads,
- n_output=self.num_outputs,
+ quantiles=self.quantiles,
  d_past_feat_dynamic_real=_default_feat_args(
  list(self.past_dynamic_feature_dims.values())
  ),

src/gluonts/mx/model/tft/_network.py

@@ -151,7 +151,7 @@ def __init__(
  d_var: int,
  d_hidden: int,
  n_head: int,
- n_output: int,
+ quantiles: List[float],
  d_past_feat_dynamic_real: List[int],
  c_past_feat_dynamic_cat: List[int],
  d_feat_dynamic_real: List[int],
@@ -167,11 +167,7 @@ def __init__(
  self.d_var = d_var
  self.d_hidden = d_hidden
  self.n_head = n_head
- self.n_output = n_output
- self.quantiles = sum(
- ([i / 10, 1.0 - i / 10] for i in range(1, (n_output + 1) // 2)),
- [0.5],
- )
+ self.quantiles = quantiles
  self.normalize_eps = 1e-5
 
  self.d_past_feat_dynamic_real = d_past_feat_dynamic_real