def model_vae3(zdim, dist):
n_params, obs = observation(dist)
e_layers = [
InputLayer(IMAGE_SHAPE),
CenterAt0(),
Conv(32, 3, 1, name='e1'), # (32, 32, 32)
Conv(32, 4, 2, name='e2'), # (16, 16, 32)
Conv(64, 3, 1, name='e3'), # (16, 16, 64)
Conv(64, 4, 2, name='e4'), # (8, 8, 64)
Conv(128, 3, 1, name='e5'), # (8, 8, 128)
Conv(128, 4, 2, name='e6'), # (4, 4, 128)
Flatten(),
Dense(512, name='encoder_proj'),
]
d_layers = [
InputLayer([zdim]),
Dense(512, name='decoder_proj'),
Reshape([4, 4, 32]),
Tonv(128, 4, 2, name='d6'), # (8, 8, 128)
Conv(128, 3, 1, name='d5'), # (8, 8, 128)
Tonv(64, 4, 2, name='d4'), # (16, 16, 64)
Conv(64, 3, 1, name='d3'), # (16, 16, 64)
Tonv(32, 4, 2, name='d2'), # (32, 32, 32)
Conv(32, 3, 1, name='d1'), # (32, 32, 32)
Conv2D(IMAGE_SHAPE[-1] * n_params, 1, 1)
]
return VariationalAutoencoder(encoder=Sequential(e_layers, 'Encoder'),
decoder=Sequential(d_layers, 'Decoder'),
latents=latents(zdim),
observation=obs)
def model_bdvae2(zdim, dist):
n_params, obs = observation(dist)
e_layers = [
InputLayer(IMAGE_SHAPE),
CenterAt0(),
Conv(32, 3, 1, name='e0'), # (32, 32, 32)
Conv(32, 4, 2, name='e1'), # (16, 16, 32)
Conv(64, 3, 1, name='e2'), # (16, 16, 64)
Conv(64, 4, 2, name='e3'), # (8, 8, 64)
Flatten(),
Dense(512, name='encoder_proj')
]
d_layers = [
InputLayer([zdim]),
Dense(512, name='decoder_proj'),
Reshape([8, 8, 8]),
Tonv(64, 4, 2, name='d3'), # (16, 16, 64)
Conv(64, 3, 1, name='d2'), # (16, 16, 64)
Tonv(32, 4, 2, name='d1'), # (32, 32, 32)
Conv(32, 3, 1, name='d0'), # (32, 32, 32)
Conv2D(IMAGE_SHAPE[-1] * n_params, 1, 1)
]
bi_latents(2, 64, 8, 4, e_layers, d_layers, 'latents1', pre_affine=True)
bi_latents(0, 32, 8, 4, e_layers, d_layers, 'latents2', pre_affine=True)
return HierarchicalVAE(encoder=Sequential(e_layers, 'Encoder'),
decoder=Sequential(d_layers, 'Decoder'),
latents=latents(zdim),
observation=obs)
def model_vae2(zdim, dist):
n_params, obs = observation(dist)
e_layers = [
InputLayer(IMAGE_SHAPE),
CenterAt0(),
Conv(32, 3, 1, name='e0'), # (32, 32, 32)
AveragePooling2D(), # (16, 16)
Conv(32, 3, 1, name='e1'), # (16, 16, 32)
Conv(64, 3, 1, name='e2'), # (16, 16, 64)
AveragePooling2D(), # (8, 8)
Conv(64, 3, 1, name='e3'), # (8, 8, 64)
Flatten(),
Dense(512, name='encoder_proj')
]
d_layers = [
InputLayer([zdim]),
Dense(512, name='decoder_proj'),
Reshape([8, 8, 8]),
Conv(64, 3, 1, name='d3'), # (8, 8, 64)
UpSampling2D(), # (16, 16)
Conv(64, 3, 1, name='d2'), # (16, 16, 64)
Conv(32, 3, 1, name='d1'), # (16, 16, 32)
UpSampling2D(), # (32, 32)
Conv(32, 3, 1, name='d0'), # (32, 32, 32)
Conv2D(IMAGE_SHAPE[-1] * n_params, 1, 1),
]
return VariationalAutoencoder(encoder=Sequential(e_layers, 'Encoder'),
decoder=Sequential(d_layers, 'Decoder'),
latents=latents(zdim),
observation=obs)
def get_dense_networks(args: Arguments):
networks = get_networks('mnist',
is_semi_supervised=False,
is_hierarchical=False,
zdim=args.zdim)
networks['encoder'] = SequentialNetwork([
InputLayer(input_shape=[28, 28, 1]),
CenterAt0(),
Flatten(),
Dense(1024, activation='relu'),
Dense(1024, activation='relu'),
Dense(1024, activation='relu'),
],
name='Encoder')
networks['decoder'] = SequentialNetwork([
InputLayer(input_shape=[args.zdim]),
Dense(1024, activation='relu'),
Dense(1024, activation='relu'),
Dense(1024, activation='relu'),
Dense(28 * 28 * 1, activation='linear'),
Reshape([28, 28, 1]),
],
name='Decoder')
return networks