# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.
from collections import OrderedDict
from typing import Optional, Tuple
from mxnet import gluon
from gluonts.core.component import validated
from gluonts.mx import Tensor
from . import LowrankMultivariateGaussian
from .distribution_output import (
ArgProj,
DistributionOutput,
AffineTransformedDistribution,
)
from .lowrank_multivariate_gaussian import inv_softplus, sigma_minimum
[docs]class GPArgProj(gluon.HybridBlock):
@validated()
def __init__(
self,
rank: int,
sigma_init: float = 1.0,
sigma_minimum: float = sigma_minimum,
mu_ratio: float = 1.0,
dropout_rate: float = 0.0,
prefix: Optional[str] = None,
) -> None:
super().__init__(prefix=prefix)
self.param_dim_args = OrderedDict({"mu": 1, "sigma": 1, "w": rank})
self.mu_ratio = mu_ratio
self.sigma_init = sigma_init
self.sigma_minimum = sigma_minimum
self.W_ratio = 1.0
self.rank = rank
def make(name, param_dim):
net = gluon.nn.HybridSequential()
net.add(
gluon.nn.Dense(
param_dim,
flatten=False,
prefix=f"distr_{name}_",
weight_initializer="Xavier",
)
)
if dropout_rate > 0:
net.add(gluon.nn.Dropout(dropout_rate))
return net
self.proj = [
make(name, param_dim)
for name, param_dim in self.param_dim_args.items()
]
for dense in self.proj:
self.register_child(dense)
[docs] def hybrid_forward(self, F, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
"""
Parameters
----------
F
x : (..., dim, hidden_dim)
Returns
-------
Returns (mu, D, W) where shapes are (..., dim), (..., dim),
(..., dim, rank)
"""
# TODO 2 concatenate inputs features to x, better names would be great
# (..., dim)
mu_vector = self.proj[0](x).squeeze(axis=-1)
mu = mu_vector * self.mu_ratio
# (..., dim, rank)
W_matrix = self.proj[2](x) * self.W_ratio
# (..., |x| + 1)
x_plus_w = F.concat(
x, W_matrix.square().sum(axis=-1, keepdims=True), dim=-1
)
# (..., dim)
D_vector = self.proj[1](x_plus_w).squeeze(axis=-1)
d_bias = (
0.0 if self.sigma_init == 0.0 else inv_softplus(self.sigma_init**2)
)
D_positive = (
F.Activation(D_vector + d_bias, act_type="softrelu")
+ self.sigma_minimum
)
return mu, D_positive, W_matrix
[docs]class LowrankGPOutput(DistributionOutput):
@validated()
def __init__(
self,
rank: int,
dim: Optional[int] = None, # needed to compute variance
sigma_init: float = 1.0,
mu_ratio: float = 1.0,
dropout_rate: float = 0.0,
) -> None:
super().__init__(self)
self.dist_cls = LowrankMultivariateGaussian
self.dim = dim
self.rank = rank
self.args_dim = {"mu": 1, "sigma": 1, "w": rank}
self.mu_bias = 0.0
self.sigma_init = sigma_init
self.mu_ratio = mu_ratio
self.dropout_rate = dropout_rate
[docs] def get_args_proj(self, prefix: Optional[str] = None) -> ArgProj:
return GPArgProj(
rank=self.rank,
mu_ratio=self.mu_ratio,
sigma_init=self.sigma_init,
dropout_rate=self.dropout_rate,
prefix=prefix,
)
[docs] def distribution(self, distr_args, loc=None, scale=None, dim=None):
dist = LowrankMultivariateGaussian(dim, self.rank, *distr_args)
if loc is None and scale is None:
return dist
else:
return AffineTransformedDistribution(dist, loc=loc, scale=scale)
@property
def event_shape(self) -> Tuple:
return (self.dim,)
@property
def event_dim(self) -> int:
return 1