Tensorflow 2 + Keras 的知识蒸馏损失

Question

我正在尝试实现一个非常简单的 keras 模型，该模型使用来自另一个模型的知识蒸馏 [1]。粗略地说，我需要更换原来的损失L(y_true, y_pred)由L(y_true, y_pred)+L(y_teacher_pred, y_pred)哪里y_teacher_pred是另一个模型的预测。

我试过做

def create_student_model_with_distillation(teacher_model):

  inp = tf.keras.layers.Input(shape=(21,))

  model = tf.keras.models.Sequential()
  model.add(inp)

  model.add(...) 
  model.add(tf.keras.layers.Dense(units=1))

  teacher_pred = teacher_model(inp)

  def my_loss(y_true,y_pred):
      loss = tf.keras.losses.mean_squared_error(y_true, y_pred)
      loss += tf.keras.losses.mean_squared_error(teacher_pred, y_pred)
      return loss

  model.compile(loss=my_loss, optimizer='adam')

  return model

但是，当我尝试调用fit我的模型时，我得到了

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.

我该如何解决这个问题？

参考资料

[1] https://arxiv.org/abs/1503.02531

Answer 1

实际上，这篇博文可以回答您的问题：keras blog

但简而言之 - 您应该使用新的 TF2 API 并在块predict之前调用老师的tf.GradientTape()：

def train_step(self, data):
        # Unpack data
        x, y = data

        # Forward pass of teacher
        teacher_predictions = self.teacher(x, training=False)

        with tf.GradientTape() as tape:
            # Forward pass of student
            student_predictions = self.student(x, training=True)

            # Compute losses
            student_loss = self.student_loss_fn(y, student_predictions)
            distillation_loss = self.distillation_loss_fn(
                tf.nn.softmax(teacher_predictions / self.temperature, axis=1),
                tf.nn.softmax(student_predictions / self.temperature, axis=1),
            )
            loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss