The code for the model, inference, and training was implemented using PyTorch [1]}, Catalyst [2]}, and NumPy [3]}. The VAE uses a convolutional encoder and decoder pair, each of the layers uses a 3-voxel kernel, a stride of 2 and 1-voxel padding. The channel sizes in the encoder are \(1\rightarrow 64, 64\rightarrow 128, 128\rightarrow 256, 256\rightarrow 512\)  and \(512\rightarrow 256, 256\rightarrow 128, 128\rightarrow 64, 64\rightarrow 32, 32\rightarrow 16, 16\rightarrow 1\)  in the decoder, the last layer in the decoder uses stride 1 and no bias parameters. Each convolutional layer uses a ReLU [4]} as its activation function, except for the last layer in the decoder, which uses a hyperbolic tangent function to map the output between [-1, 1] to match the input range. The last convolutional layer in the encoder produces an output with shape: \(4\text{-by-}4\text{-by-}4\)  and 256 channels, this output is flattened and mapped to the mean \(\mu \)  and variance \(\sigma \) , which are used to construct a diagonal multivariate Gaussian from which \(z\)  is sampled. To make sure the VAE is fully differentiable, we use the reparameterization trick to train it [5]}. The classification evaluations in the latent space are implemented using RAPIDS AI [6]} to make sure highly parallelizable computations are performed on the GPU and to minimize costly CPU\(\rightarrow \) GPU and GPU\(\rightarrow \) CPU data transfers. The experiments were performed on an NVIDIA DGX-1 V100.
