title: In Tensorflow, how do you include a pre-trained model in a custom model?
body: <p>I've got a custom model that I've pre-trained in a <a href="https://vehiclereidjupyternotebook.s3.eu-west-2.amazonaws.com/vae_source_code.html" rel="nofollow noreferrer">separate notebook</a> called <code>vae</code> which I've saved using <code>vae.save()</code>.</p>
  <p>I'm now looking to implement a model which is a full training pipeline containing the pre-trained <code>vae</code> and some other (including a ResNet-50 from the Tensorflow gallery).</p>
  <p>In the Tensorflow documentation about <a href="https://www.tensorflow.org/guide/keras/custom_layers_and_models" rel="nofollow noreferrer">Making new Layers and Models via subclassing</a> it doesn't mention anything about including models in subclassed models.</p>
  <p>My question is <strong>how do I include my pre-trained VAE in a new subclassed model?</strong></p>
  <p>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 <code>reid.save()</code> or <code>reid.summary()</code> on the model, so I presume there is a better way to do it:</p>
  <pre><code>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=&quot;total_loss&quot;)
          self.vae_loss_tracker = keras.metrics.Mean(
              name=&quot;vae_loss&quot;
          )
          self.classification_loss_tracker = keras.metrics.Mean(name=&quot;classification_loss&quot;)
          self.triplet_loss_tracker = keras.metrics.Mean(name=&quot;triplet_loss&quot;)
  
      @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 {
              &quot;loss&quot;: self.total_loss_tracker.result(),
              &quot;classification loss:&quot;: self.classification_loss_tracker.result(),
              &quot;triplet loss:&quot;: self.triplet_loss_tracker.result(),
              &quot;vae loss:&quot;: self.vae_loss_tracker.result(),
          }
  
      def get_config(self):
          config = super(ReId, self).get_config()
          config.update({&quot;num_ids&quot;: 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
  </code></pre>