def _train_generator(self, features, labels, step, optimizer, params):
# Set the random offset tensor for operations in tpu_random.py.
tpu_random.set_random_offset_from_features(features)
# create_loss will set self.g_loss.
self.create_loss(features, labels, params=params)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(
self.g_loss,
var_list=self.generator.trainable_variables,
global_step=step)
if self._g_use_ema:
g_vars = self.generator.trainable_variables
with tf.name_scope("generator_ema"):
logging.info("Creating moving averages of weights: %s",
g_vars)
# The decay value is set to 0 if we're before the moving-average start
# point, so that the EMA vars will be the normal vars.
decay = self._ema_decay * tf.cast(
tf.greater_equal(step, self._ema_start_step),
tf.float32)
ema = tf.train.ExponentialMovingAverage(decay=decay)
with tf.control_dependencies([train_op]):
train_op = ema.apply(g_vars)
with tf.control_dependencies([train_op]):
return tf.identity(self.g_loss)
def _train_discriminator(self, features, labels, step, optimizer, params):
features = features.copy()
features["generated"] = tf.stop_gradient(features["generated"])
# Set the random offset tensor for operations in tpu_random.py.
tpu_random.set_random_offset_from_features(features)
# create_loss will set self.d_loss.
self.create_loss(features, labels, params=params)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(
self.d_loss,
var_list=self.discriminator.trainable_variables,
global_step=step)
with tf.control_dependencies([train_op]):
return tf.identity(self.d_loss)
def model_fn(features, labels, mode, params):
# Set the random offset tensor for operations in tpu_random.py.
tpu_random.set_random_offset_from_features(features)
test_op = create_op_fn()
predictions = tf.layers.dense(features["x"], 1)
loss = tf.losses.mean_squared_error(labels, predictions)
optimizer = tf.train.GradientDescentOptimizer(0.01)
if params["use_tpu"]:
optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
with tf.control_dependencies([test_op]):
train_op = optimizer.minimize(
loss, global_step=tf.train.get_or_create_global_step())
return tf.contrib.tpu.TPUEstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def _create_sub_step_loss(sub_step_idx=0, reuse=True):
"""Creates the loss for a slice of the current batch.
Args:
sub_step_idx: Index of the slice of the batch to use to construct the
loss. If self.unroll_disc_iters is True this must be 0 and the whole
batch will be used.
reuse: Bool, whether to reuse existing variables for the models.
Should be False for the first call and True on all other calls.
Returns:
Fake images created by the generator.
"""
logging.info("sub_step_idx: %s, params: %s", sub_step_idx, params)
# Set the random offset tensor for operations in tpu_random.py.
tpu_random.set_random_offset_from_features(fs[sub_step_idx])
self.create_loss(fs[sub_step_idx],
ls[sub_step_idx],
params,
is_training=is_training,
reuse=reuse)