def test_svi_vae():
# define the encoder and decoder
encoder = Encoder()
decoder = Decoder()
# define the vae model
vae = VAE(encoder, decoder, hidden_size=400, latent_size=20)
# define the loss function
net_loss = ELBO(latent_prior='Normal', output_prior='Normal')
# define the optimizer
optimizer = nn.Adam(params=vae.trainable_params(), learning_rate=0.001)
# define the training dataset
ds_train = create_dataset(image_path, 128, 1)
net_with_loss = nn.WithLossCell(vae, net_loss)
# define the variational inference
vi = SVI(net_with_loss=net_with_loss, optimizer=optimizer)
# run the vi to return the trained network.
vae = vi.run(train_dataset=ds_train, epochs=5)
# get the trained loss
trained_loss = vi.get_train_loss()
# test function: generate_sample
generated_sample = vae.generate_sample(64, IMAGE_SHAPE)
# test function: reconstruct_sample
for sample in ds_train.create_dict_iterator():
sample_x = Tensor(sample['image'], dtype=mstype.float32)
reconstructed_sample = vae.reconstruct_sample(sample_x)
print('The loss of the trained network is ', trained_loss)
print('The hape of the generated sample is ', generated_sample.shape)
print('The shape of the reconstructed sample is ',
reconstructed_sample.shape)
def test_svi_cvae():
# define the encoder and decoder
encoder = Encoder(num_classes=10)
decoder = Decoder()
# define the cvae model
cvae = ConditionalVAE(encoder,
decoder,
hidden_size=400,
latent_size=20,
num_classes=10)
# define the loss function
net_loss = ELBO(latent_prior='Normal', output_prior='Normal')
# define the optimizer
optimizer = nn.Adam(params=cvae.trainable_params(), learning_rate=0.001)
# define the training dataset
ds_train = create_dataset(image_path, 128, 1)
# define the WithLossCell modified
net_with_loss = CVAEWithLossCell(cvae, net_loss)
# define the variational inference
vi = SVI(net_with_loss=net_with_loss, optimizer=optimizer)
# run the vi to return the trained network.
cvae = vi.run(train_dataset=ds_train, epochs=5)
# get the trained loss
trained_loss = vi.get_train_loss()
# test function: generate_sample
sample_label = Tensor([i for i in range(0, 8)] * 8, dtype=mstype.int32)
generated_sample = cvae.generate_sample(sample_label, 64, IMAGE_SHAPE)
# test function: reconstruct_sample
for sample in ds_train.create_dict_iterator(output_numpy=True,
num_epochs=1):
sample_x = Tensor(sample['image'], dtype=mstype.float32)
sample_y = Tensor(sample['label'], dtype=mstype.int32)
reconstructed_sample = cvae.reconstruct_sample(sample_x, sample_y)
print('The loss of the trained network is ', trained_loss)
print('The shape of the generated sample is ', generated_sample.shape)
print('The shape of the reconstructed sample is ',
reconstructed_sample.shape)
return mnist_ds
if __name__ == "__main__":
# define the encoder and decoder
encoder = Encoder(num_classes=10)
decoder = Decoder()
# define the cvae model
cvae = ConditionalVAE(encoder,
decoder,
hidden_size=400,
latent_size=20,
num_classes=10)
# define the loss function
net_loss = ELBO(latent_prior='Normal', output_prior='Normal')
# define the optimizer
optimizer = nn.Adam(params=cvae.trainable_params(), learning_rate=0.001)
# define the training dataset
ds_train = create_dataset(image_path, 128, 1)
# define the WithLossCell modified
net_with_loss = WithLossCell(cvae, net_loss)
# define the variational inference
vi = SVI(net_with_loss=net_with_loss, optimizer=optimizer)
# run the vi to return the trained network.
cvae = vi.run(train_dataset=ds_train, epochs=10)
# get the trained loss
trained_loss = vi.get_train_loss()
# test function: generate_sample
sample_label = Tensor([i for i in range(0, 8)] * 8, dtype=mstype.int32)
generated_sample = cvae.generate_sample(sample_label, 64, IMAGE_SHAPE)
def __init__(self):
super(VaeGanLoss, self).__init__()
self.zeros = P.ZerosLike()
self.mse = nn.MSELoss(reduction='sum')
self.elbo = ELBO(latent_prior='Normal', output_prior='Normal')