diff --git a/pts/model/deepar/deepar_estimator.py b/pts/model/deepar/deepar_estimator.py index 5266aad..d2975b4 100644 --- a/pts/model/deepar/deepar_estimator.py +++ b/pts/model/deepar/deepar_estimator.py @@ -32,6 +32,7 @@ class DeepAREstimator(PTSEstimator): self, freq: str, prediction_length: int, + input_size: int, trainer: Trainer = Trainer(), context_length: Optional[int] = None, num_layers: int = 2, @@ -58,6 +59,7 @@ class DeepAREstimator(PTSEstimator): self.prediction_length = prediction_length self.distr_output = distr_output self.distr_output.dtype = dtype + self.input_size = input_size self.num_layers = num_layers self.num_cells = num_cells self.cell_type = cell_type @@ -65,9 +67,7 @@ class DeepAREstimator(PTSEstimator): self.use_feat_dynamic_real = use_feat_dynamic_real self.use_feat_static_cat = use_feat_static_cat self.use_feat_static_real = use_feat_static_real - self.cardinality = cardinality if cardinality and use_feat_static_cat else [ - 1 - ] + self.cardinality = cardinality if cardinality and use_feat_static_cat else [1] self.embedding_dimension = ( embedding_dimension if embedding_dimension is not None else [min(50, (cat + 1) // 2) for cat in self.cardinality]) @@ -153,6 +153,7 @@ class DeepAREstimator(PTSEstimator): def create_training_network(self, device: torch.device) -> DeepARTrainingNetwork: return DeepARTrainingNetwork( + input_size=self.input_size, num_layers=self.num_layers, num_cells=self.num_cells, cell_type=self.cell_type, diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py index c5955fd..dabf715 100644 --- a/pts/model/deepar/deepar_network.py +++ b/pts/model/deepar/deepar_network.py @@ -17,6 +17,7 @@ def prod(xs): class DeepARNetwork(nn.Module): def __init__( self, + input_size: int, num_layers: int, num_cells: int, cell_type: str, @@ -32,6 +33,7 @@ class DeepARNetwork(nn.Module): dtype: np.dtype = np.float32, ) -> None: super().__init__() + self.input_size = input_size self.num_layers = num_layers self.num_cells = num_cells self.cell_type = cell_type @@ -49,7 +51,7 @@ class DeepARNetwork(nn.Module): self.distr_output = distr_output rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type] - self.rnn = rnn(input_size=48, + self.rnn = rnn(input_size=input_size, hidden_size=num_cells, num_layers=num_layers, dropout=dropout_rate, @@ -72,8 +74,7 @@ class DeepARNetwork(nn.Module): sequence: torch.Tensor, sequence_length: int, indices: List[int], - subsequences_length: int = 1, - ) -> torch.Tensor: + subsequences_length: int = 1) -> torch.Tensor: """ Returns lagged subsequences of a given sequence. Parameters diff --git a/test/model/deepar/test_auxillary_outputs.py b/test/model/deepar/test_auxillary_outputs.py index 2cdfeea..4d55a5b 100644 --- a/test/model/deepar/test_auxillary_outputs.py +++ b/test/model/deepar/test_auxillary_outputs.py @@ -35,6 +35,7 @@ def test_distribution(): estimator = DeepAREstimator( freq=freq, prediction_length=prediction_length, + input_size=48, trainer=Trainer(epochs=1, num_batches_per_epoch=1), distr_output=StudentTOutput(), )