为何TensorFlow实现Adam优化算法的训练效率远超手写版本？

我尝试自行实现Adam优化算法来解决线性回归问题并在数据集上运行。在超参数一致的情况下，我的手写版本需要约100个epochs才能达到可接受的损失值，而TensorFlow仅用3个epochs就将损失降至我的十分之一。此外需要说明的是，TensorFlow完成1个epoch的计算量，我的算法大约要执行130次相关运算。

手写Adam实现代码 #

import numpy as np
import math

def grad(x, y, w, b, par):
  if par == "w":
    return -2*x*(y-w*x-b)
  if par == "b":
    return -2*(y-w*x-b)

def cost(x, y, w, b):
    # 假设cost为均方误差计算
    return np.mean((y - w*x - b)**2)

def Adam(x, y, alpha, beta1=0.9, beta2=0.999, epsilon=10e-8, epochs=1000, batch_size=32):
  mw = 0
  vw = 0
  mb = 0
  vb = 0

  mmw = 0
  vvw = 0
  mmb = 0
  vvb = 0

  w = 0
  b = 0

  bestw = 0
  bestb = 0
  bestcost = cost(x, y, w, b)

  n = len(x)
  sw = int(n/batch_size)-1
  counter = 0
  
  for t in range(epochs):
    if t % sw == 0:
      arr = np.random.randint(0, n, batch_size)
    
    if t % 100 == 0:
      print(cost(x, y, w, b))

    if t % 15 == 0:
      if cost(x, y, w, b) < bestcost:
        bestw = w
        bestb = b
        bestcost = cost(x, y, w, b)

    for i in range(len(arr)):
      mw = beta1 * mw + (1 - beta1) * grad(x[arr[i]], y[arr[i]], w, b, 'w')
      vw = beta2 * vw + (1 - beta2) * ((grad(x[arr[i]], y[arr[i]], w, b, 'w'))**2)

      mmw = mw/(1 - beta1**(t+1))
      vvw = vw/(1 - beta2**(t+1))

      w = w - (alpha * mmw)/(math.sqrt(vvw)+epsilon)

      mb = beta1 * mb + (1 - beta1) * grad(x[arr[i]], y[arr[i]], w, b, 'b')
      vb = beta2 * vb + (1 - beta2) * ((grad(x[arr[i]], y[arr[i]], w, b, 'b'))**2)

      mmb = mb/(1 - beta1**(t+1))
      vvb = vb/(1 - beta2**(t+1))

      b = b - (alpha * mmb)/(math.sqrt(vvb)+epsilon)

  return bestw, bestb, bestcost # 返回最优参数与对应损失

TensorFlow实现代码 #

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(1)
])
model.compile(loss="mse", optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001))
model.fit(tf.expand_dims(x, axis=-1), y, epochs=100)

核心差异与原因分析 #

• 批量更新逻辑完全错位：
  TensorFlow的fit函数中，每个epoch会将整个数据集按batch_size划分成多个独立batch，每个batch计算一次梯度均值后完成一次参数更新。而你的手写代码仅在t % sw == 0时采样一个batch，后续多个epochs都重复使用这个batch的样本做单样本更新——相当于大部分时间在对同一小部分数据反复迭代，完全没有遍历整个数据集，收敛效率自然极低。

• 单样本更新vs批量更新的本质区别：
  你代码的内层循环是对batch中的每个样本单独计算梯度并更新参数，属于单样本SGD，梯度方差大、更新方向不稳定；而TensorFlow的Adam是小批量SGD，用整个batch的梯度均值做更新，梯度更平滑，收敛速度远快于单样本模式。

• Epoch定义完全不一致：
  你代码里的epochs是外层循环次数，每次循环可能只完成一次单样本更新；但TensorFlow的1个epoch是遍历完整数据集一次，包含n//batch_size次批量更新。这就是为什么你觉得“TensorFlow1个epoch的计算量等于我130次运算”——你的一次外层循环只是一次单样本更新，和TensorFlow的epoch完全不是一个量级。

• 偏差修正的步数计算错误：
  Adam的偏差修正需要基于参数更新的总步数，但你用的是外层epoch的计数t+1。比如每个epoch你做32次单样本更新，总步数应该是t*32 + i，但你一直用t+1计算修正分母，导致修正值偏差，直接影响收敛效率。

• 额外的全数据集评估损耗：
  你代码中每15个epoch就调用一次全数据集的cost函数评估，每100次循环还打印一次，这些全量计算会占用额外时间，进一步拉慢训练的感知速度。

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