cvae_encoder.py

set_seed(13)

z_size = 1
n_filters = 32
in_channels = 1
img_size = 28
input_shape = (in_channels, img_size, img_size)

base_model = nn.Sequential(
    # in_channels@28x28 -> n_filters@28x28
    nn.Conv2d(in_channels, n_filters, kernel_size=3, stride=1, padding=1),
    nn.LeakyReLU(),
            
    # n_filters@28x28 -> (n_filters*2)@14x14
    nn.Conv2d(n_filters, n_filters*2, kernel_size=3, stride=2, padding=1),
    nn.LeakyReLU(),
    
    # (n_filters*2)@14x14 -> (n_filters*2)@7x7
    nn.Conv2d(n_filters*2, n_filters*2, kernel_size=3, stride=2, padding=1),
    nn.LeakyReLU(),
    
    # (n_filters*2)@7x7 -> (n_filters*2)@7x7
    nn.Conv2d(n_filters*2, n_filters*2, kernel_size=3, stride=1, padding=1),
    nn.LeakyReLU(),
    
    # (n_filters*2)@7x7 -> (n_filters*2)*7*7
    nn.Flatten(),
)

encoder_var_cnn = EncoderVar(input_shape, z_size, base_model)