使用核心 API 的优化器

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

简介

本笔记本介绍了使用 TensorFlow 核心低级 API 创建自定义优化器的过程。访问核心 API 概述了解有关 TensorFlow 核心及其预期用例的更多信息。

Keras 优化器模块是许多通用训练目的的推荐优化工具包。它包括各种预构建的优化器以及用于自定义的子类化功能。Keras 优化器还与使用核心 API 构建的自定义层、模型和训练循环兼容。这些预构建和可自定义的优化器适用于大多数情况，但核心 API 允许完全控制优化过程。例如，诸如 Sharpness-Aware Minimization (SAM) 之类的技术需要模型和优化器耦合，这与机器学习优化器的传统定义不符。本指南将逐步介绍使用核心 API 从头开始构建自定义优化器的过程，让您能够完全控制优化器的结构、实现和行为。

优化器概述

优化器是一种算法，用于针对模型的可训练参数最小化损失函数。最简单的优化技术是梯度下降，它通过在损失函数最陡下降的方向上迈出一步来迭代地更新模型的参数。它的步长与梯度的大小成正比，当梯度过大或过小时，这可能会出现问题。还有许多其他基于梯度的优化器，例如 Adam、Adagrad 和 RMSprop，它们利用梯度的各种数学特性来提高内存效率和快速收敛。

设置

import matplotlib
from matplotlib import pyplot as plt
# Preset Matplotlib figure sizes.
matplotlib.rcParams['figure.figsize'] = [9, 6]

import tensorflow as tf
print(tf.__version__)
# set random seed for reproducible results 
tf.random.set_seed(22)

梯度下降

基本优化器类应具有初始化方法和一个函数，用于在给定一组梯度的情况下更新一组变量。首先实现基本的梯度下降优化器，该优化器通过减去其梯度乘以学习率来更新每个变量。

class GradientDescent(tf.Module):

  def __init__(self, learning_rate=1e-3):
    # Initialize parameters
    self.learning_rate = learning_rate
    self.title = f"Gradient descent optimizer: learning rate={self.learning_rate}"

  def apply_gradients(self, grads, vars):
    # Update variables
    for grad, var in zip(grads, vars):
      var.assign_sub(self.learning_rate*grad)

为了测试此优化器，创建一个示例损失函数以针对单个变量 $x$ 进行最小化。计算其梯度函数并求解其最小化参数值

$L = 2x^4 + 3x^3 + 2$

$\frac{dL}{dx} = 8x^3 + 9x^2$

$\frac{dL}{dx}$ 在 $x = 0$ 处为 0，这是一个鞍点，在 $x = - \frac{9}{8}$ 处为 0，这是一个全局最小值。因此，损失函数在 $x^\star = - \frac{9}{8}$ 处被优化。

x_vals = tf.linspace(-2, 2, 201)
x_vals = tf.cast(x_vals, tf.float32)

def loss(x):
  return 2*(x**4) + 3*(x**3) + 2

def grad(f, x):
  with tf.GradientTape() as tape:
    tape.watch(x)
    result = f(x)
  return tape.gradient(result, x)

plt.plot(x_vals, loss(x_vals), c='k', label = "Loss function")
plt.plot(x_vals, grad(loss, x_vals), c='tab:blue', label = "Gradient function")
plt.plot(0, loss(0),  marker="o", c='g', label = "Inflection point")
plt.plot(-9/8, loss(-9/8),  marker="o", c='r', label = "Global minimum")
plt.legend()
plt.ylim(0,5)
plt.xlabel("x")
plt.ylabel("loss")
plt.title("Sample loss function and gradient");

编写一个函数来测试单变量损失函数的优化器收敛性。假设当时间步 $t$ 的更新参数值与时间步 $t-1$ 的值相同时，收敛已实现。在设定的迭代次数后终止测试，并跟踪过程中出现的任何梯度爆炸。为了真正挑战优化算法，请对参数进行不良初始化。在上面的示例中， $x = 2$ 是一个不错的选择，因为它涉及陡峭的梯度，并且会导致拐点。

def convergence_test(optimizer, loss_fn, grad_fn=grad, init_val=2., max_iters=2000):
  # Function for optimizer convergence test
  print(optimizer.title)
  print("-------------------------------")
  # Initializing variables and structures
  x_star = tf.Variable(init_val)
  param_path = []
  converged = False

  for iter in range(1, max_iters + 1):
    x_grad = grad_fn(loss_fn, x_star)

    # Case for exploding gradient
    if tf.math.is_nan(x_grad):
      print(f"Gradient exploded at iteration {iter}\n")
      return []

    # Updating the variable and storing its old-version
    x_old = x_star.numpy()
    optimizer.apply_gradients([x_grad], [x_star])
    param_path.append(x_star.numpy())

    # Checking for convergence
    if x_star == x_old:
      print(f"Converged in {iter} iterations\n")
      converged = True
      break

  # Print early termination message
  if not converged:
    print(f"Exceeded maximum of {max_iters} iterations. Test terminated.\n")
  return param_path

