基本训练循环

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

在之前的指南中,您已经了解了 张量变量梯度带模块。在本指南中,您将把所有这些内容整合在一起以训练模型。

TensorFlow 还包括 tf.Keras API,这是一个高级神经网络 API,它提供了有用的抽象来减少样板代码。但是,在本指南中,您将使用基本类。

设置

import tensorflow as tf

import matplotlib.pyplot as plt

colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

解决机器学习问题

解决机器学习问题通常包括以下步骤

  • 获取训练数据。
  • 定义模型。
  • 定义损失函数。
  • 遍历训练数据,计算理想值的损失
  • 计算该损失的梯度,并使用优化器调整变量以适应数据。
  • 评估您的结果。

为了说明目的,在本指南中,您将开发一个简单的线性模型,\(f(x) = x * W + b\),它有两个变量:\(W\)(权重)和 \(b\)(偏差)。

这是最基本的机器学习问题:给定 \(x\) 和 \(y\),尝试通过 简单线性回归 找到一条线的斜率和偏移量。

数据

监督学习使用输入(通常表示为x)和输出(表示为y,通常称为标签)。目标是从配对的输入和输出中学习,以便您可以根据输入预测输出的值。

在 TensorFlow 中,您的数据的每个输入几乎总是由张量表示,并且通常是一个向量。在监督训练中,输出(或您要预测的值)也是一个张量。

以下是一些通过向一条线上的点添加高斯(正态)噪声而合成的数据。

# The actual line
TRUE_W = 3.0
TRUE_B = 2.0

NUM_EXAMPLES = 201

# A vector of random x values
x = tf.linspace(-2,2, NUM_EXAMPLES)
x = tf.cast(x, tf.float32)

def f(x):
  return x * TRUE_W + TRUE_B

# Generate some noise
noise = tf.random.normal(shape=[NUM_EXAMPLES])

# Calculate y
y = f(x) + noise
# Plot all the data
plt.plot(x, y, '.')
plt.show()

张量通常被收集在一起形成批次,或将多个输入和输出堆叠在一起的组。批处理可以带来一些训练优势,并且与加速器和矢量化计算配合良好。鉴于此数据集很小,您可以将整个数据集视为单个批次。

定义模型

使用 tf.Variable 来表示模型中的所有权重。一个 tf.Variable 存储一个值,并在需要时以张量形式提供该值。有关更多详细信息,请参阅 变量指南

使用 tf.Module 来封装变量和计算。您可以使用任何 Python 对象,但这样可以轻松保存。

在这里,您将wb都定义为变量。

class MyModel(tf.Module):
  def __init__(self, **kwargs):
    super().__init__(**kwargs)
    # Initialize the weights to `5.0` and the bias to `0.0`
    # In practice, these should be randomly initialized
    self.w = tf.Variable(5.0)
    self.b = tf.Variable(0.0)

  def __call__(self, x):
    return self.w * x + self.b

model = MyModel()

# List the variables tf.modules's built-in variable aggregation.
print("Variables:", model.variables)

# Verify the model works
assert model(3.0).numpy() == 15.0

初始变量在这里以固定方式设置,但 Keras 附带了许多您可以使用的 初始化器,无论是否使用其他 Keras 功能。

定义损失函数

损失函数衡量模型对给定输入的输出与目标输出匹配的程度。目标是在训练期间最小化这种差异。定义标准 L2 损失,也称为“均方”误差

# This computes a single loss value for an entire batch
def loss(target_y, predicted_y):
  return tf.reduce_mean(tf.square(target_y - predicted_y))

在训练模型之前,您可以通过将模型的预测以红色绘制,将训练数据以蓝色绘制来可视化损失值

plt.plot(x, y, '.', label="Data")
plt.plot(x, f(x), label="Ground truth")
plt.plot(x, model(x), label="Predictions")
plt.legend()
plt.show()

print("Current loss: %1.6f" % loss(y, model(x)).numpy())

定义训练循环

训练循环包括按顺序重复执行三个任务

  • 将一批输入通过模型生成输出
  • 通过将输出与输出(或标签)进行比较来计算损失
  • 使用梯度带找到梯度
  • 使用这些梯度优化变量

对于此示例,您可以使用 梯度下降 训练模型。

梯度下降方案有许多变体,这些变体都包含在 tf.keras.optimizers 中。但本着从第一性原理构建的精神,您将使用 tf.GradientTape 进行自动微分和 tf.assign_sub 递减值(它结合了 tf.assigntf.sub)来实现基本数学运算。

