initial gluonts dependency

This commit is contained in:
Dr. Kashif Rasul
2020-12-17 17:04:56 +01:00
parent ecc31f6082
commit b072ab227b
88 changed files with 498 additions and 11571 deletions
-4
View File
@@ -1,7 +1,4 @@
from .distribution_output import (
ArgProj,
Output,
DistributionOutput,
NormalOutput,
StudentTOutput,
BetaOutput,
@@ -20,5 +17,4 @@ from .distribution_output import (
)
from .feature import FeatureEmbedder, FeatureAssembler
from .flows import RealNVP, MAF
from .lambda_layer import LambdaLayer
from .scaler import MeanScaler, NOPScaler
+39 -103
View File
@@ -19,7 +19,8 @@ from torch.distributions import (
MultivariateNormal,
TransformedDistribution,
AffineTransform,
Poisson)
Poisson,
)
from pts.distributions import (
ZeroInflatedPoisson,
@@ -29,79 +30,13 @@ from pts.distributions import (
ImplicitQuantile,
TransformedImplicitQuantile,
)
from pts.core.component import validated
from gluonts.core.component import validated
from gluonts.torch.modules.distribution_output import (
DistributionOutput,
LambdaLayer,
PtArgProj,
)
from pts.modules.iqn_modules import ImplicitQuantileModule
from .lambda_layer import LambdaLayer
class ArgProj(nn.Module):
def __init__(
self,
in_features: int,
args_dim: Dict[str, int],
domain_map: Callable[..., Tuple[torch.Tensor]],
dtype: np.dtype = np.float32,
prefix: Optional[str] = None,
**kwargs,
):
super().__init__(**kwargs)
self.args_dim = args_dim
self.dtype = dtype
self.proj = nn.ModuleList(
[nn.Linear(in_features, dim) for dim in args_dim.values()]
)
self.domain_map = domain_map
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
params_unbounded = [proj(x) for proj in self.proj]
return self.domain_map(*params_unbounded)
class Output(ABC):
in_features: int
args_dim: Dict[str, int]
_dtype: np.dtype = np.float32
@property
def dtype(self):
return self._dtype
@dtype.setter
def dtype(self, dtype: np.dtype):
self._dtype = dtype
def get_args_proj(self, in_features: int, prefix: Optional[str] = None) -> ArgProj:
return ArgProj(
in_features=in_features,
args_dim=self.args_dim,
domain_map=LambdaLayer(self.domain_map),
prefix=prefix,
dtype=self.dtype,
)
@abstractclassmethod
def domain_map(cls, *args: torch.Tensor):
pass
class DistributionOutput(Output, ABC):
distr_cls: type
@validated()
def __init__(self) -> None:
pass
def distribution(
self, distr_args, scale: Optional[torch.Tensor] = None
) -> Distribution:
distr = self.distr_cls(*distr_args)
if scale is None:
return distr
else:
return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)])
class IndependentDistributionOutput(DistributionOutput):
@@ -364,7 +299,9 @@ class PiecewiseLinearOutput(DistributionOutput):
return gamma.squeeze(axis=-1), slopes_proj, knot_spacings_proj
def distribution(
self, distr_args, scale: Optional[torch.Tensor] = None,
self,
distr_args,
scale: Optional[torch.Tensor] = None,
) -> PiecewiseLinear:
if scale is None:
return self.distr_cls(*distr_args)
@@ -415,7 +352,11 @@ class NormalMixtureOutput(DistributionOutput):
class LowRankMultivariateNormalOutput(DistributionOutput):
@validated()
def __init__(
self, dim: int, rank: int, sigma_init: float = 1.0, sigma_minimum: float = 1e-3,
self,
dim: int,
rank: int,
sigma_init: float = 1.0,
sigma_minimum: float = 1e-3,
) -> None:
self.distr_cls = LowRankMultivariateNormal
self.dim = dim
@@ -508,25 +449,16 @@ class FlowOutput(DistributionOutput):
return (self.dim,)
class QuantileArgProj(ArgProj):
class QuantilePtArgProj(PtArgProj):
def __init__(
self,
in_features: int,
output_domain_cls: nn.Module,
args_dim: Dict[str, int],
domain_map: Callable[..., Tuple[torch.Tensor]],
dtype: np.dtype = np.float32,
prefix: Optional[str] = None,
**kwargs,
self,
in_features: int,
output_domain_cls: nn.Module,
args_dim: Dict[str, int],
domain_map: Callable[..., Tuple[torch.Tensor]],
**kwargs,
):
super().__init__(
in_features,
args_dim,
domain_map,
dtype,
prefix,
**kwargs
)
super().__init__(in_features, args_dim, domain_map, **kwargs)
self.output_domain_cls = output_domain_cls
self.proj = ImplicitQuantileModule(in_features, output_domain_cls)
@@ -535,8 +467,8 @@ class QuantileArgProj(ArgProj):
forecast_length = x.shape[1]
device = x.device
taus = torch.rand(size=(batch_size, forecast_length), device=device)
self.register_buffer('taus', taus)
self.register_buffer('nn_ouput', x.clone().detach())
self.register_buffer("taus", taus)
self.register_buffer("nn_ouput", x.clone().detach())
predicted_quantiles = self.proj(x, taus)
return self.domain_map(predicted_quantiles)
@@ -548,6 +480,7 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
output_domain_cls: type = nn.Module
quantile_arg_proj: type = nn.Module
@validated()
def __init__(self, output_domain: str) -> None:
super().__init__()
self.set_output_domain_map(output_domain)
@@ -559,14 +492,17 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
"Positive": nn.Softplus,
"Real": nn.Identity,
}
assert output_domain in available_domain_map_cls.keys(), \
"Only the following output domains are allowed: {}".format(available_domain_map_cls.keys())
assert (
output_domain in available_domain_map_cls.keys()
), "Only the following output domains are allowed: {}".format(
available_domain_map_cls.keys()
)
output_domain_cls = available_domain_map_cls[output_domain]
cls.output_domain_cls = output_domain_cls
@classmethod
def set_args_proj(cls):
cls.quantile_arg_proj = QuantileArgProj(
cls.quantile_arg_proj = QuantilePtArgProj(
in_features=cls.in_features,
output_domain_cls=cls.output_domain_cls,
args_dim=cls.args_dim,
@@ -584,11 +520,13 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
cls.set_args_proj()
return cls.quantile_arg_proj
def get_args_proj(self, in_features: int, prefix: Optional[str] = None) :
def get_args_proj(self, in_features: int, prefix: Optional[str] = None):
return self.args_proj(in_features)
def distribution(
self, distr_args, scale: Optional[torch.Tensor] = None,
self,
distr_args,
scale: Optional[torch.Tensor] = None,
) -> ImplicitQuantile:
args_proj = self.get_args_proj(self.in_features)
@@ -597,7 +535,8 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
implicit_quantile_function=implicit_quantile_function,
taus=list(args_proj.buffers())[0],
nn_output=list(args_proj.buffers())[1],
predicted_quantiles=distr_args)
predicted_quantiles=distr_args,
)
if scale is None:
return distr
else:
@@ -608,6 +547,3 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
@property
def event_shape(self) -> Tuple:
return ()
+5 -1
View File
@@ -5,7 +5,11 @@ import torch.nn as nn
class FeatureEmbedder(nn.Module):
def __init__(self, cardinalities: List[int], embedding_dims: List[int],) -> None:
def __init__(
self,
cardinalities: List[int],
embedding_dims: List[int],
) -> None:
super().__init__()
self.__num_features = len(cardinalities)
+1 -1
View File
@@ -52,6 +52,6 @@ class QuantileLayer(nn.Module):
integers = torch.repeat_interleave(
torch.arange(0, self.n_cos_embedding).unsqueeze(dim=0),
repeats=tau.shape[-1],
dim=0
dim=0,
).to(tau.device)
return torch.cos(pi * tau.unsqueeze(dim=-1) * integers)
-10
View File
@@ -1,10 +0,0 @@
import torch.nn as nn
class LambdaLayer(nn.Module):
def __init__(self, function):
super().__init__()
self._func = function
def forward(self, x, *args):
return self._func(x, *args)
+1 -1
View File
@@ -37,7 +37,7 @@ class Scaler(ABC, nn.Module):
Tensor
Tensor containing the "scaled" data, shape: (N, T, C) or (N, C, T).
Tensor
Tensor containing the scale, of shape (N, C) if ``keepdim == False``,
Tensor containing the scale, of shape (N, C) if ``keepdim == False``,
and shape (N, 1, C) or (N, C, 1) if ``keepdim == True``.
"""