Example #1
0
    def __init__(self, stage_id, batch_size, config):
        """Build graph stage from config dictionary.

    Stage_id and batch_size change during training so they are kept separate
    from the global config. This function is also called by 'load_from_path()'.

    Args:
      stage_id: (int) Build generator/discriminator with this many stages.
      batch_size: (int) Build graph with fixed batch size.
      config: (dict) All the global state.
    """
        data_helper = data_helpers.registry[config['data_type']](config)
        real_images, real_one_hot_labels = data_helper.provide_data(batch_size)

        # gen_one_hot_labels = real_one_hot_labels
        gen_one_hot_labels = data_helper.provide_one_hot_labels(batch_size)
        num_tokens = real_one_hot_labels.shape[1].value

        current_image_id = tf.train.get_or_create_global_step()
        current_image_id_inc_op = current_image_id.assign_add(batch_size)
        tf.summary.scalar('current_image_id', current_image_id)

        train_time = tf.Variable(0., dtype=tf.float32, trainable=False)
        tf.summary.scalar('train_time', train_time)

        resolution_schedule = train_util.make_resolution_schedule(**config)
        num_blocks, num_images = train_util.get_stage_info(stage_id, **config)

        num_stages = (2 * config['num_resolutions']) - 1
        if config['train_time_limit'] is not None:
            stage_times = np.zeros(num_stages, dtype='float32')
            stage_times[0] = 1.
            for i in range(1, num_stages):
                stage_times[i] = (stage_times[i - 1] *
                                  config['train_time_stage_multiplier'])
            stage_times *= config['train_time_limit'] / np.sum(stage_times)
            stage_times = np.cumsum(stage_times)
            print('Stage times:')
            for t in stage_times:
                print('\t{}'.format(t))

        if config['train_progressive']:
            if config['train_time_limit'] is not None:
                progress = networks.compute_progress_from_time(
                    train_time, config['num_resolutions'], num_blocks,
                    stage_times)
            else:
                progress = networks.compute_progress(
                    current_image_id, config['stable_stage_num_images'],
                    config['transition_stage_num_images'], num_blocks)
        else:
            progress = num_blocks - 1.  # Maximum value, must be float.
            num_images = 0
            for stage_id_idx in train_util.get_stage_ids(**config):
                _, n = train_util.get_stage_info(stage_id_idx, **config)
                num_images += n

        # Add to config
        config['resolution_schedule'] = resolution_schedule
        config['num_blocks'] = num_blocks
        config['num_images'] = num_images
        config['progress'] = progress
        config['num_tokens'] = num_tokens
        tf.summary.scalar('progress', progress)

        real_images = networks.blend_images(real_images,
                                            progress,
                                            resolution_schedule,
                                            num_blocks=num_blocks)

        ########## Define model.
        noises = train_util.make_latent_vectors(batch_size, **config)

        # Get network functions and wrap with hparams
        g_fn = lambda x: net_fns.g_fn_registry[config['g_fn']](x, **config)
        d_fn = lambda x: net_fns.d_fn_registry[config['d_fn']](x, **config)

        # Extra lambda functions to conform to tfgan.gan_model interface
        gan_model = tfgan.gan_model(
            generator_fn=lambda inputs: g_fn(inputs)[0],
            discriminator_fn=lambda images, unused_cond: d_fn(images)[0],
            real_data=real_images,
            generator_inputs=(noises, gen_one_hot_labels))

        ########## Define loss.
        gan_loss = train_util.define_loss(gan_model, **config)

        ########## Auxiliary loss functions
        def _compute_ac_loss(images, target_one_hot_labels):
            with tf.variable_scope(gan_model.discriminator_scope, reuse=True):
                _, end_points = d_fn(images)
            return tf.reduce_mean(
                tf.nn.softmax_cross_entropy_with_logits_v2(
                    labels=tf.stop_gradient(target_one_hot_labels),
                    logits=end_points['classification_logits']))

        def _compute_gl_consistency_loss(data):
            """G&L consistency loss."""
            sh = data_helper.specgrams_helper
            is_mel = isinstance(data_helper, data_helpers.DataMelHelper)
            if is_mel:
                stfts = sh.melspecgrams_to_stfts(data)
            else:
                stfts = sh.specgrams_to_stfts(data)
            waves = sh.stfts_to_waves(stfts)
            new_stfts = sh.waves_to_stfts(waves)
            # Magnitude loss
            mag = tf.abs(stfts)
            new_mag = tf.abs(new_stfts)
            # Normalize loss to max
            get_max = lambda x: tf.reduce_max(x, axis=(1, 2), keepdims=True)
            mag_max = get_max(mag)
            new_mag_max = get_max(new_mag)
            mag_scale = tf.maximum(1.0, tf.maximum(mag_max, new_mag_max))
            mag_diff = (mag - new_mag) / mag_scale
            mag_loss = tf.reduce_mean(tf.square(mag_diff))
            return mag_loss

        with tf.name_scope('losses'):
            # Loss weights
            gen_ac_loss_weight = config['generator_ac_loss_weight']
            dis_ac_loss_weight = config['discriminator_ac_loss_weight']
            gen_gl_consistency_loss_weight = config[
                'gen_gl_consistency_loss_weight']

            # AC losses.
            fake_ac_loss = _compute_ac_loss(gan_model.generated_data,
                                            gen_one_hot_labels)
            real_ac_loss = _compute_ac_loss(gan_model.real_data,
                                            real_one_hot_labels)

            # GL losses.
            is_fourier = isinstance(data_helper,
                                    (data_helpers.DataSTFTHelper,
                                     data_helpers.DataSTFTNoIFreqHelper,
                                     data_helpers.DataMelHelper))
            if isinstance(data_helper, data_helpers.DataWaveHelper):
                is_fourier = False

            if is_fourier:
                fake_gl_loss = _compute_gl_consistency_loss(
                    gan_model.generated_data)
                real_gl_loss = _compute_gl_consistency_loss(
                    gan_model.real_data)

            # Total losses.
            wx_fake_ac_loss = gen_ac_loss_weight * fake_ac_loss
            wx_real_ac_loss = dis_ac_loss_weight * real_ac_loss
            wx_fake_gl_loss = 0.0
            if (is_fourier and gen_gl_consistency_loss_weight > 0 and stage_id
                    == train_util.get_total_num_stages(**config) - 1):
                wx_fake_gl_loss = fake_gl_loss * gen_gl_consistency_loss_weight
            # Update the loss functions
            gan_loss = gan_loss._replace(
                generator_loss=(gan_loss.generator_loss + wx_fake_ac_loss +
                                wx_fake_gl_loss),
                discriminator_loss=(gan_loss.discriminator_loss +
                                    wx_real_ac_loss))

            tf.summary.scalar('fake_ac_loss', fake_ac_loss)
            tf.summary.scalar('real_ac_loss', real_ac_loss)
            tf.summary.scalar('wx_fake_ac_loss', wx_fake_ac_loss)
            tf.summary.scalar('wx_real_ac_loss', wx_real_ac_loss)
            tf.summary.scalar('total_gen_loss', gan_loss.generator_loss)
            tf.summary.scalar('total_dis_loss', gan_loss.discriminator_loss)

            if is_fourier:
                tf.summary.scalar('fake_gl_loss', fake_gl_loss)
                tf.summary.scalar('real_gl_loss', real_gl_loss)
                tf.summary.scalar('wx_fake_gl_loss', wx_fake_gl_loss)

        ########## Define train ops.
        gan_train_ops, optimizer_var_list = train_util.define_train_ops(
            gan_model, gan_loss, **config)
        gan_train_ops = gan_train_ops._replace(
            global_step_inc_op=current_image_id_inc_op)

        ########## Generator smoothing.
        generator_ema = tf.train.ExponentialMovingAverage(decay=0.999)
        gan_train_ops, generator_vars_to_restore = \
            train_util.add_generator_smoothing_ops(generator_ema,
                                                   gan_model,
                                                   gan_train_ops)
        load_scope = tf.variable_scope(
            gan_model.generator_scope,
            reuse=True,
            custom_getter=train_util.make_var_scope_custom_getter_for_ema(
                generator_ema))

        ########## Separate path for generating samples with a placeholder (ph)
        # Mapping of pitches to one-hot labels
        pitch_counts = data_helper.get_pitch_counts()
        pitch_to_label_dict = {}
        for i, pitch in enumerate(sorted(pitch_counts.keys())):
            pitch_to_label_dict[pitch] = i

        # (label_ph, noise_ph) -> fake_wave_ph
        labels_ph = tf.placeholder(tf.int32, [batch_size])
        noises_ph = tf.placeholder(tf.float32,
                                   [batch_size, config['latent_vector_size']])
        num_pitches = len(pitch_counts)
        one_hot_labels_ph = tf.one_hot(labels_ph, num_pitches)
        with load_scope:
            fake_data_ph, _ = g_fn((noises_ph, one_hot_labels_ph))
            fake_waves_ph = data_helper.data_to_waves(fake_data_ph)

        if config['train_time_limit'] is not None:
            stage_train_time_limit = stage_times[stage_id]
            #  config['train_time_limit'] * \
            # (float(stage_id+1) / ((2*config['num_resolutions'])-1))
        else:
            stage_train_time_limit = None

        ########## Add variables as properties
        self.stage_id = stage_id
        self.batch_size = batch_size
        self.config = config
        self.data_helper = data_helper
        self.resolution_schedule = resolution_schedule
        self.num_images = num_images
        self.num_blocks = num_blocks
        self.current_image_id = current_image_id
        self.progress = progress
        self.generator_fn = g_fn
        self.discriminator_fn = d_fn
        self.gan_model = gan_model
        self.fake_ac_loss = fake_ac_loss
        self.real_ac_loss = real_ac_loss
        self.gan_loss = gan_loss
        self.gan_train_ops = gan_train_ops
        self.optimizer_var_list = optimizer_var_list
        self.generator_ema = generator_ema
        self.generator_vars_to_restore = generator_vars_to_restore
        self.real_images = real_images
        self.real_one_hot_labels = real_one_hot_labels
        self.load_scope = load_scope
        self.pitch_counts = pitch_counts
        self.pitch_to_label_dict = pitch_to_label_dict
        self.labels_ph = labels_ph
        self.noises_ph = noises_ph
        self.fake_waves_ph = fake_waves_ph
        self.saver = tf.train.Saver()
        self.sess = tf.Session()
        self.train_time = train_time
        self.stage_train_time_limit = stage_train_time_limit
