""" """
# Created by Wenjie Du <wenjay.du@gmail.com>
# License: BSD-3-Clause
from typing import Tuple, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.parameter import Parameter
[docs]
class ImplicitImputation(nn.Module):
def __init__(self, d_input: int):
super().__init__()
self.projection_layer = nn.Linear(d_input, d_input, bias=False)
[docs]
def forward(self, X: torch.Tensor, missing_mask: torch.Tensor) -> torch.Tensor:
imputation = self.projection_layer(X)
imputed_X = X * missing_mask + imputation * (1 - X)
return imputed_X
[docs]
class PeepholeLSTMCell(nn.LSTMCell):
"""
Notes
-----
This implementation is adapted from https://gist.github.com/Kaixhin/57901e91e5c5a8bac3eb0cbbdd3aba81
"""
def __init__(self, input_size: int, hidden_size: int, bias: bool = True):
super().__init__(input_size, hidden_size, bias)
self.weight_ch = Parameter(torch.Tensor(3 * hidden_size, hidden_size))
if bias:
self.bias_ch = Parameter(torch.Tensor(3 * hidden_size))
else:
self.register_parameter("bias_ch", None)
self.register_buffer("wc_blank", torch.zeros(hidden_size))
self.reset_parameters()
[docs]
def forward(
self,
X: torch.Tensor,
hx: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if hx is None:
zeros = torch.zeros(X.size(0), self.hidden_size, dtype=X.dtype, device=X.device)
hx = (zeros, zeros)
h, c = hx
wx = F.linear(X, self.weight_ih, self.bias_ih)
wh = F.linear(h, self.weight_hh, self.bias_hh)
wc = F.linear(c, self.weight_ch, self.bias_ch)
wxhc = (
wx
+ wh
+ torch.cat(
(
wc[:, : 2 * self.hidden_size],
Variable(self.wc_blank).expand_as(h),
wc[:, 2 * self.hidden_size :],
),
dim=1,
)
)
i = torch.sigmoid(wxhc[:, : self.hidden_size])
f = torch.sigmoid(wxhc[:, self.hidden_size : 2 * self.hidden_size])
g = torch.tanh(wxhc[:, 2 * self.hidden_size : 3 * self.hidden_size])
o = torch.sigmoid(wxhc[:, 3 * self.hidden_size :])
c = f * c + i * g
h = o * torch.tanh(c)
return h, c
[docs]
class GMMLayer(nn.Module):
def __init__(self, d_hidden: int, n_clusters: int):
super().__init__()
self.mu_c_unscaled = Parameter(torch.Tensor(n_clusters, d_hidden))
self.var_c_unscaled = Parameter(torch.Tensor(n_clusters, d_hidden))
self.phi_c_unscaled = Parameter(torch.Tensor(n_clusters))
def set_values(
self,
mu: torch.Tensor,
var: torch.Tensor,
phi: torch.Tensor,
) -> None:
assert mu.shape == self.mu_c_unscaled.shape
assert var.shape == self.var_c_unscaled.shape
assert phi.shape == self.phi_c_unscaled.shape
self.mu_c_unscaled = torch.nn.Parameter(mu)
self.var_c_unscaled = torch.nn.Parameter(var)
self.phi_c_unscaled = torch.nn.Parameter(phi)
[docs]
def forward(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
mu_c = self.mu_c_unscaled
var_c = F.softplus(self.var_c_unscaled)
phi_c = torch.softmax(self.phi_c_unscaled, dim=0)
return mu_c, var_c, phi_c