timesead.models.common.vae ========================== .. py:module:: timesead.models.common.vae Classes ------- .. autoapisummary:: timesead.models.common.vae.DenseVAEEncoder timesead.models.common.vae.VAE timesead.models.common.vae.VAELoss timesead.models.common.vae.NNLLoss timesead.models.common.vae.KLLoss timesead.models.common.vae.DistVAE timesead.models.common.vae.NormalVAEEncoder timesead.models.common.vae.DenseNormalVAEEncoder timesead.models.common.vae.VAEDistLoss Functions --------- .. autoapisummary:: timesead.models.common.vae.sample_normal timesead.models.common.vae.normal_standard_normal_kl timesead.models.common.vae.normal_normal_kl Module Contents --------------- .. py:function:: sample_normal(mu: torch.Tensor, std_or_log_var: torch.Tensor, log_var: bool = False, num_samples: int = 1) .. py:function:: normal_standard_normal_kl(mean: torch.Tensor, std_or_log_var: torch.Tensor, log_var: bool = False) -> torch.Tensor .. py:function:: 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 .. py:class:: DenseVAEEncoder(input_dim: int, hidden_dims: Sequence[int] = (100, 100), latent_dim: int = 10, activation=torch.nn.ReLU()) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: latent_dim :value: 10 .. py:attribute:: mlp .. py:attribute:: softplus .. py:method:: forward(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] .. py:class:: VAE(encoder: torch.nn.Module, decoder: torch.nn.Module, logvar_out: bool = True) Bases: :py:obj:`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. .. py:attribute:: encoder .. py:attribute:: decoder .. py:attribute:: log_var :value: True .. py:method:: forward(x: torch.Tensor, return_latent_sample: bool = False, num_samples: int = 1, force_sample: bool = False) -> Tuple[torch.Tensor, Ellipsis] .. py:class:: VAELoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: logvar_out :value: True .. py:method:: forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) -> torch.Tensor .. py:class:: NNLLoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: logvar_out :value: True .. py:method:: forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) -> torch.Tensor .. py:class:: KLLoss(size_average=None, reduce=None, reduction: str = 'mean', logvar_out: bool = True) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: logvar_out :value: True .. py:method:: forward(predictions: Tuple[torch.Tensor, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) -> torch.Tensor .. py:class:: DistVAE(encoder: torch.nn.Module, decoder: torch.nn.Module) Bases: :py:obj:`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. .. py:attribute:: encoder .. py:attribute:: decoder .. py:method:: get_latent_prior(shape: torch.Size) .. py:method:: sample_prior(shape: torch.Size) .. py:method:: sample_q(x: torch.Tensor) .. py:method:: forward(x: torch.Tensor, return_latent_sample: bool = False, num_samples: int = 1, force_sample: bool = False) -> Tuple[torch.distributions.Distribution, Ellipsis] .. py:class:: NormalVAEEncoder(*args, **kwargs) Bases: :py:obj:`abc.ABC`, :py:obj:`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. .. py:method:: internal_forward(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] :abstractmethod: .. py:method:: forward(x: torch.Tensor) -> torch.distributions.Normal .. py:class:: DenseNormalVAEEncoder(input_dim: int, hidden_dims: Sequence[int] = (100, 100), latent_dim: int = 10, activation=torch.nn.ReLU()) Bases: :py:obj:`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. .. py:attribute:: latent_dim :value: 10 .. py:attribute:: mlp .. py:attribute:: softplus .. py:method:: internal_forward(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] .. py:class:: VAEDistLoss(size_average=None, reduce=None, reduction: str = 'mean') Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:method:: forward(predictions: Tuple[torch.distributions.Distribution, Ellipsis], targets: Tuple[torch.Tensor, Ellipsis], *args, **kwargs) -> torch.Tensor