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
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