diff --git a/.gitignore b/.gitignore index 894a44c..d14244e 100644 --- a/.gitignore +++ b/.gitignore @@ -102,3 +102,5 @@ venv.bak/ # mypy .mypy_cache/ +.idea/ +runs/ diff --git a/pts/__init__.py b/pts/__init__.py index 623c653..63a6af7 100644 --- a/pts/__init__.py +++ b/pts/__init__.py @@ -1,2 +1,12 @@ +from pkgutil import extend_path +from pkg_resources import get_distribution, DistributionNotFound from .trainer import Trainer from .exception import assert_pts + + +__path__ = extend_path(__path__, __name__) # type: ignore + +try: + __version__ = get_distribution(__name__).version +except DistributionNotFound: + __version__ = "0.0.0-unknown" \ No newline at end of file diff --git a/pts/core/__init__.py b/pts/core/__init__.py new file mode 100644 index 0000000..9687f13 --- /dev/null +++ b/pts/core/__init__.py @@ -0,0 +1,9 @@ +# Relative imports +from ._base import fqname_for + +__all__ = ["fqname_for"] + +# fix Sphinx issues, see https://bit.ly/2K2eptM +for item in __all__: + if hasattr(item, "__module__"): + setattr(item, "__module__", __name__) \ No newline at end of file diff --git a/pts/core/_base.py b/pts/core/_base.py new file mode 100644 index 0000000..bca69cc --- /dev/null +++ b/pts/core/_base.py @@ -0,0 +1,15 @@ +def fqname_for(cls: type) -> str: + """ + Returns the fully qualified name of ``cls``. + + Parameters + ---------- + cls + The class we are interested in. + + Returns + ------- + str + The fully qualified name of ``cls``. + """ + return f"{cls.__module__}.{cls.__qualname__}" \ No newline at end of file diff --git a/pts/core/component.py b/pts/core/component.py new file mode 100644 index 0000000..6ad703e --- /dev/null +++ b/pts/core/component.py @@ -0,0 +1,161 @@ +import inspect +from pydantic import BaseConfig, BaseModel, create_model +from typing import Any +from collections import OrderedDict +from pts.core.serde import dump_code +import functools +import torch + + +class BaseValidatedInitializerModel(BaseModel): + """ + Base Pydantic model for components with :func:`validated` initializers. + + See Also + -------- + validated + Decorates an initializer methods with argument validation logic. + """ + + class Config(BaseConfig): + """ + `Config `_ for the + Pydantic model inherited by all :func:`validated` initializers. + + Allows the use of arbitrary type annotations in initializer parameters. + """ + + arbitrary_types_allowed = True + + +def validated(base_model=None): + """ + Decorates an ``__init__`` method with typed parameters with validation + and auto-conversion logic. + + >>> class ComplexNumber: + ... @validated() + ... def __init__(self, x: float = 0.0, y: float = 0.0) -> None: + ... self.x = x + ... self.y = y + + Classes with decorated initializers can be instantiated using arguments of + another type (e.g. an ``y`` argument of type ``str`` ). The decorator + handles the type conversion logic. + + >>> c = ComplexNumber(y='42') + >>> (c.x, c.y) + (0.0, 42.0) + + If the bound argument cannot be converted, the decorator throws an error. + + >>> c = ComplexNumber(y=None) + Traceback (most recent call last): + ... + pydantic.error_wrappers.ValidationError: 1 validation error for ComplexNumberModel + y + none is not an allowed value (type=type_error.none.not_allowed) + + Internally, the decorator delegates all validation and conversion logic to + `a Pydantic model `_, which can be + accessed through the ``Model`` attribute of the decorated initiazlier. + + >>> ComplexNumber.__init__.Model + + + The Pydantic model is synthesized automatically from on the parameter + names and types of the decorated initializer. In the ``ComplexNumber`` + example, the synthesized Pydantic model corresponds to the following + definition. + + >>> class ComplexNumberModel(BaseValidatedInitializerModel): + ... x: float = 0.0 + ... y: float = 0.0 + + + Clients can optionally customize the base class of the synthesized + Pydantic model using the ``base_model`` decorator parameter. The default + behavior uses :class:`BaseValidatedInitializerModel` and its + `model config `_. + + See Also + -------- + BaseValidatedInitializerModel + Default base class for all synthesized Pydantic models. + """ + + def validator(init): + init_qualname = dict(inspect.getmembers(init))["__qualname__"] + init_clsnme = init_qualname.split(".")[0] + init_params = inspect.signature(init).parameters + init_fields = { + param.name: ( + param.annotation + if param.annotation != inspect.Parameter.empty + else Any, + param.default + if param.default != inspect.Parameter.empty + else ..., + ) + for param in init_params.values() + if param.name != "self" + and param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD + } + + if base_model is None: + PydanticModel = create_model( + model_name=f"{init_clsnme}Model", + __config__=BaseValidatedInitializerModel.Config, + **init_fields, + ) + else: + PydanticModel = create_model( + model_name=f"{init_clsnme}Model", + __base__=base_model, + **init_fields, + ) + + def validated_repr(self) -> str: + return dump_code(self) + + def validated_getnewargs_ex(self): + return (), self.__init_args__ + + @functools.wraps(init) + def init_wrapper(*args, **kwargs): + self, *args = args + + nmargs = { + name: arg + for (name, param), arg in zip( + list(init_params.items()), [self] + args + ) + if name != "self" + } + model = PydanticModel(**{**nmargs, **kwargs}) + + # merge nmargs, kwargs, and the model fields into a single dict + all_args = {**nmargs, **kwargs, **model.__dict__} + + # save the merged dictionary for Representable use, but only of the + # __init_args__ is not already set in order to avoid overriding a + # value set by a subclass initializer in super().__init__ calls + if not getattr(self, "__init_args__", {}): + self.__init_args__ = OrderedDict( + { + name: arg + for name, arg in sorted(all_args.items()) + if type(arg) != torch.nn.ParameterDict + } + ) + self.__class__.__getnewargs_ex__ = validated_getnewargs_ex + self.__class__.__repr__ = validated_repr + + return init(self, **all_args) + + # attach the Pydantic model as the attribute of the initializer wrapper + setattr(init_wrapper, "Model", PydanticModel) + + return init_wrapper + + return validator \ No newline at end of file diff --git a/pts/core/serde.py b/pts/core/serde.py new file mode 100644 index 0000000..4df49ce --- /dev/null +++ b/pts/core/serde.py @@ -0,0 +1,359 @@ +from typing import Any, Optional, cast, NamedTuple +import json +from functools import singledispatch +from pts.core import fqname_for +import numpy as np +import textwrap +from pydoc import locate +import math +import itertools + + +bad_type_msg = textwrap.dedent( + """ + Cannot serialize type {}. See the documentation of the `encode` and + `validate` functions at + + http://gluon-ts.mxnet.io/api/gluonts/gluonts.html + + and the Python documentation of the `__getnewargs_ex__` magic method at + + https://docs.python.org/3/library/pickle.html#object.__getnewargs_ex__ + + for more information how to make this type serializable. + """ +).lstrip() + + +def dump_code(o: Any) -> str: + """ + Serializes an object to a Python code string. + + Parameters + ---------- + o + The object to serialize. + + Returns + ------- + str + A string representing the object as Python code. + + See Also + -------- + load_code + Inverse function. + """ + + def _dump_code(x: Any) -> str: + # r = { 'class': ..., 'args': ... } + # r = { 'class': ..., 'kwargs': ... } + if type(x) == dict and x.get("__kind__") == kind_inst: + args = x.get("args", []) + kwargs = x.get("kwargs", {}) + + fqname = x["class"] + bindings = ", ".join( + itertools.chain( + map(_dump_code, args), + [f"{k}={_dump_code(v)}" for k, v in kwargs.items()], + ) + ) + return f"{fqname}({bindings})" + + if type(x) == dict and x.get("__kind__") == kind_type: + return x["class"] + + if isinstance(x, dict): + inner = ", ".join( + f"{_dump_code(k)}: {_dump_code(v)}" for k, v in x.items() + ) + return f"{{{inner}}}" + + if isinstance(x, list): + inner = ", ".join(list(map(dump_code, x))) + return f"[{inner}]" + + if isinstance(x, tuple): + inner = ", ".join(list(map(dump_code, x))) + # account for the extra `,` in `(x,)` + if len(x) == 1: + inner += "," + return f"({inner})" + + if isinstance(x, str): + # json.dumps escapes the string + return json.dumps(x) + + if isinstance(x, float) or np.issubdtype(type(x), np.inexact): + if math.isfinite(x): + return str(x) + else: + # e.g. `nan` needs to be encoded as `float("nan")` + return 'float("{x}")' + + if isinstance(x, int) or np.issubdtype(type(x), np.integer): + return str(x) + + if x is None: + return str(x) + + raise RuntimeError( + f"Unexpected element type {fqname_for(x.__class__)}" + ) + + return _dump_code(encode(o)) + +# JSON Serialization/Deserialization +# ---------------------------------- + +# The canonical way to do this is to define and `default` and `object_hook` +# parameters to the json.dumps and json.loads methods. Unfortunately, due +# to https://bugs.python.org/issue12657 this is not possible at the moment, +# as support for custom NamedTuple serialization is broken. +# +# To circumvent the issue, we pass the input value through custom encode +# and decode functions that map nested object terms to JSON-serializable +# data structures with explicit recursion. + + + +def dump_json(o: Any, indent: Optional[int] = None) -> str: + """ + Serializes an object to a JSON string. + Parameters + ---------- + o + The object to serialize. + indent + An optional number of spaced to use as an indent. + Returns + ------- + str + A string representing the object in JSON format. + See Also + -------- + load_json + Inverse function. + """ + return json.dumps(encode(o), indent=indent, sort_keys=True) + + +def load_json(s: str) -> Any: + """ + Deserializes an object from a JSON string. + Parameters + ---------- + s + A string representing the object in JSON format. + Returns + ------- + Any + The deserialized object. + See Also + -------- + dump_json + Inverse function. + """ + return decode(json.loads(s)) + + +# Structural encoding/decoding +# ---------------------------- + +kind_type = "type" +kind_inst = "instance" + + +@singledispatch +def encode(v: Any) -> Any: + """ + Transforms a value `v` as a serializable intermediate representation (for + example, named tuples are encoded as dictionaries). The intermediate + representation is then recursively traversed and serialized either as + Python code or as JSON string. + + This function is decorated with :func:`~functools.singledispatch` and can + be specialized by clients for families of types that are not supported by + the basic implementation (explained below). + + Examples + -------- + + The conversion logic implemented by the basic implementation is used + as a fallback and is best explained by a series of examples. + + Lists (as lists). + + >>> encode([1, 2.0, '3']) + [1, 2.0, '3'] + + Tuples (as lists). + + >>> encode((1, 2.0, '3')) + [1, 2.0, '3'] + + Dictionaries (as dictionaries). + + >>> encode({'a': 1, 'b': 2.0, 'c': '3'}) + {'a': 1, 'b': 2.0, 'c': '3'} + + Named tuples (as dictionaries with a ``'__kind__': 'instance'`` member). + + >>> from pprint import pprint + >>> from typing import NamedTuple + >>> class ComplexNumber(NamedTuple): + ... x: float = 0.0 + ... y: float = 0.0 + >>> pprint(encode(ComplexNumber(4.0, 2.0))) + {'__kind__': 'instance', + 'class': 'gluonts.core.serde.ComplexNumber', + 'kwargs': {'x': 4.0, 'y': 2.0}} + + Classes with a :func:`~gluonts.core.component.validated` initializer (as + dictionaries with a ``'__kind__': 'instance'`` member). + + >>> from gluonts.core.component import validated + >>> class ComplexNumber: + ... @validated() + ... def __init__(self, x: float = 0.0, y: float = 0.0) -> None: + ... self.x = x + ... self.y = y + >>> pprint(encode(ComplexNumber(4.0, 2.0))) + {'__kind__': 'instance', + 'args': [], + 'class': 'gluonts.core.serde.ComplexNumber', + 'kwargs': {'x': 4.0, 'y': 2.0}} + + Classes with a ``__getnewargs_ex__`` magic method (as dictionaries with a + ``'__kind__': 'instance'`` member). + + >>> from gluonts.core.component import validated + >>> class ComplexNumber: + ... def __init__(self, x: float = 0.0, y: float = 0.0) -> None: + ... self.x = x + ... self.y = y + ... def __getnewargs_ex__(self): + ... return [], {'x': self.x, 'y': self.y} + >>> pprint(encode(ComplexNumber(4.0, 2.0))) + {'__kind__': 'instance', + 'args': [], + 'class': 'gluonts.core.serde.ComplexNumber', + 'kwargs': {'x': 4.0, 'y': 2.0}} + + + Types (as dictionaries with a ``'__kind__': 'type' member``). + + >>> encode(ComplexNumber) + {'__kind__': 'type', 'class': 'gluonts.core.serde.ComplexNumber'} + + Parameters + ---------- + v + The value to be encoded. + + Returns + ------- + Any + An encoding of ``v`` that can be serialized to Python code or + JSON string. + + See Also + -------- + decode + Inverse function. + dump_json + Serializes an object to a JSON string. + dump_code + Serializes an object to a Python code string. + """ + if isinstance(v, type(None)): + return None + + if isinstance(v, (float, int, str)): + return v + + if np.issubdtype(type(v), np.inexact): + return float(v) + + if np.issubdtype(type(v), np.integer): + return int(v) + + # we have to check for namedtuples first, to encode them not as plain + # tuples (which would become lists) + if isinstance(v, tuple) and hasattr(v, "_asdict"): + v = cast(NamedTuple, v) + return { + "__kind__": kind_inst, + "class": fqname_for(v.__class__), + "kwargs": encode(v._asdict()), + } + + if isinstance(v, (list, set, tuple)): + return list(map(encode, v)) + + if isinstance(v, dict): + return {k: encode(v) for k, v in v.items()} + + if isinstance(v, type): + return {"__kind__": kind_type, "class": fqname_for(v)} + + if hasattr(v, "__getnewargs_ex__"): + args, kwargs = v.__getnewargs_ex__() # mypy: ignore + return { + "__kind__": kind_inst, + "class": fqname_for(v.__class__), + "args": encode(args), + "kwargs": encode(kwargs), + } + + raise RuntimeError(bad_type_msg.format(fqname_for(v.__class__))) + + +def decode(r: Any) -> Any: + """ + Decodes a value from an intermediate representation `r`. + + Parameters + ---------- + r + An intermediate representation to be decoded. + + Returns + ------- + Any + A Python data structure corresponding to the decoded version of ``r``. + + See Also + -------- + encode + Inverse function. + """ + + # structural recursion over the possible shapes of r + # r = { 'class': ..., 'args': ... } + # r = { 'class': ..., 'kwargs': ... } + if type(r) == dict and r.get("__kind__") == kind_inst: + cls = cast(Any, locate(r["class"])) + args = decode(r["args"]) if "args" in r else [] + kwargs = decode(r["kwargs"]) if "kwargs" in r else {} + return cls(*args, **kwargs) + # r = { 'class': ..., 'args': ... } + # r = { 'class': ..., 'kwargs': ... } + if type(r) == dict and r.get("__kind__") == kind_type: + return locate(r["class"]) + # r = { k1: v1, ..., kn: vn } + elif type(r) == dict: + return {k: decode(v) for k, v in r.items()} + # r = ( y1, ..., yn ) + elif type(r) == tuple: + return tuple([decode(y) for y in r]) + # r = [ y1, ..., yn ] + elif type(r) == list: + return [decode(y) for y in r] + # r = { y1, ..., yn } + elif type(r) == set: + return {decode(y) for y in r} + # r = a + else: + return r diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py index 766a5ee..81cb319 100644 --- a/pts/model/deepar/deepar_network.py +++ b/pts/model/deepar/deepar_network.py @@ -5,10 +5,11 @@ import torch.nn as nn from torch.distributions import Distribution import numpy as np - +from pts.core.component import validated from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder from pts.model import weighted_average + def prod(xs): p = 1 for x in xs: @@ -17,6 +18,8 @@ def prod(xs): class DeepARNetwork(nn.Module): + + @validated() def __init__( self, input_size: int, diff --git a/pts/model/deepvar/deepvar_network.py b/pts/model/deepvar/deepvar_network.py index ab20c0d..754f265 100644 --- a/pts/model/deepvar/deepvar_network.py +++ b/pts/model/deepvar/deepvar_network.py @@ -5,12 +5,13 @@ import torch.nn as nn from torch.distributions import Distribution import numpy as np - +from pts.core.component import validated from pts.modules import DistributionOutput, MeanScaler, NOPScaler from pts.model import weighted_average class DeepVARTrainingNetwork(nn.Module): + @validated() def __init__( self, input_size: int, diff --git a/pts/model/forecast_generator.py b/pts/model/forecast_generator.py index 7897ce7..93e4cab 100644 --- a/pts/model/forecast_generator.py +++ b/pts/model/forecast_generator.py @@ -8,6 +8,7 @@ import torch.nn as nn from pts.dataset import InferenceDataLoader, DataEntry, FieldName from pts.modules import DistributionOutput from .forecast import Forecast, DistributionForecast, QuantileForecast, SampleForecast +from pts.core.component import validated OutputTransform = Callable[[DataEntry, np.ndarray], np.ndarray] @@ -133,6 +134,11 @@ class QuantileForecastGenerator(ForecastGenerator): class SampleForecastGenerator(ForecastGenerator): + + @validated() + def __init__(self): + pass + def __call__( self, inference_data_loader: InferenceDataLoader, diff --git a/pts/model/predictor.py b/pts/model/predictor.py index 39a6e77..e567059 100644 --- a/pts/model/predictor.py +++ b/pts/model/predictor.py @@ -1,21 +1,27 @@ -from abc import ABC, abstractmethod -from typing import Iterator, Callable, Optional - +import pts +import json import numpy as np import torch import torch.nn as nn - +from abc import ABC, abstractmethod +from pydoc import locate +from typing import Iterator, Callable, Optional from pts.dataset import Dataset, DataEntry, InferenceDataLoader from pts.transform import Transformation - +from pathlib import Path from .forecast import Forecast from .forecast_generator import ForecastGenerator, SampleForecastGenerator from .utils import get_module_forward_input_names +from pts.core.serde import dump_json, fqname_for, load_json + OutputTransform = Callable[[DataEntry, np.ndarray], np.ndarray] class Predictor(ABC): + + __version__: str = pts.__version__ + def __init__(self, prediction_length: int, freq: str) -> None: self.prediction_length = prediction_length self.freq = freq @@ -24,6 +30,42 @@ class Predictor(ABC): def predict(self, dataset: Dataset, **kwargs) -> Iterator[Forecast]: pass + def serialize(self, path: Path) -> None: + # serialize Predictor type + with (path / "type.txt").open("w") as fp: + fp.write(fqname_for(self.__class__)) + with (path / "version.json").open("w") as fp: + json.dump( + {"model": self.__version__, "pts": pts.__version__}, fp + ) + + @classmethod + def deserialize( + cls, path: Path, device: Optional[torch.device] = None + ) -> "Predictor": + """ + Load a serialized predictor from the given path + Parameters + ---------- + path + Path to the serialized files predictor. + device + Optional pytorch to be used with the predictor. + If nothing is passed will use the GPU if available and CPU otherwise. + """ + # deserialize Predictor type + with (path / "type.txt").open("r") as fp: + tpe = locate(fp.readline()) + + # ensure that predictor_cls is a subtype of Predictor + if not issubclass(tpe, Predictor): + raise IOError( + f"Class {fqname_for(tpe)} is not " + f"a subclass of {fqname_for(Predictor)}" + ) + # call deserialize() for the concrete Predictor type + return tpe.deserialize(path, device) + class PTSPredictor(Predictor): def __init__( @@ -39,7 +81,6 @@ class PTSPredictor(Predictor): 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 @@ -71,3 +112,64 @@ class PTSPredictor(Predictor): output_transform=self.output_transform, num_samples=num_samples, ) + + def serialize(self, path: Path) -> None: + + super().serialize(path) + + # serialize network + model_name = 'prediction_net' + with (path / f"{model_name}-network.json").open("w") as fp: + print(dump_json(self.prediction_net), file=fp) + torch.save(self.prediction_net.state_dict(), path / "prediction_net") + + # serialize input transformation chain + with (path / "input_transform.json").open("w") as fp: + print(dump_json(self.input_transform), file=fp) + + # serialize output transformation chain + with (path / "output_transform.json").open("w") as fp: + print(dump_json(self.output_transform), file=fp) + + # serialize all remaining constructor parameters + with (path / "parameters.json").open("w") as fp: + parameters = dict( + batch_size=self.batch_size, + prediction_length=self.prediction_length, + freq=self.freq, + dtype=self.dtype, + forecast_generator=self.forecast_generator, + input_names=self.input_names, + ) + print(dump_json(parameters), file=fp) + + @classmethod + def deserialize( + cls, path: Path, device: Optional[torch.device] = None + ) -> "PTSPredictor": + + # deserialize constructor parameters + with (path / "parameters.json").open("r") as fp: + parameters = load_json(fp.read()) + + # deserialize transformation chain + with (path / "input_transform.json").open("r") as fp: + transformation = load_json(fp.read()) + + # deserialize prediction network + model_name = 'prediction_net' + with (path / f"{model_name}-network.json").open("r") as fp: + prediction_net = load_json(fp.read()) + prediction_net.load_state_dict(torch.load(path / "prediction_net")) + + # input_names is derived from the prediction_net + if "input_names" in parameters: + del parameters["input_names"] + + parameters["device"] = device + + return PTSPredictor( + input_transform=transformation, + prediction_net=prediction_net, + **parameters + ) diff --git a/pts/model/simple_feedforward/simple_feedforward_network.py b/pts/model/simple_feedforward/simple_feedforward_network.py index ca56059..6749f59 100644 --- a/pts/model/simple_feedforward/simple_feedforward_network.py +++ b/pts/model/simple_feedforward/simple_feedforward_network.py @@ -2,8 +2,8 @@ from typing import List import torch import torch.nn as nn -import torch.nn.functional as F from torch.distributions import Distribution +from pts.core.component import validated from pts.modules import MeanScaler, NOPScaler, DistributionOutput, LambdaLayer @@ -35,7 +35,7 @@ class SimpleFeedForwardNetworkBase(nn.Module): Distribution to fit. kwargs """ - + @validated() def __init__( self, num_hidden_dimensions: List[int], @@ -104,6 +104,7 @@ class SimpleFeedForwardTrainingNetwork(SimpleFeedForwardNetworkBase): class SimpleFeedForwardPredictionNetwork(SimpleFeedForwardNetworkBase): + @validated() def __init__(self, num_parallel_samples: int = 100, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.num_parallel_samples = num_parallel_samples diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py index acbe637..94cf29e 100644 --- a/pts/model/tempflow/tempflow_network.py +++ b/pts/model/tempflow/tempflow_network.py @@ -4,12 +4,14 @@ import torch import torch.nn as nn import numpy as np - +from pts.core.component import validated from pts.modules import RealNVP, MAF, FlowOutput, MeanScaler, NOPScaler from pts.model import weighted_average class TempFlowTrainingNetwork(nn.Module): + + @validated() def __init__( self, input_size: int, diff --git a/pts/model/transformer/transformer_network.py b/pts/model/transformer/transformer_network.py index a023f51..12f6256 100644 --- a/pts/model/transformer/transformer_network.py +++ b/pts/model/transformer/transformer_network.py @@ -5,7 +5,7 @@ import torch.nn as nn from torch.distributions import Distribution import numpy as np - +from pts.core.component import validated from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder from pts.model import weighted_average @@ -18,6 +18,8 @@ def prod(xs): class TransformerNetwork(nn.Module): + + @validated() def __init__( self, input_size: int, diff --git a/pts/model/transformer_tempflow/transformer_tempflow_network.py b/pts/model/transformer_tempflow/transformer_tempflow_network.py index b9292cf..37c1b87 100644 --- a/pts/model/transformer_tempflow/transformer_tempflow_network.py +++ b/pts/model/transformer_tempflow/transformer_tempflow_network.py @@ -4,12 +4,14 @@ import torch import torch.nn as nn import numpy as np - +from pts.core.component import validated from pts.modules import RealNVP, MAF, FlowOutput, MeanScaler, NOPScaler from pts.model import weighted_average class TransformerTempFlowTrainingNetwork(nn.Module): + + @validated() def __init__( self, input_size: int, diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index 101bbba..0377f2d 100644 --- a/pts/modules/distribution_output.py +++ b/pts/modules/distribution_output.py @@ -5,6 +5,7 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +from pts.core.component import validated from torch.distributions import ( Distribution, Beta, @@ -74,8 +75,13 @@ class Output(ABC): class DistributionOutput(Output, ABC): + distr_cls: type + @validated() + def __init__(self) -> None: + pass + def distribution( self, distr_args, scale: Optional[torch.Tensor] = None ) -> Distribution: @@ -146,6 +152,7 @@ class StudentTOutput(DistributionOutput): class LowRankMultivariateNormalOutput(DistributionOutput): + def __init__( self, dim: int, rank: int, sigma_init: float = 1.0, sigma_minimum: float = 1e-3, ) -> None: @@ -179,6 +186,7 @@ class LowRankMultivariateNormalOutput(DistributionOutput): class IndependentNormalOutput(DistributionOutput): + def __init__(self, dim: int) -> None: self.dim = dim self.args_dim = {"loc": self.dim, "scale": self.dim} @@ -202,6 +210,7 @@ class IndependentNormalOutput(DistributionOutput): class MultivariateNormalOutput(DistributionOutput): + def __init__(self, dim: int) -> None: self.args_dim = {"loc": dim, "scale_tril": dim * dim} self.dim = dim @@ -239,8 +248,8 @@ class MultivariateNormalOutput(DistributionOutput): return (self.dim,) - class FlowOutput(DistributionOutput): + def __init__(self, flow, input_size, cond_size): self.args_dim = {"cond": cond_size} self.flow = flow @@ -260,3 +269,4 @@ class FlowOutput(DistributionOutput): @property def event_shape(self) -> Tuple: return (self.dim,) + diff --git a/pts/transform/sampler.py b/pts/transform/sampler.py index 40caf4c..a4c2a95 100644 --- a/pts/transform/sampler.py +++ b/pts/transform/sampler.py @@ -16,6 +16,7 @@ from abc import ABC, abstractmethod import numpy as np from pts.dataset.stat import ScaleHistogram +from pts.core.component import validated class InstanceSampler(ABC): @@ -92,7 +93,7 @@ class ExpectedNumInstanceSampler(InstanceSampler): num_instances number of training examples generated per time series on average """ - + @validated() def __init__(self, num_instances: float) -> None: self.num_instances = num_instances self.total_length = 0 diff --git a/pts/transform/split.py b/pts/transform/split.py index a5342ba..d2ad55e 100644 --- a/pts/transform/split.py +++ b/pts/transform/split.py @@ -12,7 +12,7 @@ # permissions and limitations under the License. from functools import lru_cache -from typing import Iterator, List, Optional, Union +from typing import Iterator, List, Optional import numpy as np import pandas as pd @@ -21,6 +21,7 @@ from pts.dataset import DataEntry, FieldName from .transform import FlatMapTransformation from .sampler import InstanceSampler, ContinuousTimePointSampler +from pts.core.component import validated def shift_timestamp(ts: pd.Timestamp, offset: int) -> pd.Timestamp: @@ -106,6 +107,7 @@ class InstanceSplitter(FlatMapTransformation): data is padded or not. """ + @validated() def __init__( self, target_field: str, diff --git a/pts/transform/transform.py b/pts/transform/transform.py index c250e4d..0fe680f 100644 --- a/pts/transform/transform.py +++ b/pts/transform/transform.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Callable, Iterator, Iterable, List from functools import reduce - +from pts.core.component import validated from pts.dataset import DataEntry @@ -30,7 +30,7 @@ class Chain(Transformation): """ Chain multiple transformations together. """ - + @validated() def __init__(self, trans: List[Transformation]) -> None: self.transformations = [] for transformation in trans: