added predictor

pts/model/__init__.py

@@ -1,5 +1,5 @@
 from .estimator import Estimator, PTSEstimator
 from .forecast import Forecast, SampleForecast, QuantileForecast, DistributionForecast
-from .predictor import Predictor
+from .predictor import Predictor, PTSPredictor
 from .quantile import Quantile
 from .utils import get_module_forward_input_names, copy_parameters

pts/model/deepar/deepar_estimator.py

@@ -22,7 +22,7 @@ from pts.feature import (
     InstanceSplitter,
 )
 from pts.dataset import FieldName, ExpectedNumInstanceSampler
-from pts.model import PTSEstimator, Predictor
+from pts.model import PTSEstimator, Predictor, PTSPredictor
 from pts.modules import DistributionOutput, StudentTOutput
 
 from .deepar_network import DeepARTrainingNetwork, DeepARPredictionNetwork
@@ -170,7 +170,7 @@ class DeepAREstimator(PTSEstimator):
             dtype=self.dtype).to(device)
 
     def create_predictor(
-        self, transformation: Transformation, trained_network: nn.Module
+        self, transformation: Transformation, trained_network: nn.Module, device: torch.device
     ) -> Predictor:
         prediction_network = DeepARPredictionNetwork(
             num_parallel_samples=self.num_parallel_samples,
@@ -187,8 +187,16 @@ class DeepAREstimator(PTSEstimator):
             embedding_dimension=self.embedding_dimension,
             lags_seq=self.lags_seq,
             scaling=self.scaling,
-            dtype=self.dtype)
+            dtype=self.dtype).to(device)
 
         copy_parameters(trained_network, prediction_network)
 
-        return RepresentableBlockPredictor()
+        return PTSPredictor(
+            input_transform=transformation,
+            prediction_net=prediction_network,
+            batch_size=self.trainer.batch_size,
+            freq=self.freq,
+            prediction_length=self.prediction_length,
+            device=device,
+            dtype=self.dtype
+        )

pts/model/estimator.py

@@ -81,7 +81,7 @@ class PTSEstimator(Estimator):
 
     @abstractmethod
     def create_predictor(
-        self, transformation: Transformation, trained_network: nn.Module) -> Predictor:
+        self, transformation: Transformation, trained_network: nn.Module, device: torch.device) -> Predictor:
         """
         Create and return a predictor object.
 
@@ -118,7 +118,7 @@ class PTSEstimator(Estimator):
         return TrainOutput(
             transformation=transformation,
             trained_net=trained_net,
-            predictor=self.create_predictor(transformation, trained_net),
+            predictor=self.create_predictor(transformation, trained_net, self.trainer.device),
         )
 
     def train(self, training_data: Dataset) -> Predictor:

pts/model/forecast_generator.py

@@ -5,9 +5,9 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from pts.dataset import InferenceDataLoader, DataEntry
-from pts.model import Forecast, DistributionForecast, QuantileForecast, SampleForecast
+from pts.dataset import InferenceDataLoader, DataEntry, FieldName
 from pts.modules import DistributionOutput
+from .forecast import Forecast, DistributionForecast, QuantileForecast, SampleForecast
 
 OutputTransform = Callable[[DataEntry, np.ndarray], np.ndarray]
 
@@ -90,7 +90,7 @@ class QuantileForecastGenerator(ForecastGenerator):
         self.quantiles = quantiles
 
     def __call__(self, inference_data_loader: InferenceDataLoader,
-                 prediction_net: BlockType, input_names: List[str], freq: str,
+                 prediction_net: nn.Module, input_names: List[str], freq: str,
                  output_transform: Optional[OutputTransform],
                  num_samples: Optional[int], **kwargs) -> Iterator[Forecast]:
         for batch in inference_data_loader:
@@ -115,7 +115,7 @@ class QuantileForecastGenerator(ForecastGenerator):
 
 class SampleForecastGenerator(ForecastGenerator):
     def __call__(self, inference_data_loader: InferenceDataLoader,
-                 prediction_net: BlockType, input_names: List[str], freq: str,
+                 prediction_net: nn.Module, input_names: List[str], freq: str,
                  output_transform: Optional[OutputTransform],
                  num_samples: Optional[int], **kwargs) -> Iterator[Forecast]:
         for batch in inference_data_loader:

pts/model/predictor.py

@@ -1,12 +1,18 @@
 from abc import ABC, abstractmethod
-from typing import Iterator
+from typing import Iterator, Callable, Optional
 
+import numpy as np
 import torch
 import torch.nn as nn
 
-from pts.dataset import Dataset
+from pts.dataset import Dataset, DataEntry, InferenceDataLoader
+from pts.feature import Transformation
 
 from .forecast import Forecast
+from .forecast_generator import ForecastGenerator, SampleForecastGenerator
+from .utils import get_module_forward_input_names
+
+OutputTransform = Callable[[DataEntry, np.ndarray], np.ndarray]
 
 
 class Predictor(ABC):
@@ -18,7 +24,46 @@ class Predictor(ABC):
     def predict(self, dataset: Dataset, **kwargs) -> Iterator[Forecast]:
         pass
 
-class PTSPredictor(Predictor):
-    BlockType = nn.Module
 
-    
+class PTSPredictor(Predictor):
+    def __init__(
+        self,
+        prediction_net: nn.Module,
+        batch_size: int,
+        prediction_length: int,
+        freq: str,
+        device: torch.device,
+        input_transform: Transformation,
+        forecast_generator: ForecastGenerator = SampleForecastGenerator(),
+        output_transform: Optional[OutputTransform] = None,
+        dtype: np.dtype = np.float32,
+    ) -> None:
+        super().__init__(prediction_length, freq)
+
+        self.input_names = get_module_forward_input_names(prediction_net)
+        self.prediction_net = prediction_net
+        self.batch_size = batch_size
+        self.input_transform = input_transform
+        self.forecast_generator = forecast_generator
+        self.output_transform = output_transform
+        self.device = device
+        self.dtype = dtype
+
+    def predict(self,
+                dataset: Dataset,
+                num_samples: Optional[int] = None) -> Iterator[Forecast]:
+        inference_data_loader = InferenceDataLoader(
+            dataset,
+            self.input_transform,
+            self.batch_size,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        yield from self.forecast_generator(
+            inference_data_loader=inference_data_loader,
+            prediction_net=self.prediction_net,
+            input_names=self.input_names,
+            freq=self.freq,
+            output_transform=self.output_transform,
+            num_samples=num_samples,
+        )