TensorFlow 2.x中SGD优化器报错：Cannot iterate over a scalar tensor

运行《Machine Learning with TensorFlow》书中代码时，使用TensorFlow 2.x的tf.keras.optimizers.SGD优化器报错，但切换到兼容版本的tf.compat.v1.train.GradientDescentOptimizer可正常运行。

最小复现代码 #

import tensorflow as tf
import numpy as np

X = tf.constant(np.linspace(-1, 1, 101), dtype=tf.float32)
Y = tf.constant(np.linspace(-1, 1, 101), dtype=tf.float32)
w = tf.Variable(0., name="weights", dtype=tf.float32)

cost = lambda: tf.square(Y - tf.multiply(X, w))

train_op =  tf.keras.optimizers.SGD(0.01)
#train_op =  tf.compat.v1.train.GradientDescentOptimizer(0.01)  # works with this optimizer

train_op.minimize(cost, w)   # error is here

报错信息 #

Traceback (most recent call last):
File "/home/alex/tmp/test.py", line 15, in
train_op.minimize(cost, w)
File "/home/alex/.local/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 526, in minimize
grads_and_vars = self.compute_gradients(loss, var_list, tape)
File "/home/alex/.local/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 260, in compute_gradients
return list(zip(grads, var_list))
File "/home/alex/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 583, in iter
raise TypeError("Cannot iterate over a scalar tensor.")
TypeError: Cannot iterate over a scalar tensor.

环境：Python 3.10.6，TensorFlow 2.11.0

报错核心原因是TensorFlow 2.x的tf.keras.optimizers.SGD要求损失函数返回标量值，但当前cost函数返回的是形状为(101,)的张量（每个样本的平方损失），不是标量。旧版兼容优化器会隐式处理这种情况，新版则需要显式将损失转为标量。

方法1：将损失转为标量（求平均或求和） #

修改cost函数，对所有样本的平方损失求平均或求和，得到标量损失：

# 方式A：计算平均损失
cost = lambda: tf.reduce_mean(tf.square(Y - tf.multiply(X, w)))
# 方式B：计算总损失
cost = lambda: tf.reduce_sum(tf.square(Y - tf.multiply(X, w)))

方法2：使用GradientTape手动计算梯度（TF2.x推荐模式） #

如果需要更灵活的训练流程，可以用tf.GradientTape手动计算梯度并更新变量，这是TensorFlow 2.x的标准写法：

import tensorflow as tf
import numpy as np

X = tf.constant(np.linspace(-1, 1, 101), dtype=tf.float32)
Y = tf.constant(np.linspace(-1, 1, 101), dtype=tf.float32)
w = tf.Variable(0., name="weights", dtype=tf.float32)

optimizer = tf.keras.optimizers.SGD(0.01)

# 单次训练步骤
with tf.GradientTape() as tape:
    # 先计算每个样本的损失，再转为标量
    loss = tf.reduce_mean(tf.square(Y - X * w))

# 计算梯度
grads = tape.gradient(loss, [w])
# 更新变量
optimizer.apply_gradients(zip(grads, [w]))

补充说明 #

• TF2.x的优化器设计更严谨，要求损失必须是标量，因为反向传播需要基于标量损失计算梯度。
• 旧版tf.compat.v1.train.GradientDescentOptimizer会自动对损失张量做求和处理，所以不会报错，但这种隐式行为在新版中被移除，必须显式处理损失的维度。

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