如何在TensorFlow中计算损失函数相对于所有可训练参数的Hessian矩阵

问题根源 #

你遇到的梯度为None的核心原因是：手动拼接的source张量是model.trainable_variables的下游输出，损失函数的计算链路完全不依赖该source张量，不存在梯度传导路径，因此计算结果为空。不需要额外构造拼接的参数变量，直接基于原始可训练变量列表计算后做矩阵拼接即可。

正确实现代码 #

import tensorflow as tf

# 前置条件：已定义模型model、输入x
with tf.GradientTape(persistent=True) as outer_tape:
    # 内层tape计算一阶梯度
    with tf.GradientTape() as inner_tape:
        loss = tf.reduce_mean(model(x, training=True)**2)
    # 计算损失对所有可训练变量的一阶梯度列表
    grads = inner_tape.gradient(loss, model.trainable_variables)
    # 展平所有一阶梯度并拼接为一维向量
    grads_flat = tf.concat([tf.reshape(g, [-1]) for g in grads], axis=0)

# 逐变量计算一阶梯度对参数的雅可比块
hessian_blocks = []
for var in model.trainable_variables:
    jacob = outer_tape.jacobian(grads_flat, var)
    # 展平参数对应维度，得到统一结构的Hessian子块
    hessian_blocks.append(tf.reshape(jacob, (grads_flat.shape[0], -1)))

# 拼接所有子块得到完整的[N,N] Hessian矩阵，N为模型总参数量
hessian = tf.concat(hessian_blocks, axis=1)

注意事项 #

• 当模型参数量较大时，全量Hessian矩阵会占用极高显存，容易出现OOM问题，可根据需求仅计算指定参数对应的Hessian子块。
• 若不需要完整Hessian矩阵，可直接针对每个参数单独计算Hessian后按需使用，无需拼接全量矩阵。

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