Class GaussianZ

tokenizer.py:424–466

@si_module
class GaussianZ(nn.Module):
    class Config:
        dim: int
        latent_dim: int
        bias: bool = False
        use_weight_norm: bool = False
        
    def __init__(self, c: Config):
        super().__init__()
        
        self.proj_in = nn.Linear(c.dim, c.latent_dim * 2, bias=c.bias)
        self.proj_out = nn.Linear(c.latent_dim, c.dim, bias=c.bias)

        if c.use_weight_norm:
            self.proj_in = weight_norm(self.proj_in)
            self.proj_out = weight_norm(self.proj_out)

    def reparam(self, mu, logvar):
        std = T.exp(logvar / 2)
        eps = T.randn_like(std)
        return mu + eps * std

    def kl_divergence(self, mu, logvar):
        return T.mean(-0.5 * T.sum(
                1 + logvar - mu.pow(2) - logvar.exp(),
                dim=(1, 2))
            )
        
    def repr_from_latent(self, latent: Union[dict, T.Tensor]):
        if isinstance(latent, T.Tensor):
            z = latent
        else:
            z = self.reparam(latent['mu'], latent['logvar'])
        l = self.proj_out(z)
        return l
    
    def forward(self, x: T.Tensor) -> Tuple[T.Tensor, dict]:
        mu, logvar = self.proj_in(x).chunk(2, dim=-1)
        kl_div = self.kl_divergence(mu, logvar)
        z = self.reparam(mu, logvar)
        xhat = self.proj_out(z)
        latent = {'mu': mu, 'logvar': logvar, 'z': z, 'kl_divergence': kl_div}
        return xhat, latent