Model serialization (#6)

* wip: serialization ran successfully

* wip: deserialization ran successfully

.gitignore

@@ -102,3 +102,5 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+.idea/
+runs/

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"

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__)

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__}"

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 <https://pydantic-docs.helpmanual.io/#model-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 <https://pydantic-docs.helpmanual.io/>`_, which can be
+    accessed through the ``Model`` attribute of the decorated initiazlier.
+
+    >>> ComplexNumber.__init__.Model
+    <class 'ComplexNumberModel'>
+
+    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 <https://pydantic-docs.helpmanual.io/#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

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

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,

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,

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,

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
+        )

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

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,

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,

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,

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,)
+

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

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,

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: