TensorFlow自定义梯度计算及apply_gradients接口传参问题咨询

tf.keras.optimizers.Optimizer.apply_gradients 只接收梯度与可训练变量一一配对的迭代对象，必须满足两个核心规则：

• 配对的梯度数量和model.trainable_variables的数量完全一致
• 每个梯度的形状必须和它配对的可训练变量的形状完全相同
• 不需要更新的变量，对应位置传形状一致的全0张量即可

现有代码问题排查 #

你当前定义的模型共有4个可训练变量，分别为：

• 第一层Dense(16)的卷积核：形状(1, 16)
• 第一层Dense(16)的偏置：形状(16,)
• 第二层Dense(2)的卷积核：形状(16, 2)
• 第二层Dense(2)的偏置：形状(2,)
  你直接传入形状为(2,)的单个常量梯度，数量、形状均不匹配，因此无法得到预期结果。

修复后实现 #

以下是符合你需求（仅对第一个输出应用梯度、第二个输出梯度置0）的可运行代码：

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Activation
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Nadam
import numpy as np

def cmodel():
    in_ = Input(shape=(1,))
    x = in_
    x = Dense(16)(x)
    x = Activation('relu')(x)
    x = Dense(2)(x)
    x = Activation('linear')(x)
    model = Model(inputs=in_, outputs=[x])
    return model

model = cmodel()
model.compile(optimizer=Nadam(learning_rate=0.005))

mylist2 = []
for i in range(100):
    # 按变量顺序构造对应梯度
    grads = [
        # 第一层参数不需要更新则传全0
        tf.zeros_like(model.trainable_variables[0]),
        tf.zeros_like(model.trainable_variables[1]),
        # 第二层卷积核：第一维梯度设为-0.3，第二维梯度置0
        tf.concat([tf.fill((16, 1), -0.3), tf.zeros((16, 1))], axis=-1),
        # 第二层偏置：第一个值梯度-0.3，第二个值梯度置0
        tf.constant([-0.3, 0.])
    ]
    model.optimizer.apply_gradients(zip(grads, model.trainable_variables))
    pred = model.predict(np.array([1]), verbose=0)
    print('num:', pred)
    mylist2.append(pred[0])

多输出训练通用方案 #

如果是正常训练场景需要自动计算梯度再掩蔽指定输出的梯度，可通过GradientTape计算梯度后修改对应维度：

loss_fn = tf.keras.losses.MSE

for step in range(100):
    train_x = np.array([1.])
    train_y = np.array([[4., 6.]]) # 两个输出的真实标签
    with tf.GradientTape() as tape:
        pred_y = model(train_x, training=True)
        loss = loss_fn(train_y, pred_y)
    # 计算全量梯度
    grads = tape.gradient(loss, model.trainable_variables)
    # 掩蔽第二个输出对应的梯度
    grads[2] = tf.concat([grads[2][:, :1], tf.zeros_like(grads[2][:, 1:])], axis=-1)
    grads[3] = tf.concat([grads[3][:1], tf.zeros_like(grads[3][1:])], axis=-1)
    # 应用梯度
    model.optimizer.apply_gradients(zip(grads, model.trainable_variables))

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