Example #2
0
  def __init__(self, stage_id, batch_size, config):
    """Build graph stage from config dictionary.

    Stage_id and batch_size change during training so they are kept separate
    from the global config. This function is also called by 'load_from_path()'.

    Args:
      stage_id: (int) Build generator/discriminator with this many stages.
      batch_size: (int) Build graph with fixed batch size.
      config: (dict) All the global state.
    """
    data_helper = data_helpers.registry[config['data_type']](config)
    real_images, real_one_hot_labels = data_helper.provide_data(batch_size)

    # gen_one_hot_labels = real_one_hot_labels
    gen_one_hot_labels = data_helper.provide_one_hot_labels(batch_size)
    num_tokens = real_one_hot_labels.shape[1].value

    current_image_id = tf.train.get_or_create_global_step()
    current_image_id_inc_op = current_image_id.assign_add(batch_size)
    tf.summary.scalar('current_image_id', current_image_id)

    train_time = tf.Variable(0., dtype=tf.float32, trainable=False)
    tf.summary.scalar('train_time', train_time)

    resolution_schedule = train_util.make_resolution_schedule(**config)
    num_blocks, num_images = train_util.get_stage_info(stage_id, **config)

    num_stages = (2*config['num_resolutions']) - 1
    if config['train_time_limit'] is not None:
      stage_times = np.zeros(num_stages, dtype='float32')
      stage_times[0] = 1.
      for i in range(1, num_stages):
        stage_times[i] = (stage_times[i-1] *
                          config['train_time_stage_multiplier'])
      stage_times *= config['train_time_limit'] / np.sum(stage_times)
      stage_times = np.cumsum(stage_times)
      print('Stage times:')
      for t in stage_times:
        print('\t{}'.format(t))

    if config['train_progressive']:
      if config['train_time_limit'] is not None:
        progress = networks.compute_progress_from_time(
            train_time, config['num_resolutions'], num_blocks, stage_times)
      else:
        progress = networks.compute_progress(
            current_image_id, config['stable_stage_num_images'],
            config['transition_stage_num_images'], num_blocks)
    else:
      progress = num_blocks - 1.  # Maximum value, must be float.
      num_images = 0
      for stage_id_idx in train_util.get_stage_ids(**config):
        _, n = train_util.get_stage_info(stage_id_idx, **config)
        num_images += n

    # Add to config
    config['resolution_schedule'] = resolution_schedule
    config['num_blocks'] = num_blocks
    config['num_images'] = num_images
    config['progress'] = progress
    config['num_tokens'] = num_tokens
    tf.summary.scalar('progress', progress)

    real_images = networks.blend_images(
        real_images, progress, resolution_schedule, num_blocks=num_blocks)

    ########## Define model.
    noises = train_util.make_latent_vectors(batch_size, **config)

    # Get network functions and wrap with hparams
    g_fn = lambda x: net_fns.g_fn_registry[config['g_fn']](x, **config)
    d_fn = lambda x: net_fns.d_fn_registry[config['d_fn']](x, **config)

    # Extra lambda functions to conform to tfgan.gan_model interface
    gan_model = tfgan.gan_model(
        generator_fn=lambda inputs: g_fn(inputs)[0],
        discriminator_fn=lambda images, unused_cond: d_fn(images)[0],
        real_data=real_images,
        generator_inputs=(noises, gen_one_hot_labels))

    ########## Define loss.
    gan_loss = train_util.define_loss(gan_model, **config)

    ########## Auxiliary loss functions
    def _compute_ac_loss(images, target_one_hot_labels):
      with tf.variable_scope(gan_model.discriminator_scope, reuse=True):
        _, end_points = d_fn(images)
      return tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits_v2(
              labels=tf.stop_gradient(target_one_hot_labels),
              logits=end_points['classification_logits']))

    def _compute_gl_consistency_loss(data):
      """G&L consistency loss."""
      sh = data_helper.specgrams_helper
      is_mel = isinstance(data_helper, data_helpers.DataMelHelper)
      if is_mel:
        stfts = sh.melspecgrams_to_stfts(data)
      else:
        stfts = sh.specgrams_to_stfts(data)
      waves = sh.stfts_to_waves(stfts)
      new_stfts = sh.waves_to_stfts(waves)
      # Magnitude loss
      mag = tf.abs(stfts)
      new_mag = tf.abs(new_stfts)
      # Normalize loss to max
      get_max = lambda x: tf.reduce_max(x, axis=(1, 2), keepdims=True)
      mag_max = get_max(mag)
      new_mag_max = get_max(new_mag)
      mag_scale = tf.maximum(1.0, tf.maximum(mag_max, new_mag_max))
      mag_diff = (mag - new_mag) / mag_scale
      mag_loss = tf.reduce_mean(tf.square(mag_diff))
      return mag_loss

    with tf.name_scope('losses'):
      # Loss weights
      gen_ac_loss_weight = config['generator_ac_loss_weight']
      dis_ac_loss_weight = config['discriminator_ac_loss_weight']
      gen_gl_consistency_loss_weight = config['gen_gl_consistency_loss_weight']

      # AC losses.
      fake_ac_loss = _compute_ac_loss(gan_model.generated_data,
                                      gen_one_hot_labels)
      real_ac_loss = _compute_ac_loss(gan_model.real_data, real_one_hot_labels)

      # GL losses.
      is_fourier = isinstance(data_helper, (data_helpers.DataSTFTHelper,
                                            data_helpers.DataSTFTNoIFreqHelper,
                                            data_helpers.DataMelHelper))
      if isinstance(data_helper, data_helpers.DataWaveHelper):
        is_fourier = False

      if is_fourier:
        fake_gl_loss = _compute_gl_consistency_loss(gan_model.generated_data)
        real_gl_loss = _compute_gl_consistency_loss(gan_model.real_data)

      # Total losses.
      wx_fake_ac_loss = gen_ac_loss_weight * fake_ac_loss
      wx_real_ac_loss = dis_ac_loss_weight * real_ac_loss
      wx_fake_gl_loss = 0.0
      if (is_fourier and
          gen_gl_consistency_loss_weight > 0 and
          stage_id == train_util.get_total_num_stages(**config) - 1):
        wx_fake_gl_loss = fake_gl_loss * gen_gl_consistency_loss_weight
      # Update the loss functions
      gan_loss = gan_loss._replace(
          generator_loss=(
              gan_loss.generator_loss + wx_fake_ac_loss + wx_fake_gl_loss),
          discriminator_loss=(gan_loss.discriminator_loss + wx_real_ac_loss))

      tf.summary.scalar('fake_ac_loss', fake_ac_loss)
      tf.summary.scalar('real_ac_loss', real_ac_loss)
      tf.summary.scalar('wx_fake_ac_loss', wx_fake_ac_loss)
      tf.summary.scalar('wx_real_ac_loss', wx_real_ac_loss)
      tf.summary.scalar('total_gen_loss', gan_loss.generator_loss)
      tf.summary.scalar('total_dis_loss', gan_loss.discriminator_loss)

      if is_fourier:
        tf.summary.scalar('fake_gl_loss', fake_gl_loss)
        tf.summary.scalar('real_gl_loss', real_gl_loss)
        tf.summary.scalar('wx_fake_gl_loss', wx_fake_gl_loss)

    ########## Define train ops.
    gan_train_ops, optimizer_var_list = train_util.define_train_ops(
        gan_model, gan_loss, **config)
    gan_train_ops = gan_train_ops._replace(
        global_step_inc_op=current_image_id_inc_op)

    ########## Generator smoothing.
    generator_ema = tf.train.ExponentialMovingAverage(decay=0.999)
    gan_train_ops, generator_vars_to_restore = \
        train_util.add_generator_smoothing_ops(generator_ema,
                                               gan_model,
                                               gan_train_ops)
    load_scope = tf.variable_scope(
        gan_model.generator_scope,
        reuse=True,
        custom_getter=train_util.make_var_scope_custom_getter_for_ema(
            generator_ema))

    ########## Separate path for generating samples with a placeholder (ph)
    # Mapping of pitches to one-hot labels
    pitch_counts = data_helper.get_pitch_counts()
    pitch_to_label_dict = {}
    for i, pitch in enumerate(sorted(pitch_counts.keys())):
      pitch_to_label_dict[pitch] = i

    # (label_ph, noise_ph) -> fake_wave_ph
    labels_ph = tf.placeholder(tf.int32, [batch_size])
    noises_ph = tf.placeholder(tf.float32, [batch_size,
                                            config['latent_vector_size']])
    num_pitches = len(pitch_counts)
    one_hot_labels_ph = tf.one_hot(labels_ph, num_pitches)
    with load_scope:
      fake_data_ph, _ = g_fn((noises_ph, one_hot_labels_ph))
      fake_waves_ph = data_helper.data_to_waves(fake_data_ph)

    if config['train_time_limit'] is not None:
      stage_train_time_limit = stage_times[stage_id]
      #  config['train_time_limit'] * \
      # (float(stage_id+1) / ((2*config['num_resolutions'])-1))
    else:
      stage_train_time_limit = None

    ########## Add variables as properties
    self.stage_id = stage_id
    self.batch_size = batch_size
    self.config = config
    self.data_helper = data_helper
    self.resolution_schedule = resolution_schedule
    self.num_images = num_images
    self.num_blocks = num_blocks
    self.current_image_id = current_image_id
    self.progress = progress
    self.generator_fn = g_fn
    self.discriminator_fn = d_fn
    self.gan_model = gan_model
    self.fake_ac_loss = fake_ac_loss
    self.real_ac_loss = real_ac_loss
    self.gan_loss = gan_loss
    self.gan_train_ops = gan_train_ops
    self.optimizer_var_list = optimizer_var_list
    self.generator_ema = generator_ema
    self.generator_vars_to_restore = generator_vars_to_restore
    self.real_images = real_images
    self.real_one_hot_labels = real_one_hot_labels
    self.load_scope = load_scope
    self.pitch_counts = pitch_counts
    self.pitch_to_label_dict = pitch_to_label_dict
    self.labels_ph = labels_ph
    self.noises_ph = noises_ph
    self.fake_waves_ph = fake_waves_ph
    self.saver = tf.train.Saver()
    self.sess = tf.Session()
    self.train_time = train_time
    self.stage_train_time_limit = stage_train_time_limit