# Given a callable model, inputs, outputs, and a learning rate...
def train(model, x, y, learning_rate):

  with tf.GradientTape() as t:
    # Trainable variables are automatically tracked by GradientTape
    current_loss = loss(y, model(x))

  # Use GradientTape to calculate the gradients with respect to W and b
  dw, db = t.gradient(current_loss, [model.w, model.b])

  # Subtract the gradient scaled by the learning rate
  model.w.assign_sub(learning_rate * dw)
  model.b.assign_sub(learning_rate * db)

为了查看训练情况,您可以将相同的一批xy通过训练循环,并查看Wb是如何演变的。

model = MyModel()

# Collect the history of W-values and b-values to plot later
weights = []
biases = []
epochs = range(10)

# Define a training loop
def report(model, loss):
  return f"W = {model.w.numpy():1.2f}, b = {model.b.numpy():1.2f}, loss={loss:2.5f}"


def training_loop(model, x, y):

  for epoch in epochs:
    # Update the model with the single giant batch
    train(model, x, y, learning_rate=0.1)

    # Track this before I update
    weights.append(model.w.numpy())
    biases.append(model.b.numpy())
    current_loss = loss(y, model(x))

    print(f"Epoch {epoch:2d}:")
    print("    ", report(model, current_loss))

进行训练

current_loss = loss(y, model(x))

print(f"Starting:")
print("    ", report(model, current_loss))

training_loop(model, x, y)

绘制权重随时间推移的演变

plt.plot(epochs, weights, label='Weights', color=colors[0])
plt.plot(epochs, [TRUE_W] * len(epochs), '--',
         label = "True weight", color=colors[0])

plt.plot(epochs, biases, label='bias', color=colors[1])
plt.plot(epochs, [TRUE_B] * len(epochs), "--",
         label="True bias", color=colors[1])

plt.legend()
plt.show()

可视化训练后的模型的性能

plt.plot(x, y, '.', label="Data")
plt.plot(x, f(x), label="Ground truth")
plt.plot(x, model(x), label="Predictions")
plt.legend()
plt.show()

print("Current loss: %1.6f" % loss(model(x), y).numpy())

相同的解决方案,但使用 Keras

将上面的代码与 Keras 中的等效代码进行对比非常有用。

如果您子类化 tf.keras.Model,定义模型看起来完全一样。请记住,Keras 模型最终继承自模块。

class MyModelKeras(tf.keras.Model):
  def __init__(self, **kwargs):
    super().__init__(**kwargs)
    # Initialize the weights to `5.0` and the bias to `0.0`
    # In practice, these should be randomly initialized
    self.w = tf.Variable(5.0)
    self.b = tf.Variable(0.0)

  def call(self, x):
    return self.w * x + self.b

keras_model = MyModelKeras()

# Reuse the training loop with a Keras model
training_loop(keras_model, x, y)

# You can also save a checkpoint using Keras's built-in support
keras_model.save_weights("my_checkpoint")

与其每次创建模型时都编写新的训练循环,不如使用 Keras 的内置功能作为快捷方式。当您不想编写或调试 Python 训练循环时,这很有用。

如果您这样做,您将需要使用 model.compile() 设置参数,并使用 model.fit() 进行训练。使用 Keras 实现的 L2 损失和梯度下降可以减少代码量,同样作为一种快捷方式。Keras 损失和优化器也可以在这些便利函数之外使用,前面的示例也可以使用它们。

keras_model = MyModelKeras()

# compile sets the training parameters
keras_model.compile(
    # By default, fit() uses tf.function().  You can
    # turn that off for debugging, but it is on now.
    run_eagerly=False,

    # Using a built-in optimizer, configuring as an object
    optimizer=tf.keras.optimizers.SGD(learning_rate=0.1),

    # Keras comes with built-in MSE error
    # However, you could use the loss function
    # defined above
    loss=tf.keras.losses.mean_squared_error,
)

Keras fit 期望批处理数据或作为 NumPy 数组的完整数据集。NumPy 数组被切分成批次,默认批次大小为 32。

在这种情况下,为了匹配手写循环的行为,您应该将 x 作为大小为 1000 的单个批次传递。

print(x.shape[0])
keras_model.fit(x, y, epochs=10, batch_size=1000)

请注意,Keras 在训练后打印出损失,而不是在训练前,因此第一个损失看起来更低,但除此之外,这基本上显示了相同的训练性能。

后续步骤

在本指南中,您已经了解了如何使用张量、变量、模块和梯度带的核心类来构建和训练模型,以及这些概念如何映射到 Keras。

然而,这是一个极其简单的示例。有关更实用的介绍,请参阅 自定义训练演练

有关使用内置 Keras 训练循环的更多信息,请参阅 本指南。有关训练循环和 Keras 的更多信息,请参阅 本指南。有关编写自定义分布式训练循环的信息,请参阅 本指南