mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-25 13:30:12 +08:00
25 lines
503 B
Python
25 lines
503 B
Python
from abc import ABC, abstractmethod
|
|
from typing import Iterator
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from pts.dataset import Dataset
|
|
|
|
from .forecast import Forecast
|
|
|
|
|
|
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):
|
|
BlockType = nn.Module
|
|
|
|
|