updated transformer model

pts/model/deepar/deepar_estimator.py

@@ -1,4 +1,3 @@
-from pts.model.utils import get_module_forward_input_names
 from typing import List, Optional
 
 import numpy as np
@@ -28,7 +27,7 @@ from gluonts.torch.support.util import copy_parameters
 from gluonts.torch.model.predictor import PyTorchPredictor
 from gluonts.torch.modules.distribution_output import DistributionOutput
 from gluonts.model.predictor import Predictor
-
+from pts.model.utils import get_module_forward_input_names
 from pts import Trainer
 from pts.model import PyTorchEstimator
 from pts.modules import StudentTOutput

pts/model/transformer/transformer_estimator.py

@@ -4,16 +4,13 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from pts import Trainer
-from pts.dataset import FieldName
-from pts.feature import (
-    TimeFeature,
-    fourier_time_features_from_frequency_str,
-    get_fourier_lags_for_frequency,
-)
-from pts.model import PyTorchEstimator, Predictor, PyTorchPredictor, copy_parameters
-from pts.modules import DistributionOutput, StudentTOutput
-from pts.transform import (
+from gluonts.dataset.field_names import FieldName
+from gluonts.time_feature import TimeFeature
+from gluonts.torch.modules.distribution_output import DistributionOutput
+from gluonts.torch.support.util import copy_parameters
+from gluonts.torch.model.predictor import PyTorchPredictor
+from gluonts.model.predictor import Predictor
+from gluonts.transform  import (
     Transformation,
     Chain,
     InstanceSplitter,
@@ -26,6 +23,14 @@ from pts.transform import (
     VstackFeatures,
     SetField,
 )
+from pts import Trainer
+from pts.model.utils import get_module_forward_input_names
+from pts.feature import (
+    fourier_time_features_from_frequency,
+    lags_for_fourier_time_features_from_frequency,
+)
+from pts.model import PyTorchEstimator
+from pts.modules import StudentTOutput
 from .transformer_network import (
     TransformerTrainingNetwork,
     TransformerPredictionNetwork,
@@ -77,12 +82,12 @@ class TransformerEstimator(PyTorchEstimator):
         self.lags_seq = (
             lags_seq
             if lags_seq is not None
-            else get_fourier_lags_for_frequency(freq_str=freq)
+            else lags_for_fourier_time_features_from_frequency(freq_str=freq)
         )
         self.time_features = (
             time_features
             if time_features is not None
-            else fourier_time_features_from_frequency_str(self.freq)
+            else fourier_time_features_from_frequency(self.freq)
         )
         self.history_length = self.context_length + max(self.lags_seq)
         self.scaling = scaling
@@ -198,7 +203,7 @@ class TransformerEstimator(PyTorchEstimator):
     def create_predictor(
         self,
         transformation: Transformation,
-        trained_network: nn.Module,
+        trained_network:TransformerTrainingNetwork,
         device: torch.device,
     ) -> Predictor:
 
@@ -223,13 +228,14 @@ class TransformerEstimator(PyTorchEstimator):
         ).to(device)
 
         copy_parameters(trained_network, prediction_network)
+        input_names = get_module_forward_input_names(prediction_network)
 
         return PyTorchPredictor(
             input_transform=transformation,
+            input_names=input_names,
             prediction_net=prediction_network,
             batch_size=self.trainer.batch_size,
             freq=self.freq,
             prediction_length=self.prediction_length,
             device=device,
-            output_transform=None,
         )

pts/model/transformer/transformer_network.py

@@ -4,7 +4,8 @@ import torch
 import torch.nn as nn
 
 from gluonts.core.component import validated
-from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder
+from gluonts.torch.modules.distribution_output import DistributionOutput
+from pts.modules import MeanScaler, NOPScaler, FeatureEmbedder
 
 
 def prod(xs):