 在 TensorFlow.org 上檢視
|
 在 Google Colab 中執行
|
 在 GitHub 上檢視原始碼
|
 下載筆記本
|
簡介
這個筆記本介紹如何使用 TensorFlow Core 低階 API 建立自訂最佳化工具的流程。請造訪 Core API 總覽，進一步瞭解 TensorFlow Core 及其預期用途。
Keras 最佳化工具模組是許多一般訓練用途的建議最佳化工具組。它包含各種預先建構的最佳化工具，以及用於自訂的子類別化功能。Keras 最佳化工具也與使用 Core API 建構的自訂層、模型和訓練迴圈相容。這些預先建構和可自訂的最佳化工具適用於大多數情況，但 Core API 可讓您完全控制最佳化流程。例如，「銳度感知最小化」(SAM) 等技術需要將模型和最佳化工具耦合，這不符合 ML 最佳化工具的傳統定義。本指南逐步說明如何從頭開始使用 Core 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 訓練工作流程並提高整體效能。如需更多自訂最佳化工具範例，請參閱下列文件