Repository navigation
Error in VarAutoEncoder class #8413
Copy link
Copy link
Closed
Description
I'm wondering if there isn't a bug in the implementation of the VarAutoEncoder class.
For the reparameterisation trick (lines 144-150)
MONAI/monai/networks/nets/varautoencoder.py
Lines 28 to 155 in 46a5272
| class VarAutoEncoder(AutoEncoder): | |
| """ | |
| Variational Autoencoder based on the paper - https://arxiv.org/abs/1312.6114 | |
| Args: | |
| spatial_dims: number of spatial dimensions. | |
| in_shape: shape of input data starting with channel dimension. | |
| out_channels: number of output channels. | |
| latent_size: size of the latent variable. | |
| channels: sequence of channels. Top block first. The length of `channels` should be no less than 2. | |
| strides: sequence of convolution strides. The length of `stride` should equal to `len(channels) - 1`. | |
| kernel_size: convolution kernel size, the value(s) should be odd. If sequence, | |
| its length should equal to dimensions. Defaults to 3. | |
| up_kernel_size: upsampling convolution kernel size, the value(s) should be odd. If sequence, | |
| its length should equal to dimensions. Defaults to 3. | |
| num_res_units: number of residual units. Defaults to 0. | |
| inter_channels: sequence of channels defining the blocks in the intermediate layer between encode and decode. | |
| inter_dilations: defines the dilation value for each block of the intermediate layer. Defaults to 1. | |
| num_inter_units: number of residual units for each block of the intermediate layer. Defaults to 0. | |
| act: activation type and arguments. Defaults to PReLU. | |
| norm: feature normalization type and arguments. Defaults to instance norm. | |
| dropout: dropout ratio. Defaults to no dropout. | |
| bias: whether to have a bias term in convolution blocks. Defaults to True. | |
| According to `Performance Tuning Guide <https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html>`_, | |
| if a conv layer is directly followed by a batch norm layer, bias should be False. | |
| use_sigmoid: whether to use the sigmoid function on final output. Defaults to True. | |
| Examples:: | |
| from monai.networks.nets import VarAutoEncoder | |
| # 3 layer network accepting images with dimensions (1, 32, 32) and using a latent vector with 2 values | |
| model = VarAutoEncoder( | |
| spatial_dims=2, | |
| in_shape=(32, 32), # image spatial shape | |
| out_channels=1, | |
| latent_size=2, | |
| channels=(16, 32, 64), | |
| strides=(1, 2, 2), | |
| ) | |
| see also: | |
| - Variational autoencoder network with MedNIST Dataset | |
| https://github.com/Project-MONAI/tutorials/blob/master/modules/varautoencoder_mednist.ipynb | |
| """ | |
| def __init__( | |
| self, | |
| spatial_dims: int, | |
| in_shape: Sequence[int], | |
| out_channels: int, | |
| latent_size: int, | |
| channels: Sequence[int], | |
| strides: Sequence[int], | |
| kernel_size: Sequence[int] | int = 3, | |
| up_kernel_size: Sequence[int] | int = 3, | |
| num_res_units: int = 0, | |
| inter_channels: list | None = None, | |
| inter_dilations: list | None = None, | |
| num_inter_units: int = 2, | |
| act: tuple | str | None = Act.PRELU, | |
| norm: tuple | str = Norm.INSTANCE, | |
| dropout: tuple | str | float | None = None, | |
| bias: bool = True, | |
| use_sigmoid: bool = True, | |
| ) -> None: | |
| self.in_channels, *self.in_shape = in_shape | |
| self.use_sigmoid = use_sigmoid | |
| self.latent_size = latent_size | |
| self.final_size = np.asarray(self.in_shape, dtype=int) | |
| super().__init__( | |
| spatial_dims, | |
| self.in_channels, | |
| out_channels, | |
| channels, | |
| strides, | |
| kernel_size, | |
| up_kernel_size, | |
| num_res_units, | |
| inter_channels, | |
| inter_dilations, | |
| num_inter_units, | |
| act, | |
| norm, | |
| dropout, | |
| bias, | |
| ) | |
| padding = same_padding(self.kernel_size) | |
| for s in strides: | |
| self.final_size = calculate_out_shape(self.final_size, self.kernel_size, s, padding) # type: ignore | |
| linear_size = int(np.prod(self.final_size)) * self.encoded_channels | |
| self.mu = nn.Linear(linear_size, self.latent_size) | |
| self.logvar = nn.Linear(linear_size, self.latent_size) | |
| self.decodeL = nn.Linear(self.latent_size, linear_size) | |
| def encode_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| x = self.encode(x) | |
| x = self.intermediate(x) | |
| x = x.view(x.shape[0], -1) | |
| mu = self.mu(x) | |
| logvar = self.logvar(x) | |
| return mu, logvar | |
| def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Tensor: | |
| x = F.relu(self.decodeL(z)) | |
| x = x.view(x.shape[0], self.channels[-1], *self.final_size) | |
| x = self.decode(x) | |
| if use_sigmoid: | |
| x = torch.sigmoid(x) | |
| return x | |
| def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: | |
| std = torch.exp(0.5 * logvar) | |
| if self.training: # multiply random noise with std only during training | |
| std = torch.randn_like(std).mul(std) | |
| return std.add_(mu) | |
| def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | |
| mu, logvar = self.encode_forward(x) | |
| z = self.reparameterize(mu, logvar) | |
| return self.decode_forward(z, self.use_sigmoid), mu, logvar, z |
def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
std = torch.exp(0.5 * logvar)
if self.training: # multiply random noise with std only during training
std = torch.randn_like(std).mul(std)
return std.add_(mu)
The code is correct for training. At test/inference time, we don't multiple std by a random epsilon, which is correct. However, the function still returns mu+std, which doesn't make sense to me. It should be returning mu only at test time, it makes no sense to take mu+std. Can someone check if I'm correct?
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
No labels
Activity
fix: return mu (not mu+std) in reparameterize at eval time
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(networks): return mu from VAE reparameterize at inference
fix(nets): preserve backward compat for VAE reparameterize at inference
fix(nets): preserve backward compat for VAE reparameterize at inference
fix(nets): preserve backward compat for VAE reparameterize at inference
fix(nets): preserve backward compat for VAE reparameterize at inference
fix(nets): preserve backward compat for VAE reparameterize at inference
Fix VarAutoEncoder reparameterize returning mu+std at inference (#8413)…