From 085ee835a52c72816d8e22dd4c1f31657d66a4a1 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 15 May 2020 13:04:45 +0200 Subject: [PATCH] Mixture output (#13) * get dataset returns shuffled training data by default * added Gaussian mixture output distribution * added StudentTMixtureOutput --- pts/dataset/repository/datasets.py | 4 ++ pts/dataset/utils.py | 2 + pts/modules/__init__.py | 2 + pts/modules/distribution_output.py | 97 ++++++++++++++++++++++++++---- 4 files changed, 94 insertions(+), 11 deletions(-) diff --git a/pts/dataset/repository/datasets.py b/pts/dataset/repository/datasets.py index e173e6f..58cc7c9 100644 --- a/pts/dataset/repository/datasets.py +++ b/pts/dataset/repository/datasets.py @@ -133,6 +133,7 @@ def materialize_dataset( def get_dataset( dataset_name: str, path: Path = default_dataset_path, regenerate: bool = False, + shuffle: bool = True, ) -> TrainDatasets: """ Get a repository dataset. @@ -151,6 +152,8 @@ def get_dataset( be downloaded again. path where the dataset should be saved + shuffle + wheather to shuffle the training time series Returns ------- dataset obtained by either downloading or reloading from local file. @@ -161,6 +164,7 @@ def get_dataset( metadata=dataset_path / "metadata.json", train=dataset_path / "train" / "*.json", test=dataset_path / "test" / "*.json", + shuffle=shuffle ) diff --git a/pts/dataset/utils.py b/pts/dataset/utils.py index 5178a1d..f7e30c0 100644 --- a/pts/dataset/utils.py +++ b/pts/dataset/utils.py @@ -67,6 +67,8 @@ def load_datasets(metadata, train, test, shuffle: bool = False) -> TrainDatasets Path to the training dataset files. test Path to the test dataset files. + shuffle + Return shuffled train data. Returns ------- TrainDatasets diff --git a/pts/modules/__init__.py b/pts/modules/__init__.py index df0beba..4e0fadc 100644 --- a/pts/modules/__init__.py +++ b/pts/modules/__init__.py @@ -6,6 +6,8 @@ from .distribution_output import ( StudentTOutput, BetaOutput, NegativeBinomialOutput, + NormalMixtureOutput, + StudentTMixtureOutput, IndependentNormalOutput, LowRankMultivariateNormalOutput, MultivariateNormalOutput, diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index c124e2c..66f4345 100644 --- a/pts/modules/distribution_output.py +++ b/pts/modules/distribution_output.py @@ -1,4 +1,4 @@ -from abc import ABC, abstractmethod +from abc import ABC, abstractclassmethod from typing import Callable, Dict, Optional, Tuple import numpy as np @@ -11,6 +11,8 @@ from torch.distributions import ( NegativeBinomial, StudentT, Normal, + Categorical, + MixtureSameFamily, Independent, LowRankMultivariateNormal, MultivariateNormal, @@ -68,8 +70,8 @@ class Output(ABC): dtype=self.dtype, ) - @abstractmethod - def domain_map(self, *args: torch.Tensor): + @abstractclassmethod + def domain_map(cls, *args: torch.Tensor): pass @@ -85,10 +87,10 @@ class DistributionOutput(Output, ABC): self, distr_args, scale: Optional[torch.Tensor] = None ) -> Distribution: + distr = self.distr_cls(*distr_args) if scale is None: - return self.distr_cls(*distr_args) + return distr else: - distr = self.distr_cls(*distr_args) return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) @@ -97,7 +99,7 @@ class NormalOutput(DistributionOutput): distr_cls: type = Normal @classmethod - def domain_map(self, loc, scale): + def domain_map(cls, loc, scale): scale = F.softplus(scale) return loc.squeeze(-1), scale.squeeze(-1) @@ -165,6 +167,75 @@ class StudentTOutput(DistributionOutput): return () +class StudentTMixtureOutput(DistributionOutput): + def __init__(self, components: int = 1) -> None: + self.components = components + self.args_dim = { + "mix_logits": components, + "df": components, + "loc": components, + "scale": components, + } + + @classmethod + def domain_map(cls, mix_logits, df, loc, scale): + scale = F.softplus(scale) + df = 2.0 + F.softplus(df) + return ( + mix_logits.squeeze(-1), + df.squeeze(-1), + loc.squeeze(-1), + scale.squeeze(-1), + ) + + def distribution( + self, distr_args, scale: Optional[torch.Tensor] = None + ) -> Distribution: + mix_logits, df, loc, scale = distr_args + + distr = MixtureSameFamily( + Categorical(logits=mix_logits), StudentT(df, loc, scale) + ) + if scale is None: + return distr + else: + return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) + + @property + def event_shape(self) -> Tuple: + return () + + +class NormalMixtureOutput(DistributionOutput): + def __init__(self, components: int = 1) -> None: + self.components = components + self.args_dim = { + "mix_logits": components, + "loc": components, + "scale": components, + } + + @classmethod + def domain_map(cls, mix_logits, loc, scale): + scale = F.softplus(scale) + return mix_logits.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1) + + def distribution( + self, distr_args, scale: Optional[torch.Tensor] = None + ) -> Distribution: + mix_logits, loc, scale = distr_args + + distr = MixtureSameFamily(Categorical(logits=mix_logits), Normal(loc, scale)) + if scale is None: + return distr + else: + return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) + + @property + def event_shape(self) -> Tuple: + return () + + class LowRankMultivariateNormalOutput(DistributionOutput): def __init__( self, dim: int, rank: int, sigma_init: float = 1.0, sigma_minimum: float = 1e-3, @@ -176,7 +247,8 @@ class LowRankMultivariateNormalOutput(DistributionOutput): self.sigma_minimum = sigma_minimum self.args_dim = {"loc": dim, "cov_factor": dim * rank, "cov_diag": dim} - def domain_map(self, loc, cov_factor, cov_diag): + @classmethod + def domain_map(cls, loc, cov_factor, cov_diag): diag_bias = ( self.inv_softplus(self.sigma_init ** 2) if self.sigma_init > 0.0 else 0.0 ) @@ -203,7 +275,8 @@ class IndependentNormalOutput(DistributionOutput): self.dim = dim self.args_dim = {"loc": self.dim, "scale": self.dim} - def domain_map(self, loc, scale): + @classmethod + def domain_map(cls, loc, scale): return loc, F.softplus(scale) @property @@ -226,8 +299,9 @@ class MultivariateNormalOutput(DistributionOutput): self.args_dim = {"loc": dim, "scale_tril": dim * dim} self.dim = dim - def domain_map(self, loc, scale): - d = self.dim + @classmethod + def domain_map(cls, loc, scale): + d = len(loc) device = scale.device shape = scale.shape[:-1] + (d, d) @@ -264,7 +338,8 @@ class FlowOutput(DistributionOutput): self.flow = flow self.dim = input_size - def domain_map(self, cond): + @classmethod + def domain_map(cls, cond): return (cond,) def distribution(self, distr_args, scale=None):