In Tensorflow, how do you include a pre-trained model in a custom model? - tensorflow

Question

I've got a custom model that I've pre-trained in a separate notebook called vae which I've saved using vae.save().
I'm now looking to implement a model which is a full training pipeline containing the pre-trained vae and some other (including a ResNet-50 from the Tensorflow gallery).
In the Tensorflow documentation about Making new Layers and Models via subclassing it doesn't mention anything about including models in subclassed models.
My question is how do I include my pre-trained VAE in a new subclassed model?
I've already tried writing the code below, which works fine for training, but I haven't seen implementations like it elsewhere and I'm getting errors when trying to do reid.save() or reid.summary() on the model, so I presume there is a better way to do it:
class ReId(keras.Model):
def __init__(self, vae, num_ids, **kwargs):
super(ReId, self).__init__(**kwargs)
self.vae = vae
self.convex_combination = ConvexCombination()
self.resnet_50 = resnet50
self.glob_avg_pool = keras.layers.GlobalAveragePooling2D()
self.bnneck = keras.layers.BatchNormalization()
self.num_ids = num_ids
self.final_fc_layer = keras.layers.Dense(self.num_ids)
self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
self.vae_loss_tracker = keras.metrics.Mean(
name="vae_loss"
)
self.classification_loss_tracker = keras.metrics.Mean(name="classification_loss")
self.triplet_loss_tracker = keras.metrics.Mean(name="triplet_loss")
#property
def metrics(self):
return [
self.total_loss_tracker,
self.vae_loss_tracker,
# Classification loss is cross-entropy loss here
self.classification_loss_tracker,
self.triplet_loss_tracker,
]
def train_step(self, data):
x, y = data
y = tf.cast(y, tf.int32)
with tf.GradientTape() as tape:
# Calculate VAE loss
z_mean, z_log_var, z = self.vae.encoder(x)
reconstruction = self.vae.decoder(z)
reconstruction_loss = tf.reduce_mean(
keras.losses.mean_squared_error(x, reconstruction)
)
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))
# Calculate VAE loss
vae_loss = reconstruction_loss + (LAMBDA * kl_loss)
embedding_output = self.call(x)
bnneck_output = self.bnneck(embedding_output)
training_output = self.final_fc_layer(bnneck_output)
# Calculate triplet loss
triplet_loss = tfa.losses.triplet_semihard_loss(y_true=y, y_pred=embedding_output)
# Calculate cross-entropy loss
ce_loss = keras.losses.sparse_categorical_crossentropy(y, training_output)
# Sum up the losses
total_loss = triplet_loss + ce_loss + (TOTAL_LOSS_VAE_MULTIPLIER * vae_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.classification_loss_tracker.update_state(ce_loss)
self.triplet_loss_tracker.update_state(triplet_loss)
self.vae_loss_tracker.update_state(vae_loss)
return {
"loss": self.total_loss_tracker.result(),
"classification loss:": self.classification_loss_tracker.result(),
"triplet loss:": self.triplet_loss_tracker.result(),
"vae loss:": self.vae_loss_tracker.result(),
}
def get_config(self):
config = super(ReId, self).get_config()
config.update({"num_ids": self.num_ids})
return config
def call(self, data):
recon_img = self.vae(data)
convex_combo_input = [data, (data - recon_img)]
convex_combo_output = self.convex_combination(convex_combo_input)
final_output = self.glob_avg_pool(self.resnet_50(convex_combo_output))
return final_output