TensorFlow自定义train_step中计算损失时如何正确访问层权重？

问题场景 #

我在实现一篇论文的模型，损失函数需要用到某个Keras Dense层的权重矩阵。继承Keras.Model自定义train_step()方法时，在with tf.GradientTape() as tape:代码块里用self.fc_layer.get_weights()[0]获取权重得到的是NumPy数组，导致梯度胶带无法关联损失和可训练权重，报错：

ValueError: No gradients provided for any variable: (['conv1_conv/kernel:0', 'conv1_conv/bias:0', ...]

需求是在不脱离TensorFlow计算图的前提下访问self.fc_layer的权重，当前启用Eager Execution但代码要兼容图模式。额外说明：权重矩阵每行代表一个类中心，需计算其与嵌入的余弦相似度，以最大化车辆重识别任务中的类间差异。

解决方案 #

• 直接访问层的weights属性获取Tensor
  不要用get_weights()（该方法会把Tensor转换为NumPy数组，脱离计算图追踪），直接调用self.fc_layer.weights[0]，得到的是原生TensorFlow Variable，会被梯度胶带正常追踪。

• 示例代码片段

  class CustomModel(tf.keras.Model):
      def __init__(self, num_classes):
          super().__init__()
          # 特征提取骨干网络示例
          self.backbone = tf.keras.applications.ResNet50(include_top=False, pooling='avg')
          self.fc_layer = tf.keras.layers.Dense(num_classes)
  
      def train_step(self, data):
          x, y = data
          with tf.GradientTape() as tape:
              # 获取特征嵌入
              embeddings = self.backbone(x, training=True)
              # 直接获取fc层权重Tensor（weights[0]为kernel矩阵，weights[1]为bias）
              class_centers = self.fc_layer.weights[0]
              # 计算余弦相似度：先归一化再做点积
              embeddings_norm = tf.nn.l2_normalize(embeddings, axis=1)
              centers_norm = tf.nn.l2_normalize(class_centers, axis=1)
              cos_sim = tf.matmul(embeddings_norm, centers_norm, transpose_b=True)
              # 代入自定义损失逻辑（替换为论文中的损失公式）
              loss = self.compiled_loss(y, cos_sim, regularization_losses=self.losses)
  
          # 计算梯度并更新可训练权重
          trainable_vars = self.trainable_variables
          gradients = tape.gradient(loss, trainable_vars)
          self.optimizer.apply_gradients(zip(gradients, trainable_vars))
          # 更新训练指标
          self.compiled_metrics.update_state(y, cos_sim)
          return {m.name: m.result() for m in self.metrics}
  

• 关键注意点

  • self.fc_layer.weights返回包含kernel和bias的Variable列表，weights[0]即为所需的权重矩阵Tensor，全程在TensorFlow计算图内，梯度可被正常追踪。
  • 对权重的所有后续计算（如归一化）都使用TensorFlow原生操作，避免转换为NumPy数组，确保同时兼容Eager模式和图模式。

内容的提问来源于stack exchange，提问作者magmacollaris