测试以下学习率的梯度下降优化器的收敛性：1e-3、1e-2、1e-1

param_map_gd = {}
learning_rates = [1e-3, 1e-2, 1e-1]
for learning_rate in learning_rates:
  param_map_gd[learning_rate] = (convergence_test(
      GradientDescent(learning_rate=learning_rate), loss_fn=loss))

将参数路径可视化在损失函数的等高线图上。

def viz_paths(param_map, x_vals, loss_fn, title, max_iters=2000):
  # Creating a controur plot of the loss function
  t_vals = tf.range(1., max_iters + 100.)
  t_grid, x_grid = tf.meshgrid(t_vals, x_vals)
  loss_grid = tf.math.log(loss_fn(x_grid))
  plt.pcolormesh(t_vals, x_vals, loss_grid, vmin=0, shading='nearest')
  colors = ['r', 'w', 'c']
  # Plotting the parameter paths over the contour plot
  for i, learning_rate in enumerate(param_map):
    param_path = param_map[learning_rate]
    if len(param_path) > 0:
      x_star = param_path[-1]
      plt.plot(t_vals[:len(param_path)], param_path, c=colors[i])
      plt.plot(len(param_path), x_star, marker='o', c=colors[i], 
              label = f"x*: learning rate={learning_rate}")
  plt.xlabel("Iterations")
  plt.ylabel("Parameter value")
  plt.legend()
  plt.title(f"{title} parameter paths")

viz_paths(param_map_gd, x_vals, loss, "Gradient descent")

使用较小的学习率时，梯度下降似乎卡在了拐点。由于步长较大，增加学习率可以鼓励在高原区域更快地移动；但是，这会带来在损失函数极其陡峭的早期迭代中出现梯度爆炸的风险。

带动量的梯度下降

带动量的梯度下降不仅使用梯度来更新变量，还涉及基于变量先前更新的变量位置变化。动量参数决定了时间步 $t-1$ 的更新对时间步 $t$ 的更新的影响程度。累积动量有助于使变量比基本梯度下降更快地越过高原区域。动量更新规则如下

$\Delta_x^{[t]} = lr \cdot L^\prime(x^{[t-1]}) + p \cdot \Delta_x^{[t-1]}$

$x^{[t]} = x^{[t-1]} - \Delta_x^{[t]}$

其中

$x$ ：正在优化的变量
$\Delta_x$ ： $x$ 的变化
$lr$ ：学习率
$L^\prime(x)$ ：损失函数相对于 x 的梯度
$p$ ：动量参数

class Momentum(tf.Module):

  def __init__(self, learning_rate=1e-3, momentum=0.7):
    # Initialize parameters
    self.learning_rate = learning_rate
    self.momentum = momentum
    self.change = 0.
    self.title = f"Gradient descent optimizer: learning rate={self.learning_rate}"

  def apply_gradients(self, grads, vars):
    # Update variables 
    for grad, var in zip(grads, vars):
      curr_change = self.learning_rate*grad + self.momentum*self.change
      var.assign_sub(curr_change)
      self.change = curr_change

测试以下学习率的动量优化器的收敛性：1e-3、1e-2、1e-1

param_map_mtm = {}
learning_rates = [1e-3, 1e-2, 1e-1]
for learning_rate in learning_rates:
  param_map_mtm[learning_rate] = (convergence_test(
      Momentum(learning_rate=learning_rate),
      loss_fn=loss, grad_fn=grad))

将参数路径可视化在损失函数的等高线图上。

viz_paths(param_map_mtm, x_vals, loss, "Momentum")

自适应矩估计 (Adam)

自适应矩估计 (Adam) 算法是一种高效且高度可泛化的优化技术，它利用了两种关键的梯度下降方法：动量和均方根传播 (RMSP)。动量通过使用一阶矩（梯度之和）以及衰减参数来加速梯度下降。RMSP 类似；但是，它利用二阶矩（梯度平方之和）。

Adam 算法结合了一阶矩和二阶矩，以提供更可泛化的更新规则。变量 $x$ 的符号可以通过计算 $\frac{x}{\sqrt{x^2} }$ 来确定。Adam 优化器使用此事实来计算一个更新步骤，该步骤实际上是一个平滑的符号。优化器不是计算 $\frac{x}{\sqrt{x^2} }$ ，而是为每个变量更新计算 $x$ （一阶矩）和 $x^2$ （二阶矩）的平滑版本。

Adam 算法

$\beta_1 \gets 0.9 \; \triangleright \text{文献值}$

$\beta_2 \gets 0.999 \; \triangleright \text{文献值}$

$lr \gets \text{1e-3} \; \triangleright \text{可配置学习率}$

