mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-10 02:57:15 +08:00
Kashif Rasul
74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Iterator, Callable, Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from pts.dataset import Dataset, DataEntry, InferenceDataLoader
|
|
from pts.transform 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):
|
|
def __init__(self, prediction_length: int, freq: str) -> None:
|
|
self.prediction_length = prediction_length
|
|
self.freq = freq
|
|
|
|
@abstractmethod
|
|
def predict(self, dataset: Dataset, **kwargs) -> Iterator[Forecast]:
|
|
pass
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
self.prediction_net.eval()
|
|
|
|
with torch.no_grad():
|
|
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,
|
|
)
|