timesead.models.common.vae

Classes

DenseVAEEncoder	Base class for all neural network modules.
VAE	VAE Implementation that supports normal distribution with diagonal cov matrix in the latent space
VAELoss	Base class for all neural network modules.
NNLLoss	Base class for all neural network modules.
KLLoss	Base class for all neural network modules.
DistVAE	VAE Implementation that supports arbitrary torch.distributions as long as the latent distribution supports
NormalVAEEncoder	Generic class that returns a normal distribution parameterized by mean and standard deviation.
DenseNormalVAEEncoder	Generic class that returns a normal distribution parameterized by mean and standard deviation.
VAEDistLoss	Base class for all neural network modules.

Functions

sample_normal(mu, std_or_log_var[, log_var, num_samples])
normal_standard_normal_kl(→ torch.Tensor)
normal_normal_kl(→ torch.Tensor)

Module Contents

timesead.models.common.vae.sample_normal(mu: torch.Tensor, std_or_log_var: torch.Tensor, log_var: bool = False, num_samples: int = 1)

Parameters:

• mu (torch.Tensor)

• std_or_log_var (torch.Tensor)

• log_var (bool)

• num_samples (int)

timesead.models.common.vae.normal_standard_normal_kl(mean: torch.Tensor, std_or_log_var: torch.Tensor, log_var: bool = False) → torch.Tensor

Parameters:

• mean (torch.Tensor)

• std_or_log_var (torch.Tensor)

• log_var (bool)

Return type:

torch.Tensor

timesead.models.common.vae.normal_normal_kl(mean_1: torch.Tensor, std_or_log_var_1: torch.Tensor, mean_2: torch.Tensor, std_or_log_var_2: torch.Tensor, log_var: bool = False) → torch.Tensor

Parameters:

• mean_1 (torch.Tensor)

• std_or_log_var_1 (torch.Tensor)

• mean_2 (torch.Tensor)

• std_or_log_var_2 (torch.Tensor)

• log_var (bool)

Return type:

torch.Tensor

class timesead.models.common.vae.DenseVAEEncoder(input_dim: int, hidden_dims: Sequence[int] = (100, 100), latent_dim: int = 10, activation=torch.nn.ReLU())

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

• input_dim (int)

• hidden_dims (Sequence[int])

• latent_dim (int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

latent_dim = 10

mlp

softplus

forward(x: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Parameters:

x (torch.Tensor)

Return type:

Tuple[torch.Tensor, torch.Tensor]

class timesead.models.common.vae.VAE(encoder: torch.nn.Module, decoder: torch.nn.Module, logvar_out: bool = True)

Bases: torch.nn.Module

VAE Implementation that supports normal distribution with diagonal cov matrix in the latent space and the output

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• encoder (torch.nn.Module)

• decoder (torch.nn.Module)

• logvar_out (bool)

encoder

decoder

log_var = True

forward(x: torch.Tensor, return_latent_sample: bool = False, num_samples: int = 1, force_sample: bool = False) → Tuple[torch.Tensor, Ellipsis]

Parameters:

• x (torch.Tensor)

• return_latent_sample (bool)

• num_samples (int)

• force_sample (bool)

Return type:

Tuple[torch.Tensor, Ellipsis]

class timesead.models.common.vae.VAELoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True)

Bases: timesead.optim.loss.Loss

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

• reduction (str)

• logvar_out (bool)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

logvar_out = True

forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) → torch.Tensor

Parameters:

• predictions (Tuple[torch.Tensor, Ellipsis])

• targets (Tuple[torch.Tensor, Ellipsis])

Return type:

torch.Tensor

class timesead.models.common.vae.NNLLoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True)

Bases: timesead.optim.loss.Loss

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

• reduction (str)

• logvar_out (bool)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

logvar_out = True

forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) → torch.Tensor

Parameters:

• predictions (Tuple[torch.Tensor, Ellipsis])

• targets (Tuple[torch.Tensor, Ellipsis])

Return type:

torch.Tensor

class timesead.models.common.vae.KLLoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True)

Bases: timesead.optim.loss.Loss

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

• reduction (str)

• logvar_out (bool)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

logvar_out = True

forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) → torch.Tensor

Parameters:

• predictions (Tuple[torch.Tensor, Ellipsis])

• targets (Tuple[torch.Tensor, Ellipsis])

Return type:

torch.Tensor

class timesead.models.common.vae.DistVAE(encoder: torch.nn.Module, decoder: torch.nn.Module)

Bases: torch.nn.Module

VAE Implementation that supports arbitrary torch.distributions as long as the latent distribution supports rsample

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• encoder (torch.nn.Module)

• decoder (torch.nn.Module)

encoder

decoder

get_latent_prior(shape: torch.Size)

Parameters:

shape (torch.Size)

sample_prior(shape: torch.Size)

Parameters:

shape (torch.Size)

sample_q(x: torch.Tensor)

Parameters:

x (torch.Tensor)

forward(x: torch.Tensor, return_latent_sample: bool = False, num_samples: int = 1, force_sample: bool = False) → Tuple[torch.distributions.Distribution, Ellipsis]

Parameters:

• x (torch.Tensor)

• return_latent_sample (bool)

• num_samples (int)

• force_sample (bool)

Return type:

Tuple[torch.distributions.Distribution, Ellipsis]

class timesead.models.common.vae.NormalVAEEncoder(*args, **kwargs)

Bases: abc.ABC, torch.nn.Module

Generic class that returns a normal distribution parameterized by mean and standard deviation. Subclasses should override internal_forward to generate these parameters from some input x

Initialize internal Module state, shared by both nn.Module and ScriptModule.

abstract internal_forward(x: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Parameters:

x (torch.Tensor)

Return type:

Tuple[torch.Tensor, torch.Tensor]

forward(x: torch.Tensor) → torch.distributions.Normal

Parameters:

x (torch.Tensor)

Return type:

torch.distributions.Normal

class timesead.models.common.vae.DenseNormalVAEEncoder(input_dim: int, hidden_dims: Sequence[int] = (100, 100), latent_dim: int = 10, activation=torch.nn.ReLU())

Bases: NormalVAEEncoder

Generic class that returns a normal distribution parameterized by mean and standard deviation. Subclasses should override internal_forward to generate these parameters from some input x

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• input_dim (int)

• hidden_dims (Sequence[int])

• latent_dim (int)

latent_dim = 10

mlp

softplus

internal_forward(x: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Parameters:

x (torch.Tensor)

Return type:

Tuple[torch.Tensor, torch.Tensor]

class timesead.models.common.vae.VAEDistLoss(size_average=None, reduce=None, reduction: str = 'mean')

Bases: timesead.optim.loss.Loss

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

reduction (str)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(predictions: Tuple[torch.distributions.Distribution, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) → torch.Tensor

Parameters:

• predictions (Tuple[torch.distributions.Distribution, Ellipsis])

• targets (Tuple[torch.Tensor, Ellipsis])

Return type:

torch.Tensor