def __init__( self, in_channels=1, out_channels=1, latent_channels=16, strides=[4], hidden_channels=64, residual_channels=32, ): """Initializes a new VAE instance. Args: in_channels: Number of input channels. out_channels: Number of output channels. latent_channels: Number of channels for each latent variable. strides: List of encoder/decoder strides. For each stride, we create an encoder (decoder) which downsamples (upsamples) the input by the stride. hidden_channels: Number of channels in (non residual block) hidden layers. residual_channels: Number of hidden channels in residual blocks. """ super().__init__() self._latent_channels = latent_channels self._total_stride = sum(strides) encoder = [] for i, stride in enumerate(strides): in_c = in_channels if i == 0 else hidden_channels out_c = ( hidden_channels if i < len(strides) - 1 else 2 * self._latent_channels ) encoder.append( vaes.Encoder( in_channels=in_c, out_channels=out_c, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=2, stride=stride, ) ) self._encoder = nn.Sequential(*encoder) decoder = [] for i, stride in enumerate(reversed(strides)): in_c = self._latent_channels if i == 0 else hidden_channels out_c = hidden_channels if i < len(strides) - 1 else out_channels decoder.append( vaes.Decoder( in_channels=in_c, out_channels=out_c, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=2, stride=stride, ) ) self._decoder = nn.Sequential(*decoder)
def __init__( self, in_channels=1, out_channels=1, in_size=28, latent_dim=10, hidden_channels=32, n_residual_blocks=2, residual_channels=16, ): """Initializes a new VAE instance. Args: in_channels: Number of input channels. out_channels: Number of output channels. in_size: Size of the input images. Used to create bottleneck layers. latent_dim: The dimensionality of each latent variable. hidden_channels: Number of channels in (non residual block) hidden layers. n_residual_blocks: Number of residual blocks in each residual stack. residual_channels: Number of hidden channels in residual blocks. """ super().__init__() self._latent_dim = latent_dim self._encoder_out_dims = (hidden_channels, in_size // 4, in_size // 4) self._encoder = vaes.Encoder( in_channels=in_channels, out_channels=hidden_channels, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=n_residual_blocks, stride=4, ) flat_dim = hidden_channels * in_size // 4 * in_size // 4 self._mean = nn.Linear(flat_dim, self._latent_dim) self._log_var = nn.Linear(flat_dim, self._latent_dim) self._bottleneck = nn.Linear(self._latent_dim, flat_dim) self._decoder_ = vaes.Decoder( in_channels=hidden_channels, out_channels=out_channels, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=n_residual_blocks, stride=4, )
def __init__( self, in_channels=1, out_channels=1, hidden_channels=128, n_residual_blocks=2, residual_channels=32, n_embeddings=128, embedding_dim=16, ): """Initializes a new VQVAE instance. Args: in_channels: Number of input channels. out_channels: Number of output channels. hidden_channels: Number of channels in (non residual block) hidden layers. n_residual_blocks: Number of residual blocks in each residual stack. residual_channels: Number of hidden channels in residual blocks. n_embeddings: Number of VectorQuantizer embeddings. embedding_dim: Dimension of the VectorQuantizer embeddings. """ super().__init__() self._encoder = vaes.Encoder( in_channels=in_channels, out_channels=hidden_channels, hidden_channels=hidden_channels, n_residual_blocks=n_residual_blocks, residual_channels=residual_channels, stride=4, ) self._quantizer = vaes.Quantizer( in_channels=hidden_channels, n_embeddings=n_embeddings, embedding_dim=embedding_dim, ) self._decoder = vaes.Decoder( in_channels=embedding_dim, out_channels=out_channels, hidden_channels=hidden_channels, n_residual_blocks=n_residual_blocks, residual_channels=residual_channels, stride=4, )
def __init__( self, in_channels=1, out_channels=1, latent_channels=2, hidden_channels=128, residual_channels=32, ): """Initializes a new VAE instance. Args: in_channels: Number of input channels. out_channels: Number of output channels. latent_channels: Number of channels for each latent variable. hidden_channels: Number of channels in (non residual block) hidden layers. residual_channels: Number of hidden channels in residual blocks. """ super().__init__() self._latent_channels = latent_channels self._stride = 4 self._encoder = vaes.Encoder( in_channels=in_channels, out_channels=2 * self._latent_channels, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=2, stride=self._stride, ) self._decoder = vaes.Decoder( in_channels=self._latent_channels, out_channels=out_channels, hidden_channels=hidden_channels, residual_channels=residual_channels, n_residual_blocks=2, stride=self._stride, )