# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
local_rank=args.local_rank,
optimization_level=args.amp_opt_level,
log_dir=work_dir,
create_tb_writer=True,
files_to_copy=[__file__],
)
mnist_data = nemo_simple_gan.MnistGanDataLayer(batch_size=batch_size,
shuffle=True,
train=True,
root=args.train_dataset)
generator = nemo_simple_gan.SimpleGenerator(batch_size)
discriminator = nemo_simple_gan.SimpleDiscriminator()
neg_disc_loss = nemo_simple_gan.DiscriminatorLoss(neg=True)
disc_loss = nemo_simple_gan.DiscriminatorLoss()
disc_grad_penalty = nemo_simple_gan.GradientPenalty(lambda_=3)
interpolater = nemo_simple_gan.InterpolateImage()
# Create generator DAG
latents, real_data, _ = mnist_data()
generated_image = generator(latents=latents)
generator_decision = discriminator(image=generated_image)
generator_loss = neg_disc_loss(decision=generator_decision)
# Create discriminator DAG
interpolated_image = interpolater(image1=real_data, image2=generated_image)
interpolated_decision = discriminator(image=interpolated_image)
# instantiate Neural Factory with supported backend
neural_factory = nemo.core.NeuralModuleFactory(
local_rank=args.local_rank,
optimization_level=args.amp_opt_level,
log_dir=work_dir,
create_tb_writer=True,
files_to_copy=[__file__],
)
mnist_data = nemo_simple_gan.MnistGanDataLayer(batch_size=batch_size,
shuffle=True,
train=True,
root=args.train_dataset)
generator = nemo_simple_gan.SimpleGenerator()
discriminator = nemo_simple_gan.SimpleDiscriminator()
neg_disc_loss = nemo_simple_gan.DiscriminatorLoss(neg=True)
disc_loss = nemo_simple_gan.DiscriminatorLoss()
disc_grad_penalty = nemo_simple_gan.GradientPenalty(lambda_=3)
interpolater = nemo_simple_gan.InterpolateImage()
# Create generator DAG
latents, real_data, _ = mnist_data()
generated_image = generator(latents=latents)
generator_decision = discriminator(image=generated_image)
generator_loss = neg_disc_loss(decision=generator_decision)
# Create discriminator DAG
interpolated_image = interpolater(image1=real_data, image2=generated_image)
interpolated_decision = discriminator(image=interpolated_image)