GAN_jax module

class GAN_jax.DataLoader(data, batch_size, shuffle=True)

Bases: object

class GAN_jax.GAN(divergence, generator, gen_optimizer, noise_source, epochs, disc_steps_per_gen_step, batch_size=None, reverse_order=False, include_penalty_in_gen_loss=False, cnn=False)

Bases: object

Class for training a GAN using one of the provided divergences If reverse_order=False the GAN works to minimize min_theta D(P||g_theta(Z)) where P is the distribution to be leared, Z is the noise source and g_theta is the generator (with parameters theta). If reverse_order=True the GAN works to minimize min_theta D(g_theta(Z)||P) where P is the distribution to be leared, Z is the noise source and g_theta is the generator (with parameters theta).

disc_train_step(x, z, disc_params, disc_opt_state)

discriminator’s parameters update

estimate_loss(x, z, gen_params)

Estimating the loss

gen_train_step(x, z, gen_params, gen_opt_state)

generator’s parameters update

generate_samples(num_samples, gen_params)

train(data_P, disc_params, gen_params, disc_opt_state, gen_opt_state, save_frequency=None, num_gen_samples_to_save=None, save_loss_estimates=False)

training function of our GAN