$\epsilon \gets \text{1e-7} \; \triangleright \text{防止除以 0 错误}$

$V_{dv} \gets \vec {\underset{n\times1}{0} } \;\triangleright \text{存储每个变量的动量更新}$

$S_{dv} \gets \vec {\underset{n\times1}{0} } \; \triangleright \text{存储每个变量的 RMSP 更新}$

$t \gets 1$

$\text{在迭代 } t:$

$\;\;\;\; \text{对于} (\frac{dL}{dv}, v) \text{ 在梯度变量对中}:$

$\;\;\;\;\;\;\;\; V_{dv\_i} = \beta_1V_{dv\_i} + (1 - \beta_1)\frac{dL}{dv} \; \triangleright \text{动量更新}$

$\;\;\;\;\;\;\;\; S_{dv\_i} = \beta_2V_{dv\_i} + (1 - \beta_2)(\frac{dL}{dv})^2 \; \triangleright \text{RMSP 更新}$

$\;\;\;\;\;\;\;\; v_{dv}^{bc} = \frac{V_{dv\_i} }{(1-\beta_1)^t} \; \triangleright \text{动量偏差校正}$

$\;\;\;\;\;\;\;\; s_{dv}^{bc} = \frac{S_{dv\_i} }{(1-\beta_2)^t} \; \triangleright \text{RMSP 偏差校正}$

$\;\;\;\;\;\;\;\; v = v - lr\frac{v_{dv}^{bc} }{\sqrt{s_{dv}^{bc} } + \epsilon} \; \triangleright \text{参数更新}$

$\;\;\;\;\;\;\;\; t = t + 1$

算法结束

鉴于 $V_{dv}$ 和 $S_{dv}$ 初始化为 0，并且 $\beta_1$ 和 $\beta_2$ 接近 1，动量和 RMSP 更新自然会偏向于 0；因此，变量可以从偏差校正中受益。偏差校正还有助于控制权重在接近全局最小值时的振荡。

class Adam(tf.Module):

    def __init__(self, learning_rate=1e-3, beta_1=0.9, beta_2=0.999, ep=1e-7):
      # Initialize the Adam parameters
      self.beta_1 = beta_1
      self.beta_2 = beta_2
      self.learning_rate = learning_rate
      self.ep = ep
      self.t = 1.
      self.v_dvar, self.s_dvar = [], []
      self.title = f"Adam: learning rate={self.learning_rate}"
      self.built = False

    def apply_gradients(self, grads, vars):
      # Set up moment and RMSprop slots for each variable on the first call
      if not self.built:
        for var in vars:
          v = tf.Variable(tf.zeros(shape=var.shape))
          s = tf.Variable(tf.zeros(shape=var.shape))
          self.v_dvar.append(v)
          self.s_dvar.append(s)
        self.built = True
      # Perform Adam updates
      for i, (d_var, var) in enumerate(zip(grads, vars)):
        # Moment calculation
        self.v_dvar[i] = self.beta_1*self.v_dvar[i] + (1-self.beta_1)*d_var
        # RMSprop calculation
        self.s_dvar[i] = self.beta_2*self.s_dvar[i] + (1-self.beta_2)*tf.square(d_var)
        # Bias correction
        v_dvar_bc = self.v_dvar[i]/(1-(self.beta_1**self.t))
        s_dvar_bc = self.s_dvar[i]/(1-(self.beta_2**self.t))
        # Update model variables
        var.assign_sub(self.learning_rate*(v_dvar_bc/(tf.sqrt(s_dvar_bc) + self.ep)))
      # Increment the iteration counter
      self.t += 1.

测试 Adam 优化器在与梯度下降示例相同的学习率下的性能。

param_map_adam = {}
learning_rates = [1e-3, 1e-2, 1e-1]
for learning_rate in learning_rates:
  param_map_adam[learning_rate] = (convergence_test(
      Adam(learning_rate=learning_rate), loss_fn=loss))

将参数路径可视化在损失函数的等高线图上。

viz_paths(param_map_adam, x_vals, loss, "Adam")

在这个特定示例中，当使用较小的学习率时，Adam 优化器的收敛速度比传统的梯度下降慢。但是，当使用较大的学习率时，该算法成功地越过了高原区域并收敛到全局最小值。由于 Adam 在遇到大梯度时会动态缩放学习率，因此梯度爆炸不再是问题。

结论

本笔记本介绍了使用 TensorFlow Core API 编写和比较优化器的基础知识。尽管像 Adam 这样的预构建优化器是可泛化的，但它们并不总是每个模型或数据集的最佳选择。对优化过程进行细粒度控制可以帮助简化 ML 训练工作流程并提高整体性能。有关自定义优化器的更多示例，请参阅以下文档

此 Adam 优化器在多层感知器教程和分布式训练中使用
模型花园包含使用 Core API 编写的各种自定义优化器。