def __init__(self, cfg, data_shapes):
"""Constructs the autoencoder.
Args:
cfg: ConfigDict with model hyperparameters.
data_shapes: Dict of shapes of model input tensors, as returned by
datasets.get_sequence_dataset.
"""
input_sequence = tf.keras.Input(shape=data_shapes['image'][1:],
name='image')
image_shape = data_shapes['image'][1:]
keypoints, _ = vision.build_images_to_keypoints_net(
cfg, image_shape)(input_sequence)
reconstructed_sequence = vision.build_keypoints_to_images_net(
cfg, image_shape)([
keypoints, input_sequence[:, 0, Ellipsis], keypoints[:, 0,
Ellipsis]
])
super(Autoencoder, self).__init__(inputs=input_sequence,
outputs=reconstructed_sequence,
name='autoencoder')
self.add_loss(tf.nn.l2_loss(input_sequence - reconstructed_sequence))
def testImagesToKeypointsNetShapes(self):
model = vision.build_images_to_keypoints_net(
self.cfg, self.data_shapes['image'][1:])
images = tf.zeros((self.cfg.batch_size,) + self.data_shapes['image'][1:])
keypoints, heatmaps = model(images)
self.assertEqual(
keypoints.shape.as_list(),
[self.cfg.batch_size, self.time_steps, self.cfg.num_keypoints, 3])
self.assertEqual(
heatmaps.shape.as_list(),
[self.cfg.batch_size, self.time_steps, self.cfg.heatmap_width,
self.cfg.heatmap_width, 3])
def build_model(cfg, data_shapes):
"""Builds the complete model with image encoder plus dynamics model.
This architecture is meant for testing/illustration only.
Model architecture:
image_sequence --> keypoints --> reconstructed_image_sequence
|
V
dynamics_model --> predicted_keypoints
The model takes a [batch_size, timesteps, H, W, C] image sequence as input. It
"observes" all frames, detects keypoints, and reconstructs the images. The
dynamics model learns to predict future keypoints based on the detected
keypoints.
Args:
cfg: ConfigDict with model hyperparameters.
data_shapes: Dict of shapes of model input tensors, as returned by
datasets.get_sequence_dataset.
Returns:
tf.keras.Model object.
"""
input_shape_no_batch = data_shapes['image'][
1:] # Keras uses shape w/o batch.
input_images = tf.keras.Input(shape=input_shape_no_batch, name='image')
# Vision model:
observed_keypoints, _ = vision.build_images_to_keypoints_net(
cfg, input_shape_no_batch)(input_images)
keypoints_to_images_net = vision.build_keypoints_to_images_net(
cfg, input_shape_no_batch)
reconstructed_images = keypoints_to_images_net([
observed_keypoints, input_images[:, 0, Ellipsis],
observed_keypoints[:, 0, Ellipsis]
])
# Dynamics model:
observed_keypoints_stop = tf.keras.layers.Lambda(
tf.stop_gradient)(observed_keypoints)
dynamics_model = dynamics.build_vrnn(cfg)
predicted_keypoints, kl_divergence = dynamics_model(
observed_keypoints_stop)
model = tf.keras.Model(inputs=[input_images],
outputs=[
reconstructed_images, observed_keypoints,
predicted_keypoints
],
name='autoencoder')
# Losses:
image_loss = tf.nn.l2_loss(input_images - reconstructed_images)
# Normalize by batch size and sequence length:
image_loss /= tf.to_float(
tf.shape(input_images)[0] * tf.shape(input_images)[1])
model.add_loss(image_loss)
separation_loss = losses.temporal_separation_loss(
cfg, observed_keypoints[:, :cfg.observed_steps, Ellipsis])
model.add_loss(cfg.separation_loss_scale * separation_loss)
vrnn_coord_pred_loss = tf.nn.l2_loss(observed_keypoints_stop -
predicted_keypoints)
# Normalize by batch size and sequence length:
vrnn_coord_pred_loss /= tf.to_float(
tf.shape(input_images)[0] * tf.shape(input_images)[1])
model.add_loss(vrnn_coord_pred_loss)
kl_loss = tf.reduce_mean(kl_divergence) # Mean over batch and timesteps.
model.add_loss(cfg.kl_loss_scale * kl_loss)
return model