| """Conditioning Augmentation Module""" |
|
|
| from typing import Any |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class CondAugmentation(nn.Module): |
| """Conditioning Augmentation Module""" |
|
|
| def __init__(self, D: int, conditioning_dim: int): |
| """ |
| :param D: Dimension of the text embedding space [D from AttnGAN paper] |
| :param conditioning_dim: Dimension of the conditioning space |
| """ |
| super().__init__() |
| self.cond_dim = conditioning_dim |
| self.cond_augment = nn.Linear(D, conditioning_dim * 4, bias=True) |
| self.glu = nn.GLU(dim=1) |
|
|
| def encode(self, text_embedding: torch.Tensor) -> Any: |
| """ |
| This function encodes the text embedding into the conditioning space |
| :param text_embedding: Text embedding |
| :return: Conditioning embedding |
| """ |
| x_tensor = self.glu(self.cond_augment(text_embedding)) |
| mu_tensor = x_tensor[:, : self.cond_dim] |
| logvar = x_tensor[:, self.cond_dim :] |
| return mu_tensor, logvar |
|
|
| def sample(self, mu_tensor: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: |
| """ |
| This function samples from the Gaussian distribution |
| :param mu: Mean of the Gaussian distribution |
| :param logvar: Log variance of the Gaussian distribution |
| :return: Sample from the Gaussian distribution |
| """ |
| std = torch.exp(0.5 * logvar) |
| eps = torch.randn_like( |
| std |
| ) |
| return mu_tensor + eps * std |
|
|
| def forward(self, text_embedding: torch.Tensor) -> Any: |
| """ |
| This function encodes the text embedding into the conditioning space, |
| and samples from the Gaussian distribution. |
| :param text_embedding: Text embedding |
| :return c_hat: Conditioning embedding (C^ from StackGAN++ paper) |
| :return mu: Mean of the Gaussian distribution |
| :return logvar: Log variance of the Gaussian distribution |
| """ |
| mu_tensor, logvar = self.encode(text_embedding) |
| c_hat = self.sample(mu_tensor, logvar) |
| return c_hat, mu_tensor, logvar |
|
|
