def test_client_and_server_computations(self): train_generator_fn, train_discriminator_fn = ( _get_train_generator_and_discriminator_fns()) # N.B. The way we are using datasets and re-using the same # generator and discriminator doesn't really "make sense" from an ML # perspective, but it's sufficient for testing. For more proper usage of # these functions, see training_loops.py. generator = one_dim_gan.create_generator() discriminator = one_dim_gan.create_discriminator() gen_inputs = one_dim_gan.create_generator_inputs() real_data = one_dim_gan.create_real_data() server_state = gan_training_tf_fns.server_initial_state( generator, discriminator, INIT_DP_AVERAGING_STATE) # DP averaging aggregation state is initialized properly in # server_initial_state(). self.assertEqual(server_state.dp_averaging_state, INIT_DP_AVERAGING_STATE) client_output = gan_training_tf_fns.client_computation( gen_inputs.take(3), real_data.take(3), gan_training_tf_fns.FromServer( generator_weights=server_state.generator_weights, discriminator_weights=server_state.discriminator_weights), generator, discriminator, train_discriminator_fn) server_disc_update_optimizer = tf.keras.optimizers.Adam() for _ in range(2): # Train for 2 rounds server_state = gan_training_tf_fns.server_computation( server_state, gen_inputs.take(3), client_output, generator, discriminator, server_disc_update_optimizer, train_generator_fn, NEW_DP_AVERAGING_STATE) counters = self.evaluate(server_state.counters) self.assertDictEqual( counters, { 'num_rounds': 2, 'num_discriminator_train_examples': 2 * 3 * one_dim_gan.BATCH_SIZE, 'num_generator_train_examples': 2 * 3 * one_dim_gan.BATCH_SIZE }) # DP averaging aggregation state updates properly in server_computation(). self.assertEqual(server_state.dp_averaging_state, NEW_DP_AVERAGING_STATE)
def server_computation(server_state, gen_inputs, client_output, new_dp_averaging_state): """The wrapped server_computation.""" # initialize the optimizers beforehand so you don't create them within the tf.function steps = server_state.counters['num_rounds'] scheduler = tf.keras.optimizers.schedules.PiecewiseConstantDecay( [1000], [0.001, 0.0005]) state_gen_optimizer = gan.state_gen_optimizer_fn( scheduler.__call__(steps)) generator = gan.generator_model_fn() gan_training_tf_fns.initialize_optimizer_vars(generator, state_gen_optimizer) discriminator = gan.discriminator_model_fn(0.0002) state_disc_optimizer = gan.state_disc_optimizer_fn(steps) gan_training_tf_fns.initialize_optimizer_vars(discriminator, state_disc_optimizer) if gan.disc_status == 'fedadam': return gan_training_tf_fns.server_computation_fedadam( server_state=server_state, gen_inputs_ds=gen_inputs, client_output=client_output, generator=generator, discriminator=discriminator, state_disc_optimizer=state_disc_optimizer, state_gen_optimizer=state_gen_optimizer, new_dp_averaging_state=new_dp_averaging_state) else: return gan_training_tf_fns.server_computation( server_state=server_state, gen_inputs_ds=gen_inputs, client_output=client_output, generator=generator, discriminator=discriminator, state_disc_optimizer=state_disc_optimizer, state_gen_optimizer=state_gen_optimizer, new_dp_averaging_state=new_dp_averaging_state)