TensorFlow自定义层与梯度实现报错：OperatorNotAllowedInGraphError

错误原因 #

你的自定义梯度函数grad返回单个张量，但当层包含可训练变量时，TensorFlow期望梯度函数返回与层可训练变量一一对应的梯度列表。原代码中，优化器尝试迭代你的返回值来匹配self.trainable_variables（此处为self.w），但图模式下不允许直接迭代张量，因此触发该错误。

修正后的代码 #

import tensorflow as tf
from tensorflow import keras
import sys

print("Python version")
print(sys.version)
print("Version info.")
print(sys.version_info)
print("Tensorflow version")
print(tf.__version__)

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )

    @tf.custom_gradient
    def call(self, inputs):
        def grad(dy, variables=None):
            # 返回与可训练变量对应的梯度列表，此处仅self.w，故返回单元素列表
            return [tf.matmul(tf.transpose(inputs), dy)]
        return tf.matmul(inputs, self.w), grad

model = tf.keras.models.Sequential([
    Linear(1),
])
model.compile(optimizer='sgd',loss='mean_squared_error')

xs = tf.constant([[-1.0],  [0.0], [1.0], [2.0], [3.0], [4.0]], dtype=float) 
print(model(xs))

ys = tf.constant([[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0]], dtype=float) 

model.fit(xs, ys, epochs=10)  

关键修改说明 #

1. 梯度返回格式调整：将grad函数的返回值从单个张量改为列表，列表元素与层中可训练变量一一对应。这里层仅含self.w，所以返回包含一个梯度张量的列表。
2. 梯度计算修正：线性层权重w的正确梯度应为输入转置与上游梯度dy的乘积（tf.transpose(inputs) @ dy），原代码的tf.matmul(inputs, dy)维度不匹配，修正后才能正确计算梯度并完成训练。

运行修正后的代码，模型将正常执行训练流程，不再触发原错误。

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