def train_rgan(gan_model, dataset, n_epochs):
generator_optimizer = keras.optimizers.Adam(lr=0.01, beta_1=0.9, beta_2=0.999, clipnorm=1.)
discriminator_optimizer = keras.optimizers.SGD(lr=0.1, momentum=0.9, nesterov=True, clipnorm=1.)
recurrent_generator, recurrent_discriminator = gan_model.layers
# Keep results for plotting
train_discriminator_loss_results = []
train_generator_loss_results = []
for epoch in range(n_epochs):
epoch_discriminator_loss_avg = Mean()
epoch_generator_loss_avg = Mean()
for x_batch, mask_batch in dataset:
no, seq_len , dim = x_batch.shape
x_batch = cast(x_batch, float32)
# phase 1 - training the discriminator
noise = noise_generator(no, seq_len, dim)
generated_samples = recurrent_generator(noise)
x_fake_and_real = concat([generated_samples, x_batch], axis=1)
y1 = cast(reshape(constant([[0.]] * seq_len + [[1.]] * seq_len), [seq_len*2, 1]), float32)
y1 = tf.broadcast_to(y1, [no, seq_len*2, 1])
mask1 = tf.ones([no, seq_len])
mask_fake_and_real = concat([mask1, mask_batch], axis=1)
recurrent_discriminator.trainable = True
discriminator_loss_value, discriminator_grads = grad(recurrent_discriminator, x_fake_and_real, y1, mask_fake_and_real)
discriminator_optimizer.apply_gradients(zip(discriminator_grads, recurrent_discriminator.trainable_variables))
# phase 2 - training the generator
noise = noise_generator(no, seq_len, dim)
y2 = cast(reshape( constant([[1.]] * seq_len), [seq_len, 1]), float32)
y2 = tf.broadcast_to(y2, [no, seq_len, 1])
recurrent_discriminator.trainable = False
generator_loss_value, generator_grads = grad(gan_model, noise, y2, mask1)
generator_optimizer.apply_gradients(zip(generator_grads, gan_model.trainable_variables))
# Track progress: Add current batch loss
epoch_discriminator_loss_avg.update_state(discriminator_loss_value)
epoch_generator_loss_avg.update_state(generator_loss_value)
# End epoch
train_discriminator_loss_results.append(epoch_discriminator_loss_avg.result())
train_generator_loss_results.append(epoch_generator_loss_avg.result())
if epoch % 50 == 0:
print("RGAN Epoch {:03d}: Discriminator Loss: {:.3f}".format(epoch, epoch_discriminator_loss_avg.result() ) , file=sys.stdout)
print("RGAN Epoch {:03d}: Generator Loss: {:.3f}".format(epoch, epoch_generator_loss_avg.result() ) , file=sys.stdout)
return gan_model, train_discriminator_loss_results, train_generator_loss_results
class VFIB(keras.Model):
def __init__(self, encoder, predictor, feature_dim,loss_type, **kwargs):
super(VFIB, self).__init__(**kwargs)
self.encoder = encoder
self.classifier = predictor
self.loss_type = loss_type
self.total_loss_tracker = Mean(name="total_loss")
self.prediction_loss_tracker = Mean(name="prediction_loss")
self.kl_loss_tracker = Mean(name="kl_loss")
self.mmd_loss_tracker = Mean(name="mmd_loss")
@property
def metrics(self):
return [
self.total_loss_tracker,
self.prediction_loss_tracker,
self.kl_loss_tracker,
self.mmd_loss_tracker
]
def call(self, inputs):
# 0 refers to first column with sensitive feature 'Age'
sens, _ = split_sensitive_X(inputs, 0, 1)
mu, sig, z = self.encoder(inputs)
preds = self.classifier(tf.concat([z, sens], 1))
return mu, sig, z, preds
def train_step(self, data):
X, y = data
with tf.GradientTape() as tape:
z_mean, z_log_sigma, z, preds = self.call(X)
prediction_loss = neg_log_bernoulli(y, preds)
kl_loss = KL(z_mean, z_log_sigma)
mmd_loss = mmd_loss(X, z)
if self.loss_type=='all':
total_loss = prediction_loss+ kl_loss + mmd_loss
elif self.loss_type=='kl':
total_loss = prediction_loss+ kl_loss
else:
total_loss = prediction_loss
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
self.total_loss_tracker.update_state(total_loss)
self.prediction_loss_tracker.update_state(prediction_loss)
self.kl_loss_tracker.update_state(kl_loss)
self.mmd_loss_tracker.update_state(mmd_loss)
return {
"loss": self.total_loss_tracker.result(),
"classification_loss": self.prediction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
"mmd_loss": self.mmd_loss_tracker.result()
}
class MeanBasedMetric(Metric):
def __init__(self, name, dtype):
super().__init__(name, dtype=dtype)
self._mean = Mean(dtype=dtype)
@abstractmethod
def _objective_function(self, y_true, y_pred):
pass
def update_state(self, y_true, y_pred, sample_weight=None):
values = self._objective_function(y_true, y_pred)
self._mean.update_state(values=values, sample_weight=sample_weight)
def result(self):
return self._mean.result()
def reset_states(self):
self._mean.reset_states()
class StandardVarianceBasedMetric(Metric):
def __init__(self, name, dtype):
super().__init__(name, dtype=dtype)
self._mean = Mean(dtype=dtype)
self._square_mean = Mean(dtype=dtype)
@abstractmethod
def _objective_function(self, y_true, y_pred):
pass
def update_state(self, y_true, y_pred, sample_weight=None):
values = self._objective_function(y_true, y_pred)
self._mean.update_state(values=values, sample_weight=sample_weight)
self._square_mean.update_state(values=tf.square(values),
sample_weight=sample_weight)
def result(self):
return tf.sqrt(self._square_mean.result() -
tf.square(self._mean.result()))
def reset_states(self):
self._mean.reset_states()
self._square_mean.reset_states()
def gain_train_step(dataset, gain, n_epochs):
generator, discriminator = gain.layers
discriminator_optimizer = keras.optimizers.SGD(momentum=0.9, nesterov=True)
generator_optimizer = keras.optimizers.Adam()
# Keep results for plotting
train_discriminator_loss_results = []
train_generator_loss_results = []
for epoch in range(n_epochs):
epoch_discriminator_loss_avg = Mean()
epoch_generator_loss_avg = Mean()
for x_batch, mask_batch in dataset:
x_batch = cast(x_batch, float32)
mask_batch = cast(mask_batch, float32)
# phase 1: train discriminator
hint = hint_generator(x_batch, mask_batch)
generated_samples = generator(concat( [hint, mask_batch], axis = 1))
discriminator.trainable = True
discriminator_loss_value, discriminator_grads = gain_grad(discriminator, generated_samples, mask_batch)
discriminator_optimizer.apply_gradients(zip(discriminator_grads, discriminator.trainable_variables))
# phase 2 - training the generator
hint = hint_generator(x_batch, mask_batch)
discriminator.trainable = False
generator_loss_value, generator_grads = gain_grad(gain, concat( [hint, mask_batch], axis = 1), mask_batch)
generator_optimizer.apply_gradients(zip(generator_grads, gain.trainable_variables))
# Track progress: Add current batch loss
epoch_discriminator_loss_avg.update_state(discriminator_loss_value)
epoch_generator_loss_avg.update_state(generator_loss_value)
# End epoch
train_discriminator_loss_results.append(epoch_discriminator_loss_avg.result())
train_generator_loss_results.append(epoch_generator_loss_avg.result())
if epoch % 50 == 0:
print("GAIN Epoch {:03d}: Discriminator Loss: {:.3f}".format(epoch, epoch_discriminator_loss_avg.result() ) , file=sys.stdout)
print("GAIN Epoch {:03d}: Generator Loss: {:.3f}".format(epoch, epoch_generator_loss_avg.result() ) , file=sys.stdout)
return gain, train_discriminator_loss_results, train_generator_loss_results
def train_wgain(dataset, gain, n_epoch, n_critic, alpha):
'''Train wgain function
Args:
- dataset: A dataset TF2 object.
- gain: a gain model.
- n_epoch: number of iterations.
- alpha: hyper-parameter
Returns:
- gain: Trained model
- critic loss, generator loss and reconstruction loss for monitoring.
'''
generator, discriminator = gain.layers
d_optimizer = keras.optimizers.RMSprop(lr=0.00005)
g_optimizer = keras.optimizers.Adam()
# Keep results for plotting
train_d_loss_results = []
train_g_loss_results = []
train_rec_loss_results = []
for epoch in range(n_epoch):
epoch_d_loss_avg = Mean()
epoch_g_loss_avg = Mean()
epoch_rec_loss_avg = Mean()
for x_batch, mask_batch in dataset:
batch_size, dim = x_batch.shape
# phase 1: train discriminator
for _ in range(n_critic):
hint = hint_generator(x_batch, mask_batch)
generated_samples = generator(hint, training = True)
discriminator.trainable = True
d_loss, d_grads = discriminator_grad(discriminator, generated_samples, mask_batch[:,1:])
d_optimizer.apply_gradients(zip(d_grads, discriminator.trainable_variables))
# phase 2 - training the generator
hint = hint_generator(x_batch, mask_batch)
discriminator.trainable = False
g_loss, g_grads = gain_grad(gain, hint)
d_optimizer.apply_gradients(zip(g_grads, gain.trainable_variables))
hint = hint_generator(x_batch, mask_batch)
rec_loss, rec_grads = rec_grad(generator, hint, mask_batch, alpha)
g_optimizer.apply_gradients(zip(rec_grads, gain.trainable_variables))
# Track progress: Add current batch loss
epoch_d_loss_avg.update_state(d_loss)
epoch_g_loss_avg.update_state(g_loss)
epoch_rec_loss_avg.update_state(rec_loss)
# End epoch
train_d_loss_results.append(epoch_d_loss_avg.result())
train_g_loss_results.append(epoch_g_loss_avg.result())
train_rec_loss_results.append(epoch_rec_loss_avg.result())
return gain, train_d_loss_results, train_g_loss_results, train_rec_loss_results
class Pix2Pose(Model):
def __init__(self, image_shape, discriminator, generator, latent_dim):
super(Pix2Pose, self).__init__()
self.image_shape = image_shape
self.discriminator = discriminator
self.generator = generator
self.latent_dim = latent_dim
@property
def metrics(self):
return [self.generator_loss, self.discriminator_loss]
def compile(self, optimizers, losses, loss_weights):
super(Pix2Pose, self).compile()
self.optimizer_generator = optimizers['generator']
self.optimizer_discriminator = optimizers['discriminator']
self.compute_reconstruction_loss = losses['weighted_reconstruction']
self.compute_error_prediction_loss = losses['error_prediction']
self.compute_discriminator_loss = losses['discriminator']
self.generator_loss = Mean(name='generator_loss')
self.discriminator_loss = Mean(name='discriminator_loss')
self.reconstruction_loss = Mean(name='weighted_reconstruction')
self.error_prediction_loss = Mean(name='error_prediction')
self.reconstruction_weight = loss_weights['weighted_reconstruction']
self.error_prediction_weight = loss_weights['error_prediction']
def _build_discriminator_labels(self, batch_size):
return tf.concat([tf.ones(batch_size, 1), tf.zeros(batch_size, 1)], 0)
def _add_noise_to_labels(self, labels):
noise = tf.random.uniform(tf.shape(labels))
labels = labels + 0.05 * noise
return labels
def _get_batch_size(self, values):
return tf.shape(values)[0]
def _train_discriminator(self, RGB_inputs, RGBA_true):
RGB_true = RGBA_true[:, :, :, 0:3]
RGB_fake = self.generator(RGB_inputs)[:, :, :, 0:3]
RGB_fake_true = tf.concat([RGB_fake, RGB_true], axis=0)
batch_size = self._get_batch_size(RGB_inputs)
y_true = self._build_discriminator_labels(batch_size)
y_true = self._add_noise_to_labels(y_true)
with tf.GradientTape() as tape:
y_pred = self.discriminator(RGB_fake_true)
discriminator_loss = self.compute_discriminator_loss(
y_true, y_pred)
gradients = tape.gradient(discriminator_loss,
self.discriminator.trainable_weights)
self.optimizer_discriminator.apply_gradients(
zip(gradients, self.discriminator.trainable_weights))
return discriminator_loss
def _train_generator(self, RGB_inputs):
batch_size = tf.shape(RGB_inputs)[0]
y_misleading = tf.zeros((batch_size, 1))
with tf.GradientTape() as tape:
RGBE_preds = self.generator(RGB_inputs)
y_pred = self.discriminator(RGBE_preds[..., 0:3])
generator_loss = self.compute_discriminator_loss(
y_misleading, y_pred)
gradients = tape.gradient(generator_loss,
self.generator.trainable_weights)
self.optimizer_generator.apply_gradients(
zip(gradients, self.generator.trainable_weights))
return generator_loss
def _train_reconstruction(self, RGB_inputs, RGBA_true):
with tf.GradientTape() as tape:
RGBE_pred = self.generator(RGB_inputs)
reconstruction_loss = self.compute_reconstruction_loss(
RGBA_true, RGBE_pred)
reconstruction_loss = (self.reconstruction_weight *
reconstruction_loss)
gradients = tape.gradient(reconstruction_loss,
self.generator.trainable_weights)
self.optimizer_generator.apply_gradients(
zip(gradients, self.generator.trainable_weights))
return reconstruction_loss
def _train_error_prediction(self, RGB_inputs, RGBA_true):
with tf.GradientTape() as tape:
RGBE_pred = self.generator(RGB_inputs)
error_prediction_loss = self.compute_error_prediction_loss(
RGBA_true, RGBE_pred)
error_prediction_loss = (self.error_prediction_weight *
error_prediction_loss)
gradients = tape.gradient(error_prediction_loss,
self.generator.trainable_weights)
self.optimizer_generator.apply_gradients(
zip(gradients, self.generator.trainable_weights))
return error_prediction_loss
def train_step(self, data):
RGB_inputs, RGBA_true = data[0]['RGB_input'], data[1]['RGB_with_error']
reconstruction_loss = self._train_reconstruction(RGB_inputs, RGBA_true)
self.reconstruction_loss.update_state(reconstruction_loss)
error_loss = self._train_error_prediction(RGB_inputs, RGBA_true)
self.error_prediction_loss.update_state(error_loss)
discriminator_loss = self._train_discriminator(RGB_inputs, RGBA_true)
self.discriminator_loss.update_state(discriminator_loss)
generator_loss = self._train_generator(RGB_inputs)
self.generator_loss.update_state(generator_loss)
return {
'discriminator_loss': self.discriminator_loss.result(),
'generator_loss': self.generator_loss.result(),
'reconstruction_loss': self.reconstruction_loss.result(),
'error_prediction_loss': self.error_prediction_loss.result()
}
class StyleTransfer(Model):
def __init__(self, *args, encoder=None, decoder=None, **kwargs):
super().__init__(*args, **kwargs)
self.init_encoder(encoder)
self.init_decoder(decoder)
self.build_metrics()
def init_encoder(self, encoder):
if encoder is not None:
self.encoder = encoder
return
self.encoder = Sequential([Input((224, 224, 3))])
self.encoder.trainable = False
for layer_name in STYLE_LAYERS:
layer = vgg19.get_layer(layer_name)
self.encoder.add(layer)
self.encoder.compile()
def init_decoder(self, decoder):
if decoder is not None:
self.decoder = decoder
return
# Build the decoder if it wasn't provided
input_shape = self.encoder.layers[-1].output_shape[1:]
self.decoder = Sequential([Input(input_shape)])
# The decoder is the trimed inverse of the encoder
for layer_name in DECODER_LAYERS:
layer = vgg19.get_layer(layer_name)
# Add the upsampling to double the image size
if 'pool' in layer.name:
block_name = layer.name.split("_")[0]
self.decoder.add(UpSampling2D(name=f'{block_name}_upsampling'))
# Add some reflective padding followed by a Conv2D layer
elif 'conv' in layer.name:
self.decoder.add(
Conv2DReflectivePadding(filters=layer.output_shape[-1],
kernel_size=layer.kernel_size,
strides=layer.strides,
activation='relu',
name=layer.name))
# Add one final Conv2D to reduce the feature maps to 3 (N,W,H,3)
self.decoder.add(
Conv2DReflectivePadding(3, (3, 3), name='output_conv1'))
def build_metrics(self):
self.c_loss_metric = Mean(name='c_loss')
self.s_loss_metric = Mean(name='s_loss')
self.loss_metric = Mean(name='loss')
def compile(self, optimizer, content_loss, style_loss, **kwargs):
super().compile(**kwargs)
if not getattr(self, 'decoder_compiled', False):
self.decoder.compile(optimizer=optimizer)
self.content_loss = content_loss
self.style_loss = style_loss
self.adain = AdaIN()
@tf.function
def train_step(self, data, training=True):
c_encoded_outputs, s_encoded_outputs = data
# The outputs of the encoded style images and the encoded generated images
# Retrieved from the selected encoder layers used for the loss
s_loss_outputs = []
g_loss_outputs = []
with tf.GradientTape(watch_accessed_variables=training) as tape:
# 1. Encode the content and style image
for layer in self.encoder.layers:
# Encode the content image
c_encoded_outputs = layer(c_encoded_outputs)
# Encode the style image
s_encoded_outputs = layer(s_encoded_outputs)
# If this layer is used to calculate the loss save its outputs
if layer.name in LOSS_LAYERS:
s_loss_outputs.append(s_encoded_outputs)
# 2. Adaptive Instance Normalization
adain_outputs = self.adain(c_encoded_outputs, s_encoded_outputs)
# 3. Decode the feature maps generated by AdaIN to get the final generated image
generated_imgs = self.decoder(adain_outputs, training=training)
# 4. Encode the generated image to calculate the loss
g_encoded_outputs = generated_imgs
for layer in self.encoder.layers:
# Encode the generated image
g_encoded_outputs = layer(g_encoded_outputs)
# If this layer is used to calculate the loss save its outputs
if layer.name in LOSS_LAYERS:
g_loss_outputs.append(g_encoded_outputs)
# 5. Calculate the content loss
c_per_replica_loss = self.content_loss(g_encoded_outputs,
adain_outputs) # (N,W,H)
# Reduce the loss (we do this ourselves in order to be compatible with distributed training)
global_c_loss_size = tf.size(
c_per_replica_loss
) * self.distribute_strategy.num_replicas_in_sync
global_c_loss_size = tf.cast(global_c_loss_size, dtype=tf.float32)
c_loss = tf.nn.compute_average_loss(
c_per_replica_loss, global_batch_size=global_c_loss_size)
assert len(g_loss_outputs) == len(s_loss_outputs)
# 6. Calculate style loss
s_loss = 0
for i in range(len(g_loss_outputs)):
s_per_replica_loss = self.style_loss(g_loss_outputs[i],
s_loss_outputs[i]) # (N,)
# Reduce the loss (we do this ourselves in order to be compatible with distributed training)
global_s_loss_size = BATCH_SIZE * self.distribute_strategy.num_replicas_in_sync
s_loss += tf.nn.compute_average_loss(
s_per_replica_loss, global_batch_size=global_s_loss_size)
# 7. Calculate the loss
loss = c_loss + s_loss * STYLE_WEIGHT
# 8. Apply gradient descent
if training:
gradients = tape.gradient(loss, self.decoder.trainable_variables)
# gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
# tf.print('---')
# tf.print('glonorm', tf.linalg.global_norm(gradients))
# tf.print(list((i, tf.math.reduce_min(n), tf.math.reduce_max(n)) for i,n in enumerate(gradients)))
# tf.print(list((i, tf.math.reduce_min(n), tf.math.reduce_max(n)) for i,n in enumerate(self.decoder.trainable_variables)))
# tf.print('C_ENC --> ', tf.math.reduce_min(c_encoded_outputs), tf.math.reduce_max(c_encoded_outputs))
# tf.print('S_ENC --> ', tf.math.reduce_min(s_encoded_outputs), tf.math.reduce_max(s_encoded_outputs))
# tf.print('ADAIN --> ', tf.math.reduce_min(adain_outputs), tf.math.reduce_max(adain_outputs))
# tf.print('GEN_IMG --> ', tf.math.reduce_min(generated_imgs), tf.math.reduce_max(generated_imgs))
# tf.print('G_ENC --> ', tf.math.reduce_min(g_encoded_outputs), tf.math.reduce_max(g_encoded_outputs))
# tf.print('S_LOSS_E --> ', tf.math.reduce_min(s_loss_outputs[i]), tf.math.reduce_max(s_loss_outputs[i]))
# tf.print('G_LOSS_E --> ', tf.math.reduce_min(g_loss_outputs[i]), tf.math.reduce_max(g_loss_outputs[i]))
self.decoder.optimizer.apply_gradients(
zip(gradients, self.decoder.trainable_variables))
# 9. Update the metrics
self.c_loss_metric.update_state(c_loss)
self.s_loss_metric.update_state(s_loss)
self.loss_metric.update_state(loss)
return {m.name: m.result() for m in self.metrics}
@tf.function
def test_step(self, data):
c_encoded_outputs, s_encoded_outputs = data
# The outputs of the encoded style images and the encoded generated images
# Retrieved from the selected encoder layers used for the loss
s_loss_outputs = []
g_loss_outputs = []
# 1. Encode the content and style image
for layer in self.encoder.layers:
# Encode the content image
c_encoded_outputs = layer(c_encoded_outputs, training=False)
# Encode the style image
s_encoded_outputs = layer(s_encoded_outputs, training=False)
# If this layer is used to calculate the loss save its outputs
if layer.name in LOSS_LAYERS:
s_loss_outputs.append(s_encoded_outputs)
# 2. Adaptive Instance Normalization
adain_outputs = self.adain(c_encoded_outputs, s_encoded_outputs)
# 3. Decode the feature maps generated by AdaIN to get the final generated image
generated_imgs = self.decoder(adain_outputs, training=False)
# 4. Encode the generated image to calculate the loss
g_encoded_outputs = generated_imgs
for layer in self.encoder.layers:
# Encode the generated image
g_encoded_outputs = layer(g_encoded_outputs, training=False)
# If this layer is used to calculate the loss save its outputs
if layer.name in LOSS_LAYERS:
g_loss_outputs.append(g_encoded_outputs)
# 5. Calculate the content loss
c_per_replica_loss = self.content_loss(g_encoded_outputs,
adain_outputs) # (N,W,H)
# Reduce the loss (we do this ourselves in order to be compatible with distributed training)
global_c_loss_size = tf.size(
c_per_replica_loss) * self.distribute_strategy.num_replicas_in_sync
global_c_loss_size = tf.cast(global_c_loss_size, dtype=tf.float32)
c_loss = tf.nn.compute_average_loss(
c_per_replica_loss, global_batch_size=global_c_loss_size)
assert len(g_loss_outputs) == len(s_loss_outputs)
# 6. Calculate style loss
s_loss = 0
for i in range(len(g_loss_outputs)):
s_per_replica_loss = self.style_loss(g_loss_outputs[i],
s_loss_outputs[i]) # (N,)
# Reduce the loss (we do this ourselves in order to be compatible with distributed training)
global_s_loss_size = BATCH_SIZE * self.distribute_strategy.num_replicas_in_sync
s_loss += tf.nn.compute_average_loss(
s_per_replica_loss, global_batch_size=global_s_loss_size)
# 7. Calculate the loss
loss = c_loss + s_loss * STYLE_WEIGHT
# 9. Update the metrics
self.c_loss_metric.update_state(c_loss)
self.s_loss_metric.update_state(s_loss)
self.loss_metric.update_state(loss)
return {m.name: m.result() for m in self.metrics}
@tf.function
def predict_step(self, data):
content_imgs, style_imgs = data[0]
# Ensure these are batched
assert len(content_imgs.shape) == 4
assert len(style_imgs.shape) == 4
content_imgs = vgg19_preprocess_input(content_imgs) / 255.0
style_imgs = vgg19_preprocess_input(style_imgs) / 255.0
c_encoded = content_imgs
s_encoded = style_imgs
# Encode the contents and styles
for layer in self.encoder.layers:
c_encoded = layer(c_encoded)
s_encoded = layer(s_encoded)
# Apply adaptive batch normalization
adain_outputs = AdaIN()(c_encoded, s_encoded)
# Decode the images to generate them
generated_imgs = self.decoder(adain_outputs)
generated_imgs = self.deprocess_vgg19(generated_imgs)
return generated_imgs
@property
def metrics(self):
return [self.loss_metric, self.c_loss_metric, self.s_loss_metric]
def deprocess_vgg19(self, imgs):
# Ensure they are batched
assert len(imgs.shape) == 4
# Put back to 0...255
imgs *= 255.0
# Add mean
imgs += [103.939, 116.779, 123.68]
# BGR to RGB
imgs = imgs[..., ::-1]
# Clip
imgs = tf.clip_by_value(imgs, 0.0, 255.0)
# Cast
imgs = tf.cast(imgs, tf.uint8)
return imgs
def save_architecture(self, log_dir):
# Ensure the log_dir exists
pathlib.Path(log_dir).mkdir(parents=True, exist_ok=True)
with GFile(os.path.join(log_dir, 'style_transfer_architecture.json'),
'w') as f:
f.write(self.to_json())
def save_encoder(self, log_dir):
self.encoder.save(log_dir)
def save_model(self, log_dir, **kwargs):
self.decoder.save(log_dir, **kwargs)
@classmethod
def load(cls, log_dir=None):
model_found = bool(log_dir) and pathlib.Path(
os.path.join(log_dir,
'style_transfer_architecture.json')).is_file()
# If there isn't already a model create one from scratch and save it
if not model_found:
model = cls()
if log_dir:
model.save_architecture(log_dir)
model.save_encoder(os.path.join(log_dir, 'encoder'))
return model
# Load the model's architecture
with tf.keras.utils.custom_object_scope({
'StyleTransfer':
cls,
'Conv2DReflectivePadding':
Conv2DReflectivePadding
}):
saved_json = GFile(
os.path.join(log_dir, 'style_transfer_architecture.json'),
'r').read()
model = tf.keras.models.model_from_json(saved_json)
model.encoder = tf.keras.models.load_model(
os.path.join(log_dir, 'encoder'))
# Load the decoder's latest weights if there are any
ckpts = glob.glob(os.path.join(log_dir, 'weights', '*'))
if ckpts:
latest_ckpt = max(ckpts, key=os.path.getmtime)
print('Loading Checkpoint:', latest_ckpt)
model.decoder = tf.keras.models.load_model(latest_ckpt)
model.decoder_compiled = True
return model
def get_config(self):
return {'encoder': self.encoder, 'decoder': self.decoder}
@classmethod
def from_config(cls, config, **kwargs):
encoder = tf.keras.models.model_from_json(
json.dumps(config.pop('encoder')))
decoder = tf.keras.models.model_from_json(
json.dumps(config.pop('decoder')))
style_transfer = cls(encoder=encoder, decoder=decoder)
return style_transfer
class VAE(Model):
def __init__(self, input_dim, encoder_kwargs, decoder_kwargs, warmup_steps,
**kwargs):
super(VAE, self).__init__(**kwargs)
self.input_dim = input_dim
# Generate the encoder and decoder networks
self.encoder = vae_utils.generate_encoder(**encoder_kwargs)
self.decoder = vae_utils.generate_decoder(**decoder_kwargs)
self.latent_dim = encoder_kwargs['latent_dim']
# Combine encoder and decoder to create VAE structure
self.input_tensor = Input(shape=(input_dim, ))
if encoder_kwargs.get('is_variational', True):
self.latent_tensor = self.encoder(self.input_tensor)[-1]
else:
self.latent_tensor = self.encoder(self.input_tensor)
self.output_tensor = self.decoder(self.latent_tensor)
self.vae = Model(inputs=self.input_tensor,
outputs=[self.output_tensor, self.latent_tensor],
name='vae')
# We changed call() method so that self is self.vae. We compile the above model.
self.compile(optimizer=Adam(learning_rate=1e-3))
# Logging
self.total_loss_tracker = Mean(name='total_loss')
self.reconstruction_loss_tracker = Mean(name='reconstruction_loss')
self.kl_loss_tracker = Mean(name="kl_loss")
# Hyper-parameters
self.warmup_steps = tf.constant(warmup_steps, dtype=tf.int32)
self.it = tf.Variable(0, dtype=tf.int32, trainable=False)
# The model is actually the self.vae Model inside this class.
def call(self, inputs, training=None, mask=None):
return self.vae(inputs)
@property
def metrics(self):
return [
self.total_loss_tracker,
self.reconstruction_loss_tracker,
self.kl_loss_tracker,
]
# @tf.function
def train_step(self, data):
return self.train_step_deterministic(data)
# return self.train_step_variational(data)
def train_step_deterministic(self, data):
with tf.GradientTape() as tape:
# Get encoder outputs
z = self.encoder(data)
# How good are we at reconstruction?
reconstruction = self.decoder(z)
reconstruction_loss = tf.reduce_mean(
tf.square(data - reconstruction))
self.it.assign_add(1)
tf.print(tf.reduce_mean(data - reconstruction, axis=0))
# Gradients w.r.t. the total_loss in the GradientTape
grads = tape.gradient(reconstruction_loss, self.trainable_weights)
# Update the network parameters
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
# Logging
# self.total_loss_tracker.update_state(total_loss)
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
# self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": self.total_loss_tracker.result(),
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
}
def train_step_variational(self, data):
with tf.GradientTape() as tape:
# Get encoder outputs
z_mean, z_log_var, z = self.encoder(data)
# How good are we at reconstruction?
reconstruction = self.decoder(z)
reconstruction_loss = tf.reduce_mean(
tf.square(data - reconstruction))
# How much regularized is our latent space?
kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) -
tf.exp(z_log_var))
kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
# Mask kl_loss during warm-up phase
kl_loss = tf.cond(
tf.math.greater_equal(self.it, self.warmup_steps),
lambda: kl_loss, lambda: 0.0)
# We will minimize this loss
total_loss = reconstruction_loss + kl_loss / 10.0
self.it.assign_add(1)
# Gradients w.r.t. the total_loss in the GradientTape
grads = tape.gradient(total_loss, self.vae.trainable_weights)
# Update the network parameters
self.optimizer.apply_gradients(zip(grads, self.vae.trainable_weights))
# Logging
self.total_loss_tracker.update_state(total_loss)
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": self.total_loss_tracker.result(),
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
}
# train Dz
for _ in range(n_critic):
random_noise = generate_noise(no, latent_dim) # z
generated_samples = rnn_generator(random_noise, training = True) # x_hat
encodings = rnn_encoder(x_batch, training = True) # z_hat
encoded_noise = rnn_encoder(generated_samples, training = True) # z'
rnn_ae_critc.trainable = True
ae_loss1, ae_grads1 = critic_loss(rnn_ae_critc, encoded_noise, encodings, y1)
rmsprop_optimizer.apply_gradients(zip(ae_grads1, rnn_ae_critc.trainable_variables))
# update generator via Dz critic's error
random_noise = generate_noise(no, latent_dim) # z
rnn_ae_critc.trainable = False
ae_loss2, ae_grads2 = generator_loss(gan_model2, random_noise, y2)
rmsprop_optimizer.apply_gradients(zip(ae_grads2, gan_model2.trainable_variables))
# Track progress: Add current batch loss
epoch_dx_loss_avg.update_state(d_loss1)
epoch_dz_loss_avg.update_state(ae_loss1)
epoch_rec_loss_avg.update_state(rec_loss1)
# End epoch
train_dx_loss_results.append(epoch_dx_loss_avg.result())
train_dz_loss_results.append(epoch_dz_loss_avg.result())
train_rec_loss_results.append(epoch_rec_loss_avg.result())
return rnn_encogen, train_dx_loss_results, train_dz_loss_results, train_rec_loss_results
def rae_wgan(input_dat, seq_times, padding_mask, batch_size, n_epoch, n_critic, latent_dim):
'''
Call to RAE GAN
args:
input_dat: imputed data, inputs to RAE WGAN.
class build_prob_u_net(Model):
def __init__(self, num_classes, activation, latent_dim=6, resolution_lvl=5,
img_shape=(None, None, 1), seg_shape=(None, None, 1),
num_filters=(32, 64, 128, 256, 512), downsample_signal=(2,2,2,2,2)):
super(build_prob_u_net, self).__init__()
self.num_classes = num_classes
self.latent_dim = latent_dim
self.activation = activation
self.num_filters = num_filters
self.resolution_lvl = resolution_lvl
self.downsample_signal = downsample_signal
self.prior = self.latent_space_net(img_shape, None)
self.posterior = self.latent_space_net(img_shape, seg_shape)
self.det_unet = self.unet(img_shape)
def latent_space_net(self, img_shape, seg_shape):
if seg_shape is not None:
# Posterior inputs
inputs = [Input(shape=img_shape), Input(shape=seg_shape)]
input_ = Concatenate(name='input_con') (inputs)
name = 'prob_unet_posterior'
else:
# Prior input
inputs = Input(shape=img_shape)
input_ = inputs
name = 'prob_unet_prior'
# Encoder blocks
for i in range(self.resolution_lvl):
if i == 0:
x = conv_block(self.num_filters[i], 0, i, amount=2, type_block='encoder_latent') (input_)
else:
x = MaxPool2D(pool_size=self.downsample_signal[i],
name='encoder_latent_stage0-{}_pool'.format(i)) (x)
x = conv_block(self.num_filters[i], 0, i, amount=2, type_block='encoder_latent') (x)
# Z sample
z, mu, sigma = z_mu_sigma(self.latent_dim, 0, self.resolution_lvl+1) (x)
return Model(inputs, [z, mu, sigma], name=name)
def unet(self, img_shape):
lvl_div = np.power(2, self.resolution_lvl-1)
z_sample = Input(shape=(None, None, self.latent_dim))
inputs = Input(shape=img_shape)
skip_connections = [None] * self.resolution_lvl
# Encoder blocks
for i in range(self.resolution_lvl):
if i == 0:
x = conv_block(self.num_filters[i], 0, i, amount=2, type_block='encoder') (inputs)
else:
x = MaxPool2D(pool_size=self.downsample_signal[i],
name='encoder_stage0-{}_pool'.format(i)) (x)
x = conv_block(self.num_filters[i], 0, i, amount=2, type_block='encoder') (x)
skip_connections[i] = x
skip_connections = skip_connections[:-1]
# Decoder blocks
for i in reversed(range(self.resolution_lvl-1)):
x = UpSampling2D(size=self.downsample_signal[i],
name='decoder_stage0-{}_up'.format(i)) (x)
x = Concatenate(name='decoder_stage0-{}_con'.format(i)) ([x, skip_connections[i]])
x = conv_block(self.num_filters[i], 0, i, amount=2, type_block='decoder') (x)
# Concatenate U-Net and Z sample
broadcast_z = tf.tile(z_sample, (1, lvl_div, lvl_div, 1))
x = Concatenate(name='final_con') ([x, broadcast_z])
x = conv_block(self.num_filters[0], 0, i, amount=2, type_block='final') (x)
x = Conv2D(self.num_classes, kernel_size=1, padding='same',
activation=self.activation, name='final_conv') (x)
return Model([inputs, z_sample], x, name='prob_unet_det')
def kl_score(self, mu0, sigma0, mu1, sigma1):
# Calculate kl loss
sigma0_f = K.square(K.flatten(sigma0))
sigma1_f = K.square(K.flatten(sigma1))
logsigma0 = K.log(sigma0_f + 1e-10)
logsigma1 = K.log(sigma1_f + 1e-10)
mu0_f = K.flatten(mu0)
mu1_f = K.flatten(mu1)
return tf.reduce_mean(
0.5*tf.reduce_sum(tf.divide(sigma0_f + tf.square(mu1_f - mu0_f), sigma1_f + 1e-10)
+ logsigma1 - logsigma0 - 1, axis=-1))
def compile(self, prior_opt, posterior_opt, unet_opt, loss, metric, beta=1):
super(build_prob_u_net, self).compile()
self.posterior_opt = posterior_opt
self.prior_opt = prior_opt
self.unet_opt = unet_opt
self.beta = beta
self.metric = metric
self.compiled_loss = loss
self.metric_tracker = Mean(name='metric')
self.kl_loss_tracker = Mean(name="kl_loss")
self.total_loss_tracker = Mean(name='total_loss')
self.compiled_loss_tracker = Mean(name='compiled_loss')
@property
def metrics(self):
return [
self.compiled_loss_tracker,
self.kl_loss_tracker,
self.total_loss_tracker,
self.metric_tracker
]
def train_step(self, data):
img, seg = data
with tf.GradientTape(persistent=True) as tape:
#Prior and Posterior
_, mu_prior, sigma_prior = self.prior(img, training=True)
z_posterior, mu_posterior, sigma_posterior = self.posterior([img, seg], training=True)
#U-Net
reconstruction = self.det_unet([img, z_posterior], training=True)
#Calculate losses and metric
kl_loss = self.kl_score(mu_posterior, sigma_posterior, mu_prior, sigma_prior)
reconstruction_loss = self.compiled_loss(seg, reconstruction)
total_loss = reconstruction_loss + self.beta * kl_loss
dsc_score = self.metric(seg, reconstruction)
# Update weights
grad_prior = tape.gradient(kl_loss, self.prior.trainable_weights)
self.prior_opt.apply_gradients(zip(grad_prior, self.prior.trainable_weights))
grad_posterior = tape.gradient(kl_loss, self.posterior.trainable_weights)
self.posterior_opt.apply_gradients(zip(grad_posterior, self.posterior.trainable_weights))
grad_unet = tape.gradient(reconstruction_loss, self.det_unet.trainable_weights)
self.unet_opt.apply_gradients(zip(grad_unet, self.det_unet.trainable_weights))
self.metric_tracker.update_state(dsc_score)
self.kl_loss_tracker.update_state(kl_loss)
self.total_loss_tracker.update_state(total_loss)
self.compiled_loss_tracker.update_state(reconstruction_loss)
return {
"loss": self.compiled_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
"total_loss": self.total_loss_tracker.result(),
"dice_coef": self.metric_tracker.result()
}
def test_step(self, data):
img, seg = data
z_prior, mu_prior, sigma_prior = self.prior(img, training=False)
_, mu_posterior, sigma_posterior = self.posterior([img, seg], training=False)
reconstruction = self.det_unet([img, z_prior], training=False)
kl_loss = self.kl_score(mu_posterior, sigma_posterior, mu_prior, sigma_prior)
reconstruction_loss = tf.reduce_mean(tf.reduce_sum(self.compiled_loss(seg, reconstruction)))
total_loss = reconstruction_loss + self.beta * kl_loss
dsc_score = tf.reduce_mean(tf.reduce_sum(self.metric(seg, reconstruction)))
self.metric_tracker.update_state(dsc_score)
self.kl_loss_tracker.update_state(kl_loss)
self.total_loss_tracker.update_state(total_loss)
self.compiled_loss_tracker.update_state(reconstruction_loss)
return {
"loss": self.compiled_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
"total_loss": self.total_loss_tracker.result(),
"dice_coef": self.metric_tracker.result()
}
class Trainer():
def __init__(self,
fbnet,
input_shape,
initial_temperature=5,
temperature_decay_rate=0.956,
temperature_decay_steps=1,
latency_alpha=0.2,
latency_beta=0.6,
weight_lr=0.01,
weight_momentum=0.9,
weight_decay=1e-4,
theta_lr=1e-3,
theta_beta1=0.9,
theta_beta2=0.999,
theta_decay=5e-4):
self._epoch = 0
self.initial_temperature = initial_temperature
self.temperature = initial_temperature
self.latency_alpha = latency_alpha
self.latency_beta = latency_beta
self.exponential_decay = lambda step: exponential_decay(
initial_temperature, temperature_decay_rate,
temperature_decay_steps, step)
fbnet.build(input_shape)
self.fbnet = fbnet
self.weights = []
self.thetas = []
for trainable_weight in fbnet.trainable_weights:
if 'theta' in trainable_weight.name:
self.thetas.append(trainable_weight)
else:
self.weights.append(trainable_weight)
self.weight_opt = SGD(learning_rate=weight_lr,
momentum=weight_momentum,
decay=weight_decay)
self.theta_opt = Adam(learning_rate=theta_lr,
beta_1=theta_beta1,
beta_2=theta_beta2,
decay=theta_decay)
self.loss_fn = SparseCategoricalCrossentropy(from_logits=True)
self.accuracy_metric = SparseCategoricalAccuracy()
self.loss_metric = Mean()
@property
def epoch(self):
return self._epoch
@epoch.setter
def epoch(self, epoch):
self._epoch = epoch
self.temperature = self.exponential_decay(epoch)
def reset_metrics(self):
self.accuracy_metric.reset_states()
self.loss_metric.reset_states()
def _train(self, x, y, weights, opt, training=True):
with tf.GradientTape() as tape:
y_hat = self.fbnet(x, self.temperature, training=training)
loss = self.loss_fn(y, y_hat)
latency = sum(self.fbnet.losses)
loss += latency_loss(latency, self.latency_alpha,
self.latency_beta)
grads = tape.gradient(loss, weights)
opt.apply_gradients(zip(grads, weights))
self.accuracy_metric.update_state(y, y_hat)
self.loss_metric.update_state(loss)
@tf.function
def train_weights(self, x, y):
self._train(x, y, self.weights, self.weight_opt)
@tf.function
def train_thetas(self, x, y):
self._train(x, y, self.thetas, self.theta_opt, training=False)
@property
def training_accuracy(self):
return self.accuracy_metric.result().numpy()
@property
def training_loss(self):
return self.loss_metric.result().numpy()
@tf.function
def predict(self, x):
y_hat = self.fbnet(x, self.temperature, training=False)
return y_hat
def evaluate(self, dataset):
accuracy_metric = SparseCategoricalAccuracy()
for x, y in dataset:
y_hat = self.predict(x)
accuracy_metric.update_state(y, y_hat)
return accuracy_metric.result().numpy()
def sample_sequential_config(self):
ops = [
op.sample(self.temperature)
if isinstance(op, MixedOperation) else op for op in self.fbnet.ops
]
sequential_config = {
'name':
'sampled_fbnet',
'layers': [{
'class_name': type(op).__name__,
'config': op.get_config()
} for op in ops if not isinstance(op, Identity)]
}
return sequential_config
def save_weights(self, checkpoint):
self.fbnet.save_weights(checkpoint, save_format='tf')
def load_weights(self, checkpoint):
self.fbnet.load_weights(checkpoint)
class DualStudent(Model):
""""
Dual Student for Automatic Speech Recognition (ASR).
How to train: 1) set the optimizer by means of compile(), 2) use train()
How to test: use test()
Remarks:
- Do not use fit() by Keras, use train()
- Do not use evaluate() by Keras, use test()
- Compiled metrics and loss (i.e. set by means of compile()) are not used
Original proposal for image classification: https://arxiv.org/abs/1909.01804
"""
def __init__(self,
n_classes,
n_hidden_layers=3,
n_units=96,
consistency_loss='mse',
consistency_scale=10,
stabilization_scale=100,
xi=0.6,
padding_value=0.,
sigma=0.01,
schedule='rampup',
schedule_length=5,
version='mono_directional'):
"""
Constructs a Dual Student model.
:param n_classes: number of classes (i.e. number of units in the last layer of each student)
:param n_hidden_layers: number of hidden layers in each student (i.e. LSTM layers)
:param n_units: number of units for each hidden layer
:param consistency_loss: one of 'mse', 'kl'
:param consistency_scale: maximum value of weight for consistency constraint
:param stabilization_scale: maximum value of weight for stabilization constraint
:param xi: threshold for stable sample
:param padding_value: value used to pad input sequences (used as mask_value for Masking layer)
:param sigma: standard deviation for noisy augmentation
:param schedule: type of schedule for lambdas, one of 'rampup', 'triangular_cycling', 'sinusoidal_cycling'
:param schedule_length:
:param version: one of:
- 'mono_directional': both students have mono-directional LSTM layers
- 'bidirectional: both students have bidirectional LSTM layers
- 'imbalanced': one student has mono-directional LSTM layers, the other one bidirectional
"""
super(DualStudent, self).__init__()
# store parameters
self.n_classes = n_classes
self.padding_value = padding_value
self.n_units = n_units
self.n_hidden_layers = n_hidden_layers
self.xi = xi
self.consistency_scale = consistency_scale
self.stabilization_scale = stabilization_scale
self.sigma = sigma
self.version = version
self.schedule = schedule
self.schedule_length = schedule_length
self._lambda1 = None
self._lambda2 = None
# schedule for lambdas
if schedule == 'rampup':
self.schedule_fn = sigmoid_rampup
elif schedule == 'triangular_cycling':
self.schedule_fn = triangular_cycling
elif schedule == 'sinusoidal_cycling':
self.schedule_fn = sinusoidal_cycling
else:
raise ValueError('Invalid schedule')
# loss
self._loss_cls = SparseCategoricalCrossentropy() # classification loss
self._loss_sta = MeanSquaredError() # stabilization loss
if consistency_loss == 'mse':
self._loss_con = MeanSquaredError() # consistency loss
elif consistency_loss == 'kl':
self._loss_con = KLDivergence()
else:
raise ValueError('Invalid consistency metric')
# metrics for training
self._loss1 = Mean(
name='loss1') # we want to average the loss for each batch
self._loss2 = Mean(name='loss2')
self._loss1_cls = Mean(name='loss1_cls')
self._loss2_cls = Mean(name='loss2_cls')
self._loss1_con = Mean(name='loss1_con')
self._loss2_con = Mean(name='loss2_con')
self._loss1_sta = Mean(name='loss1_sta')
self._loss2_sta = Mean(name='loss2_sta')
self._acc1 = SparseCategoricalAccuracy(name='acc1')
self._acc2 = SparseCategoricalAccuracy(name='acc2')
# metrics for testing
self._test_loss1 = Mean(name='test_loss1')
self._test_loss2 = Mean(name='test_loss2')
self._test_acc1_train_phones = SparseCategoricalAccuracy(
name='test_acc1_train_phones')
self._test_acc2_train_phones = SparseCategoricalAccuracy(
name='test_acc2_train_phones')
self._test_acc1 = Accuracy(name='test_acc1')
self._test_acc2 = Accuracy(name='test_acc2')
self._test_per1 = PhoneErrorRate(name='test_per1')
self._test_per2 = PhoneErrorRate(name='test_per2')
# compose students
if version == 'mono_directional':
lstm_types = ['mono_directional', 'mono_directional']
elif version == 'bidirectional':
lstm_types = ['bidirectional', 'bidirectional']
elif version == 'imbalanced':
lstm_types = ['mono_directional', 'bidirectional']
else:
raise ValueError('Invalid student version')
self.student1 = self._get_student('student1', lstm_types[0])
self.student2 = self._get_student('student2', lstm_types[1])
# masking layer (just to use compute_mask and remove padding)
self.mask = Masking(mask_value=self.padding_value)
def _get_student(self, name, lstm_type):
student = Sequential(name=name)
student.add(Masking(mask_value=self.padding_value))
if lstm_type == 'mono_directional':
for i in range(self.n_hidden_layers):
student.add(LSTM(units=self.n_units, return_sequences=True))
elif lstm_type == 'bidirectional':
for i in range(self.n_hidden_layers):
student.add(
Bidirectional(
LSTM(units=self.n_units, return_sequences=True)))
else:
raise ValueError('Invalid LSTM version')
student.add(Dense(units=self.n_classes, activation="softmax"))
return student
def _noisy_augment(self, x):
return x + tf.random.normal(shape=x.shape, stddev=self.sigma)
def call(self, inputs, training=False, student='student1', **kwargs):
"""
Feed-forwards inputs to one of the students.
This function is called internally by __call__(). Do not use it directly, use the model as callable. You may
prefer to use pad_and_predict() instead of this, because it pads the sequences and splits in batches. For a big
dataset, it is strongly suggested that you use pad_and_predict().
:param inputs: tensor of shape (batch_size, n_frames, n_features)
:param training: boolean, whether the call is in inference mode or training mode
:param student: one of 'student1', 'student2'
:return: tensor of shape (batch_size, n_frames, n_classes), softmax activations (probabilities)
"""
if student == 'student1':
return self.student1(inputs, training=training)
elif student != 'student1':
return self.student2(inputs, training=training)
else:
raise ValueError('Invalid student')
def build(self, input_shape):
super(DualStudent, self).build(input_shape)
self.student1.build(input_shape)
self.student2.build(input_shape)
def train(self,
x_labeled,
x_unlabeled,
y_labeled,
x_val=None,
y_val=None,
n_epochs=10,
batch_size=32,
shuffle=True,
evaluation_mapping=None,
logs_path=None,
checkpoints_path=None,
initial_epoch=0,
seed=None):
"""
Trains the students with both labeled and unlabeled data (semi-supervised learning).
:param x_labeled: numpy array of numpy arrays (n_frames, n_features), features corresponding to y_labeled.
'n_frames' can vary, padding is added to make x_labeled a tensor.
:param x_unlabeled: numpy array of numpy arrays of shape (n_frames, n_features), features without labels.
'n_frames' can vary, padding is added to make x_unlabeled a tensor.
:param y_labeled: numpy array of numpy arrays of shape (n_frames,), labels corresponding to x_labeled.
'n_frames' can vary, padding is added to make y_labeled a tensor.
:param x_val: like x_labeled, but for validation set
:param y_val: like y_labeled, but for validation set
:param n_epochs: integer, number of training epochs
:param batch_size: integer, batch size
:param shuffle: boolean, whether to shuffle at each epoch or not
:param evaluation_mapping: dictionary {training label -> test label}, the test phones should be a subset of the
training phones
:param logs_path: path where to save logs for TensorBoard
:param checkpoints_path: path to a directory. If the directory contains checkpoints, the latest checkpoint is
restored.
:param initial_epoch: int, initial epoch from which to start the training. It can be used together with
checkpoints_path to resume the training from a previous run.
:param seed: seed for the random number generator
"""
# set seed
if seed is not None:
np.random.seed(seed)
tf.random.set_seed(seed)
# show summary
self.build(input_shape=(None, ) + x_labeled[0].shape)
self.student1.summary()
self.student2.summary()
# setup for logs
train_summary_writer = None
if logs_path is not None:
train_summary_writer = tf.summary.create_file_writer(logs_path)
# setup for checkpoints
checkpoint = None
if checkpoints_path is not None:
checkpoint = tf.train.Checkpoint(optimizer=self.optimizer,
model=self)
checkpoint_path = tf.train.latest_checkpoint(checkpoints_path)
if checkpoint_path is not None:
checkpoint.restore(checkpoint_path)
checkpoint_path = Path(checkpoints_path) / 'ckpt'
checkpoint_path = str(checkpoint_path)
# compute batch sizes
labeled_batch_size = ceil(
len(x_labeled) / (len(x_unlabeled) + len(x_labeled)) * batch_size)
unlabeled_batch_size = batch_size - labeled_batch_size
n_batches = min(ceil(len(x_unlabeled) / unlabeled_batch_size),
ceil(len(x_labeled) / labeled_batch_size))
# training loop
for epoch in trange(initial_epoch, n_epochs, desc='epochs'):
# ramp up lambda1 and lambda2
self._lambda1 = self.consistency_scale * self.schedule_fn(
epoch, self.schedule_length)
self._lambda2 = self.stabilization_scale * self.schedule_fn(
epoch, self.schedule_length)
# shuffle training set
if shuffle:
indices = np.arange(
len(x_labeled)
) # get indices to shuffle coherently features and labels
np.random.shuffle(indices)
x_labeled = x_labeled[indices]
y_labeled = y_labeled[indices]
np.random.shuffle(x_unlabeled)
for i in trange(n_batches, desc='batches'):
# select batch
x_labeled_batch = select_batch(x_labeled, i,
labeled_batch_size)
x_unlabeled_batch = select_batch(x_unlabeled, i,
unlabeled_batch_size)
y_labeled_batch = select_batch(y_labeled, i,
labeled_batch_size)
# pad batch
x_labeled_batch = pad_sequences(x_labeled_batch,
padding='post',
value=self.padding_value,
dtype='float32')
x_unlabeled_batch = pad_sequences(x_unlabeled_batch,
padding='post',
value=self.padding_value,
dtype='float32')
y_labeled_batch = pad_sequences(y_labeled_batch,
padding='post',
value=-1)
# convert to tensors
x_labeled_batch = tf.convert_to_tensor(x_labeled_batch)
x_unlabeled_batch = tf.convert_to_tensor(x_unlabeled_batch)
y_labeled_batch = tf.convert_to_tensor(y_labeled_batch)
# train step
self._train_step(x_labeled_batch, x_unlabeled_batch,
y_labeled_batch)
# put metrics in dictionary (easy management)
train_metrics = {
self._loss1.name: self._loss1.result(),
self._loss2.name: self._loss2.result(),
self._loss1_cls.name: self._loss1_cls.result(),
self._loss2_cls.name: self._loss2_cls.result(),
self._loss1_con.name: self._loss1_con.result(),
self._loss2_con.name: self._loss2_con.result(),
self._loss1_sta.name: self._loss1_sta.result(),
self._loss2_sta.name: self._loss2_sta.result(),
self._acc1.name: self._acc1.result(),
self._acc2.name: self._acc2.result(),
}
metrics = {'train': train_metrics}
# test on validation set
if x_val is not None and y_val is not None:
val_metrics = self.test(x_val,
y_val,
evaluation_mapping=evaluation_mapping)
metrics['val'] = val_metrics
# print metrics
for dataset, metrics_ in metrics.items():
print(f'Epoch {epoch + 1} - ', dataset, ' - ', sep='', end='')
for k, v in metrics_.items():
print(f'{k}: {v}, ', end='')
print()
# save logs
if train_summary_writer is not None:
with train_summary_writer.as_default():
for dataset, metrics_ in metrics.items():
for k, v in metrics_.items():
tf.summary.scalar(k, v, step=epoch)
# save checkpoint
if checkpoint is not None:
checkpoint.save(file_prefix=checkpoint_path)
# reset metrics
self._loss1.reset_states()
self._loss2.reset_states()
self._loss1_cls.reset_states()
self._loss2_cls.reset_states()
self._loss1_con.reset_states()
self._loss2_con.reset_states()
self._loss1_sta.reset_states()
self._loss2_sta.reset_states()
self._acc1.reset_states()
self._acc2.reset_states()
"""
If you want to use graph execution, pad the whole dataset externally and uncomment the decorator below.
If you uncomment the decorator without padding the dataset, the graph will be compiled for each batch,
because train() pads at batch level and so the batches have different shapes. This would result in worse
performance compared to eager execution.
"""
# @tf.function
def _train_step(self, x_labeled, x_unlabeled, y_labeled):
# noisy augmented batches (TODO: improvement with data augmentation instead of noise)
B1_labeled = self._noisy_augment(x_labeled)
B2_labeled = self._noisy_augment(x_labeled)
B1_unlabeled = self._noisy_augment(x_unlabeled)
B2_unlabeled = self._noisy_augment(x_unlabeled)
# compute masks (to remove padding)
mask_labeled = self.mask.compute_mask(x_labeled)
mask_unlabeled = self.mask.compute_mask(x_unlabeled)
y_labeled = y_labeled[mask_labeled] # remove padding from labels
# forward pass
with tf.GradientTape(persistent=True) as tape:
# predict augmented labeled samples (for classification and consistency constraint)
prob1_labeled_B1 = self.student1(B1_labeled, training=True)
prob1_labeled_B2 = self.student1(B2_labeled, training=True)
prob2_labeled_B1 = self.student2(B1_labeled, training=True)
prob2_labeled_B2 = self.student2(B2_labeled, training=True)
# predict augmented unlabeled samples (for consistency and stabilization constraints)
prob1_unlabeled_B1 = self.student1(B1_unlabeled, training=True)
prob1_unlabeled_B2 = self.student1(B2_unlabeled, training=True)
prob2_unlabeled_B1 = self.student2(B1_unlabeled, training=True)
prob2_unlabeled_B2 = self.student2(B2_unlabeled, training=True)
# remove padding
prob1_labeled_B1 = prob1_labeled_B1[mask_labeled]
prob1_labeled_B2 = prob1_labeled_B2[mask_labeled]
prob2_labeled_B1 = prob2_labeled_B1[mask_labeled]
prob2_labeled_B2 = prob2_labeled_B2[mask_labeled]
prob1_unlabeled_B1 = prob1_unlabeled_B1[mask_unlabeled]
prob1_unlabeled_B2 = prob1_unlabeled_B2[mask_unlabeled]
prob2_unlabeled_B1 = prob2_unlabeled_B1[mask_unlabeled]
prob2_unlabeled_B2 = prob2_unlabeled_B2[mask_unlabeled]
# compute classification losses
L1_cls = self._loss_cls(y_labeled, prob1_labeled_B1)
L2_cls = self._loss_cls(y_labeled, prob2_labeled_B2)
# concatenate labeled and unlabeled probability predictions (for consistency loss)
prob1_labeled_unlabeled_B1 = tf.concat(
[prob1_labeled_B1, prob1_unlabeled_B1], axis=0)
prob1_labeled_unlabeled_B2 = tf.concat(
[prob1_labeled_B2, prob1_unlabeled_B2], axis=0)
prob2_labeled_unlabeled_B1 = tf.concat(
[prob2_labeled_B1, prob2_unlabeled_B1], axis=0)
prob2_labeled_unlabeled_B2 = tf.concat(
[prob2_labeled_B2, prob2_unlabeled_B2], axis=0)
# compute consistency losses
L1_con = self._loss_con(prob1_labeled_unlabeled_B1,
prob1_labeled_unlabeled_B2)
L2_con = self._loss_con(prob2_labeled_unlabeled_B1,
prob2_labeled_unlabeled_B2)
# prediction
P1_unlabeled_B1 = tf.argmax(prob1_unlabeled_B1, axis=-1)
P1_unlabeled_B2 = tf.argmax(prob1_unlabeled_B2, axis=-1)
P2_unlabeled_B1 = tf.argmax(prob2_unlabeled_B1, axis=-1)
P2_unlabeled_B2 = tf.argmax(prob2_unlabeled_B2, axis=-1)
# confidence (probability of predicted class)
M1_unlabeled_B1 = tf.reduce_max(prob1_unlabeled_B1, axis=-1)
M1_unlabeled_B2 = tf.reduce_max(prob1_unlabeled_B2, axis=-1)
M2_unlabeled_B1 = tf.reduce_max(prob2_unlabeled_B1, axis=-1)
M2_unlabeled_B2 = tf.reduce_max(prob2_unlabeled_B2, axis=-1)
# stable samples (masks to index probabilities)
R1 = tf.logical_and(
P1_unlabeled_B1 == P1_unlabeled_B2,
tf.logical_or(M1_unlabeled_B1 > self.xi,
M1_unlabeled_B2 > self.xi))
R2 = tf.logical_and(
P2_unlabeled_B1 == P2_unlabeled_B2,
tf.logical_or(M2_unlabeled_B1 > self.xi,
M2_unlabeled_B2 > self.xi))
R12 = tf.logical_and(R1, R2)
# stabilities
epsilon1 = MSE(prob1_unlabeled_B1[R12], prob1_unlabeled_B2[R12])
epsilon2 = MSE(prob2_unlabeled_B1[R12], prob2_unlabeled_B2[R12])
# compute stabilization losses
L1_sta = self._loss_sta(
prob1_unlabeled_B1[R12][epsilon1 > epsilon2],
prob2_unlabeled_B1[R12][epsilon1 > epsilon2])
L2_sta = self._loss_sta(
prob1_unlabeled_B2[R12][epsilon1 < epsilon2],
prob2_unlabeled_B2[R12][epsilon1 < epsilon2])
L1_sta += self._loss_sta(
prob1_unlabeled_B1[tf.logical_and(tf.logical_not(R1), R2)],
prob2_unlabeled_B1[tf.logical_and(tf.logical_not(R1), R2)])
L2_sta += self._loss_sta(
prob1_unlabeled_B2[tf.logical_and(R1, tf.logical_not(R2))],
prob2_unlabeled_B2[tf.logical_and(R1, tf.logical_not(R2))])
# compute complete losses
L1 = L1_cls + self._lambda1 * L1_con + self._lambda2 * L1_sta
L2 = L2_cls + self._lambda1 * L2_con + self._lambda2 * L2_sta
# backward pass
gradients1 = tape.gradient(L1, self.student1.trainable_variables)
gradients2 = tape.gradient(L2, self.student2.trainable_variables)
self.optimizer.apply_gradients(
zip(gradients1, self.student1.trainable_variables))
self.optimizer.apply_gradients(
zip(gradients2, self.student2.trainable_variables))
del tape # to release memory (persistent tape)
# update metrics
self._loss1.update_state(L1)
self._loss2.update_state(L2)
self._loss1_cls.update_state(L1_cls)
self._loss2_cls.update_state(L2_cls)
self._loss1_con.update_state(L1_con)
self._loss2_con.update_state(L2_con)
self._loss1_sta.update_state(L1_sta)
self._loss2_sta.update_state(L2_sta)
self._acc1.update_state(y_labeled, prob1_labeled_B1)
self._acc2.update_state(y_labeled, prob2_labeled_B2)
def test(self, x, y, batch_size=32, evaluation_mapping=None):
"""
Tests the model (both students).
:param x: numpy array of numpy arrays (n_frames, n_features), features corresponding to y_labeled.
'n_frames' can vary, padding is added to make x a tensor.
:param y: numpy array of numpy arrays of shape (n_frames,), labels corresponding to x_labeled.
'n_frames' can vary, padding is added to make y a tensor.
:param batch_size: integer, batch size
:param evaluation_mapping: dictionary {training label -> test label}, the test phones should be a subset of the
training phones
:return: dictionary {metric_name -> value}
"""
# test batch by batch
n_batches = ceil(len(x) / batch_size)
for i in trange(n_batches, desc='test batches'):
# select batch
x_batch = select_batch(x, i, batch_size)
y_batch = select_batch(y, i, batch_size)
# pad batch
x_batch = pad_sequences(x_batch,
padding='post',
value=self.padding_value,
dtype='float32')
y_batch = pad_sequences(y_batch, padding='post', value=-1)
# convert to tensors
x_batch = tf.convert_to_tensor(x_batch)
y_batch = tf.convert_to_tensor(y_batch)
# test step
self._test_step(x_batch, y_batch, evaluation_mapping)
# put metrics in dictionary (easy management)
test_metrics = {
self._test_loss1.name:
self._test_loss1.result(),
self._test_loss2.name:
self._test_loss2.result(),
self._test_acc1_train_phones.name:
self._test_acc1_train_phones.result(),
self._test_acc2_train_phones.name:
self._test_acc2_train_phones.result(),
self._test_acc1.name:
self._test_acc1.result(),
self._test_acc2.name:
self._test_acc2.result(),
self._test_per1.name:
self._test_per1.result(),
self._test_per2.name:
self._test_per2.result(),
}
# reset metrics
self._test_loss1.reset_states()
self._test_loss2.reset_states()
self._test_acc1_train_phones.reset_states()
self._test_acc2_train_phones.reset_states()
self._test_acc1.reset_states()
self._test_acc2.reset_states()
self._test_per1.reset_states()
self._test_per2.reset_states()
return test_metrics
# @tf.function # see note in _train_step()
def _test_step(self, x, y, evaluation_mapping):
# compute mask (to remove padding)
mask = self.mask.compute_mask(x)
# forward pass
y_prob1_train_phones = self.student1(x, training=False)
y_prob2_train_phones = self.student2(x, training=False)
y_pred1_train_phones = tf.argmax(y_prob1_train_phones, axis=-1)
y_pred2_train_phones = tf.argmax(y_prob2_train_phones, axis=-1)
y_train_phones = tf.identity(y)
# map labels to set of test phones
if evaluation_mapping is not None:
y = tf.numpy_function(map_labels,
[y_train_phones, evaluation_mapping],
[tf.float32])
y_pred1 = tf.numpy_function(
map_labels, [y_pred1_train_phones, evaluation_mapping],
[tf.float32])
y_pred2 = tf.numpy_function(
map_labels, [y_pred2_train_phones, evaluation_mapping],
[tf.float32])
else:
y = y_train_phones
y_pred1 = y_pred1_train_phones
y_pred2 = y_pred2_train_phones
# update phone error rate
self._test_per1.update_state(y, y_pred1, mask)
self._test_per2.update_state(y, y_pred2, mask)
# remove padding
y_pred1 = y_pred1[mask]
y_pred2 = y_pred2[mask]
y_prob1_train_phones = y_prob1_train_phones[mask]
y_prob2_train_phones = y_prob2_train_phones[mask]
y_train_phones = y_train_phones[mask]
y = y[mask]
# compute loss
loss1 = self._loss_cls(y_train_phones, y_prob1_train_phones)
loss2 = self._loss_cls(y_train_phones, y_prob2_train_phones)
# update loss
self._test_loss1.update_state(loss1)
self._test_loss2.update_state(loss2)
# update accuracy using training phones
self._test_acc1_train_phones.update_state(y_train_phones,
y_prob1_train_phones)
self._test_acc2_train_phones.update_state(y_train_phones,
y_prob2_train_phones)
# update accuracy using test phones
self._test_acc1.update_state(y, y_pred1)
self._test_acc2.update_state(y, y_pred2)
def low_level_train(optimizer, yolo_loss, train_datasets, valid_datasets, train_steps, valid_steps):
"""
以底层的方式训练，这种方式更好地观察训练过程，监视变量的变化
:param optimizer: 优化器
:param yolo_loss: 自定义的loss function
:param train_datasets: 以tf.data封装好的训练集数据
:param valid_datasets: 验证集数据
:param train_steps: 迭代一个epoch的轮次
:param valid_steps: 同上
:return: None
"""
# 创建模型结构
model = yolo_body()
# 定义模型评估指标
train_loss = Mean(name='train_loss')
valid_loss = Mean(name='valid_loss')
# 设置保存最好模型的指标
best_test_loss = float('inf')
patience = 10
min_delta = 1e-3
patience_cnt = 0
history_loss = []
# 创建summary
summary_writer = tf.summary.create_file_writer(logdir=cfg.log_dir)
# low level的方式计算loss
for epoch in range(1, cfg.epochs + 1):
train_loss.reset_states()
valid_loss.reset_states()
step = 0
print("Epoch {}/{}".format(epoch, cfg.epochs))
# 处理训练集数据
for batch, (images, labels) in enumerate(train_datasets.take(train_steps)):
with tf.GradientTape() as tape:
# 得到预测
outputs = model(images, training=True)
# 计算损失(注意这里收集model.losses的前提是Conv2D的kernel_regularizer参数)
regularization_loss = tf.reduce_sum(model.losses)
pred_loss = []
# yolo_loss、label、output都是3个特征层的数据，通过for 拆包之后，一个loss_fn就是yolo_loss中一个特征层
# 然后逐一计算,
for output, label, loss_fn in zip(outputs, labels, yolo_loss):
pred_loss.append(loss_fn(label, output))
# 总损失 = yolo损失 + 正则化损失
total_train_loss = tf.reduce_sum(pred_loss) + regularization_loss
# 反向传播梯度下降
# model.trainable_variables代表把loss反向传播到每个可以训练的变量中
grads = tape.gradient(total_train_loss, model.trainable_variables)
# 将每个节点的误差梯度gradients，用于更新该节点的可训练变量值
# zip是把梯度和可训练变量值打包成元组
optimizer.apply_gradients(zip(grads, model.trainable_variables))
# 更新train_loss
train_loss.update_state(total_train_loss)
# 输出训练过程
rate = (step + 1) / train_steps
a = "*" * int(rate * 70)
b = "." * int((1 - rate) * 70)
loss = train_loss.result().numpy()
print("\r{}/{} {:^3.0f}%[{}->{}] - loss:{:.4f}".
format(batch, train_steps, int(rate * 100), a, b, loss), end='')
step += 1
# 计算验证集
for batch, (images, labels) in enumerate(valid_datasets.take(valid_steps)):
# 得到预测，不training
outputs = model(images)
regularization_loss = tf.reduce_sum(model.losses)
pred_loss = []
for output, label, loss_fn in zip(outputs, labels, yolo_loss):
pred_loss.append(loss_fn(label, output))
total_valid_loss = tf.reduce_sum(pred_loss) + regularization_loss
# 更新valid_loss
valid_loss.update_state(total_valid_loss)
print('\nLoss: {:.4f}, Test Loss: {:.4f}\n'.format(train_loss.result(), valid_loss.result()))
# 保存loss，可以选择train的loss
history_loss.append(valid_loss.result().numpy())
# 保存到tensorboard里
with summary_writer.as_default():
tf.summary.scalar('train_loss', train_loss.result(), step=optimizer.iterations)
tf.summary.scalar('valid_loss', valid_loss.result(), step=optimizer.iterations)
# 只保存最好模型
if valid_loss.result() < best_test_loss:
best_test_loss = valid_loss.result()
model.save_weights(cfg.model_path, save_format='tf')
# EarlyStopping
if epoch > 1 and history_loss[epoch - 2] - history_loss[epoch - 1] > min_delta:
patience_cnt = 0
else:
patience_cnt += 1
if patience_cnt >= patience:
tf.print("No improvement for {} times, early stopping optimization.".format(patience))
break
class VFAE(keras.Model):
def __init__(self,
encoder,
encoder_z,
reconstructor_z,
decoder,
classifier,
feature_dim,
loss_type,
**kwargs):
super(VFAE, self).__init__(**kwargs)
self.eps = tf.constant([10e-25])
self.beta=1.
self.encoder = encoder
self.encoder_z = encoder_z
self.reconstructor_z = reconstructor_z
self.decoder = decoder
self.classifier = classifier
self.loss_type = loss_type
self.total_loss_tracker = Mean(name="total_loss")
self.prediction_loss_tracker = Mean(name="pred_loss")
self.kl_loss_tracker = Mean(name="kl_loss")
self.mmd_loss_tracker = Mean(name="mmd_loss")
self.reconst_loss_tracker = Mean(name="reconst_loss")
self.reconst_z_loss_tracker = Mean(name="reconst_z_loss")
@property
def metrics(self):
return [
self.total_loss_tracker,
self.prediction_loss_tracker,
self.kl_loss_tracker,
self.mmd_loss_tracker,
self.reconst_loss_tracker,
self.reconst_z_loss_tracker
]
def call(self, inputs):
X, y = inputs
y = tf.reshape(y, (-1,1))
sens, _ = split_sensitive_X(X, 0, 1)
z_mean, z_log_sigma, z = self.encoder(X)
q_z_1_mean, q_z_1_log_sigma, z_1 = self.encoder_z(tf.concat([z,y], axis=1))
reconst = self.decoder(tf.concat([z, sens], axis=1))
z_reconst_mean, z_reconst_log_sigma, _ = self.reconstructor_z(tf.concat([z_1, y], axis=1))
preds = self.classifier(z)
return z_mean, z_log_sigma, z, \
reconst, q_z_1_mean, q_z_1_log_sigma, z_1, \
z_reconst_mean, z_reconst_log_sigma, preds
def train_step(self, data):
X, y = data
with tf.GradientTape() as tape:
z_mean, z_log_sigma, z, \
reconst, q_z_1_mean, q_z_1_log_sigma, z_1, \
z_reconst_mean, z_reconst_log_sigma, preds = self.call(data)
reconst_loss = neg_log_bernoulli(X, reconst, rec=1)
reconst_z_loss = negative_log_gaussian(z, z_reconst_mean, z_reconst_log_sigma)
classifier_loss = neg_log_bernoulli(y, preds)
kl_loss = KL(q_z_1_mean, q_z_1_log_sigma)
mmd_loss = mmd_loss(X, z)
entropy_z = entropy_gaussian(z_mean, z_log_sigma)
total_loss = reconst_loss + kl_loss + reconst_z_loss - entropy_z + self.beta*classifier_loss
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
self.total_loss_tracker.update_state(total_loss)
self.prediction_loss_tracker.update_state(classifier_loss)
self.kl_loss_tracker.update_state(kl_loss)
self.mmd_loss_tracker.update_state(mmd_loss)
self.reconst_loss_tracker.update_state(reconst_loss)
self.reconst_z_loss_tracker.update_state(reconst_z_loss)
return {
"loss": self.total_loss_tracker.result(),
"classification_loss": self.prediction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
"mmd_loss": self.mmd_loss_tracker.result(),
"reconst_loss": self.reconst_loss_tracker.result(),
"reconst_z_loss": self.reconst_z_loss_tracker.result()
}
class LFPNet(Model):
def __init__(self, encoder, planner, actor, beta) -> None:
super(LFPNet, self).__init__()
self.encoder = encoder
self.planner = planner
self.actor = actor
self.beta = beta
self.total_loss_tracker = Mean(name="total_loss")
self.action_loss_tracker = Mean(name="action_loss")
self.reg_loss_tracker = Mean(name="reg_loss")
def call(self, inputs, planner=True, training=False):
if planner:
z = self.planner(
[inputs['obs'][:, 0, :], inputs['goals'][:, 0, :]])
else:
z = self.encoder([inputs['obs'], inputs['acts']])
z_tiled = tf.tile(tf.expand_dims(z[0], 1),
(1, inputs['obs'].shape[1], 1))
acts = self.actor([inputs['obs'], z_tiled, inputs['goals']])
return acts, z
def train_step(self, inputs):
with tf.GradientTape() as tape:
acts_enc, z_enc = self(inputs, planner=False, training=True)
acts_plan, z_plan = self(inputs, planner=True, training=True)
act_loss = self.compiled_loss(inputs['acts'],
acts_enc,
regularization_losses=self.losses)
reg_loss = tfd.kl_divergence(z_enc, z_plan)
loss = act_loss + self.beta * reg_loss
gradients = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(gradients,
self.trainable_variables))
# Update metrics (includes the metric that tracks the loss)
self.total_loss_tracker.update_state(loss)
self.action_loss_tracker.update_state(act_loss)
self.reg_loss_tracker.update_state(reg_loss)
result = {m.name: m.result() for m in self.metrics}
result['beta'] = self.beta
return result
def test_step(self, inputs):
acts_enc, z_enc = self(inputs, planner=False, training=False)
acts_plan, z_plan = self(inputs, planner=True, training=False)
act_loss = self.compiled_loss(inputs['acts'],
acts_plan,
regularization_losses=self.losses)
reg_loss = tfd.kl_divergence(z_enc, z_plan)
loss = act_loss + self.beta * reg_loss
# Update metrics (includes the metric that tracks the loss)
self.total_loss_tracker.update_state(loss)
self.action_loss_tracker.update_state(act_loss)
self.reg_loss_tracker.update_state(reg_loss)
return {m.name: m.result() for m in self.metrics}
@property
def metrics(self):
return [
self.total_loss_tracker, self.action_loss_tracker,
self.reg_loss_tracker
]
class LatentStateAction(Model):
def __init__(self, n_obs, n_act, latent_dim, warmup_steps, **kwargs):
super(LatentStateAction, self).__init__(**kwargs)
self.action_encoder = Encoder(n_act,
latent_dim,
is_variational=False,
name_prefix='action_encoder')
self.state_encoder = Encoder(n_obs,
latent_dim,
name_prefix='state_encoder')
self.action_decoder = Decoder(n_act,
latent_dim,
name_prefix='action_decoder')
self.state_decoder = Decoder(n_obs,
latent_dim,
name_prefix='state_decoder')
self.action_input = Input(shape=(n_act, ))
self.action_output = self.action_decoder(
self.action_encoder(self.action_input))
self.action_ae_model = Model(self.action_input,
self.action_output,
name='action_ae')
self.state_input = Input(shape=(n_obs, ))
self.state_output = self.state_decoder(
self.state_encoder(self.state_input)[-1])
self.state_ae_model = Model(self.state_input,
self.state_output,
name='state_vae')
self.model = Model(inputs=[self.action_input, self.state_input],
outputs=[self.action_output, self.state_output])
self.compile(optimizer=Adam(learning_rate=1e-2))
# Hyper-parameters
self.warmup_steps = tf.constant(warmup_steps, dtype=tf.int32)
self.it = tf.Variable(0, dtype=tf.int32)
# Logging
self.action_loss_tracker = Mean(name='action_loss')
self.state_recon_loss_tracker = Mean(name='state_recon_loss')
self.state_kl_loss_tracker = Mean(name='state_kl_loss')
def call(self, inputs, training=None, mask=None):
return self.model([inputs['acts'], inputs['obs1']])
@property
def metrics(self):
return [
self.action_loss_tracker, self.state_recon_loss_tracker,
self.state_kl_loss_tracker
]
# @tf.function
def train_step(self, data):
data = data[0]
with tf.GradientTape(persistent=True) as tape:
# Get state encoder output
zs_mean, zs_log_var, zs = self.state_encoder(data['obs1'])
# How good are we at reconstructing state?
state_reconstruction = self.state_decoder(zs)
state_reconstruction_loss = tf.reduce_mean(
tf.square(data['obs1'] - state_reconstruction))
# How much regularized is our latent state space?
state_kl_loss = -0.5 * (1 + zs_log_var - tf.square(zs_mean) -
tf.exp(zs_log_var))
state_kl_loss = tf.reduce_mean(tf.reduce_sum(state_kl_loss,
axis=1))
# Mask state_kl_loss during warm-up phase
state_kl_loss = tf.cond(pred=tf.math.greater_equal(
self.it, self.warmup_steps),
true_fn=lambda: state_kl_loss,
false_fn=lambda: 0.0)
# Find state VAE total loss
total_state_loss = state_reconstruction_loss + 4 * state_kl_loss
# Get action encoder output
za = self.action_encoder(data['acts'])
# How good are we at reconstructing action?
action_reconstruction = self.action_decoder(za)
action_reconstruction_loss = tf.reduce_mean(
tf.square(data['acts'] - action_reconstruction))
# Get encoded next state using current state encoder
zs_tp1 = self.state_encoder(data['obs2'])[-1]
# Predict next state assuming canonical representation
pred_zs_tp1 = tf.add(za, zs_tp1)
# Get action AE total loss
latent_matching_loss = tf.reduce_mean(
tf.square(zs_tp1 - pred_zs_tp1))
total_action_loss = latent_matching_loss + action_reconstruction_loss
self.it.assign_add(1)
# Get partial derivative wrt to losses and network params
state_grads = tape.gradient(total_state_loss,
self.state_ae_model.trainable_weights)
action_grads = tape.gradient(total_action_loss,
self.action_ae_model.trainable_weights)
# Apply gradients
self.optimizer.apply_gradients(
zip(state_grads, self.state_ae_model.trainable_weights))
self.optimizer.apply_gradients(
zip(action_grads, self.action_ae_model.trainable_weights))
# Logging
self.action_loss_tracker.update_state(total_action_loss)
self.state_recon_loss_tracker.update_state(state_reconstruction_loss)
self.state_kl_loss_tracker.update_state(state_kl_loss)
return {
'action_loss': self.action_loss_tracker.result(),
'state_recon_loss': self.state_recon_loss_tracker.result(),
'state_kl_loss': self.state_kl_loss_tracker.result()
}
