4.4. Softmax回归的从零开始实现
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 SageMaker Studio Lab 中打开 Notebook

因为softmax回归是如此基础,我们认为你应该知道如何自己实现它。在这里,我们只定义模型中softmax的特定方面,并重用线性回归部分的其他组件,包括训练循环。

import torch
from d2l import torch as d2l
from mxnet import autograd, gluon, np, npx
from d2l import mxnet as d2l

npx.set_np()
from functools import partial
import jax
from flax import linen as nn
from jax import numpy as jnp
from d2l import jax as d2l
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
import tensorflow as tf
from d2l import tensorflow as d2l

4.4.1. Softmax

让我们从最重要的部分开始:从标量到概率的映射。为了复习一下,回顾一下在 2.3.6节2.3.7节 中讨论的张量中沿特定维度的求和运算符的操作。给定一个矩阵 X,我们可以对所有元素求和(默认情况下),或者只对同一轴上的元素求和。axis 变量让我们能够计算行和与列和。

X = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
X.sum(0, keepdims=True), X.sum(1, keepdims=True)
(tensor([[5., 7., 9.]]),
 tensor([[ 6.],
         [15.]]))
X = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
X.sum(0, keepdims=True), X.sum(1, keepdims=True)
[22:09:48] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU
(array([[5., 7., 9.]]),
 array([[ 6.],
        [15.]]))
X = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
X.sum(0, keepdims=True), X.sum(1, keepdims=True)
(Array([[5., 7., 9.]], dtype=float32),
 Array([[ 6.],
        [15.]], dtype=float32))
X = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
tf.reduce_sum(X, 0, keepdims=True), tf.reduce_sum(X, 1, keepdims=True)
(<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[5., 7., 9.]], dtype=float32)>,
 <tf.Tensor: shape=(2, 1), dtype=float32, numpy=
 array([[ 6.],
        [15.]], dtype=float32)>)

计算softmax需要三个步骤:(i)对每个项求幂;(ii)对每一行求和以计算每个样本的归一化常数;(iii)将每一行除以其归一化常数,确保结果总和为1。

(4.4.1)\[\mathrm{softmax}(\mathbf{X})_{ij} = \frac{\exp(\mathbf{X}_{ij})}{\sum_k \exp(\mathbf{X}_{ik})}.\]

分母的(对数)被称为(对数)配分函数(partition function)。它是在统计物理学中引入的,用于对热力学系综中所有可能的状态求和。实现很简单。

def softmax(X):
    X_exp = torch.exp(X)
    partition = X_exp.sum(1, keepdims=True)
    return X_exp / partition  # The broadcasting mechanism is applied here
def softmax(X):
    X_exp = np.exp(X)
    partition = X_exp.sum(1, keepdims=True)
    return X_exp / partition  # The broadcasting mechanism is applied here
def softmax(X):
    X_exp = jnp.exp(X)
    partition = X_exp.sum(1, keepdims=True)
    return X_exp / partition  # The broadcasting mechanism is applied here
def softmax(X):
    X_exp = tf.exp(X)
    partition = tf.reduce_sum(X_exp, 1, keepdims=True)
    return X_exp / partition  # The broadcasting mechanism is applied here

对于任何输入 X,我们将每个元素转换为一个非负数。每一行的总和为1,这是概率所要求的。注意:上面的代码对于非常大或非常小的参数并健壮。虽然它足以说明发生了什么,但你不应将此代码逐字用于任何严肃的目的。深度学习框架内置了此类保护措施,我们将继续使用内置的softmax。

X = torch.rand((2, 5))
X_prob = softmax(X)
X_prob, X_prob.sum(1)
(tensor([[0.2511, 0.1417, 0.1158, 0.2529, 0.2385],
         [0.2004, 0.1419, 0.1957, 0.2504, 0.2117]]),
 tensor([1., 1.]))
X = np.random.rand(2, 5)
X_prob = softmax(X)
X_prob, X_prob.sum(1)
(array([[0.17777154, 0.1857739 , 0.20995119, 0.23887765, 0.18762572],
        [0.24042214, 0.1757977 , 0.23786479, 0.15572716, 0.19018826]]),
 array([1., 1.]))
X = jax.random.uniform(jax.random.PRNGKey(d2l.get_seed()), (2, 5))
X_prob = softmax(X)
X_prob, X_prob.sum(1)
(Array([[0.21010879, 0.13711403, 0.17836188, 0.19433393, 0.28008142],
        [0.13819133, 0.12869641, 0.1519639 , 0.32253352, 0.25861484]],      dtype=float32),
 Array([1., 1.], dtype=float32))
X = tf.random.uniform((2, 5))
X_prob = softmax(X)
X_prob, tf.reduce_sum(X_prob, 1)
(<tf.Tensor: shape=(2, 5), dtype=float32, numpy=
 array([[0.13828081, 0.26672134, 0.12680084, 0.1782964 , 0.2899006 ],
        [0.22160178, 0.12449112, 0.22632608, 0.18417503, 0.24340603]],
       dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([1., 1.], dtype=float32)>)

4.4.2. 模型

我们现在拥有了实现softmax回归模型所需的一切。就像我们的线性回归例子一样,每个实例都将由一个固定长度的向量表示。由于这里的原始数据由 \(28 \times 28\) 像素的图像组成,我们将每个图像展平,将其视为长度为784的向量。在后面的章节中,我们将介绍卷积神经网络,它以更令人满意的方式利用了空间结构。

在softmax回归中,我们网络的输出数量应等于类别的数量。由于我们的数据集有10个类别,我们的网络输出维度为10。因此,我们的权重构成一个 \(784 \times 10\) 的矩阵,外加一个 \(1 \times 10\) 的行向量作为偏置。与线性回归一样,我们用高斯噪声初始化权重 W。偏置初始化为零。

class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = torch.normal(0, sigma, size=(num_inputs, num_outputs),
                              requires_grad=True)
        self.b = torch.zeros(num_outputs, requires_grad=True)

    def parameters(self):
        return [self.W, self.b]
class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = np.random.normal(0, sigma, (num_inputs, num_outputs))
        self.b = np.zeros(num_outputs)
        self.W.attach_grad()
        self.b.attach_grad()

    def collect_params(self):
        return [self.W, self.b]
class SoftmaxRegressionScratch(d2l.Classifier):
    num_inputs: int
    num_outputs: int
    lr: float
    sigma: float = 0.01

    def setup(self):
        self.W = self.param('W', nn.initializers.normal(self.sigma),
                            (self.num_inputs, self.num_outputs))
        self.b = self.param('b', nn.initializers.zeros, self.num_outputs)
class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = tf.random.normal((num_inputs, num_outputs), 0, sigma)
        self.b = tf.zeros(num_outputs)
        self.W = tf.Variable(self.W)
        self.b = tf.Variable(self.b)

下面的代码定义了网络如何将每个输入映射到输出。请注意,在将数据通过我们的模型之前,我们使用 reshape 将批次中的每个 \(28 \times 28\) 像素图像展平为一个向量。

@d2l.add_to_class(SoftmaxRegressionScratch)
def forward(self, X):
    X = X.reshape((-1, self.W.shape[0]))
    return softmax(torch.matmul(X, self.W) + self.b)
@d2l.add_to_class(SoftmaxRegressionScratch)
def forward(self, X):
    X = X.reshape((-1, self.W.shape[0]))
    return softmax(np.dot(X, self.W) + self.b)
@d2l.add_to_class(SoftmaxRegressionScratch)
def forward(self, X):
    X = X.reshape((-1, self.W.shape[0]))
    return softmax(jnp.matmul(X, self.W) + self.b)
@d2l.add_to_class(SoftmaxRegressionScratch)
def forward(self, X):
    X = tf.reshape(X, (-1, self.W.shape[0]))
    return softmax(tf.matmul(X, self.W) + self.b)

4.4.3. 交叉熵损失

接下来,我们需要实现交叉熵损失函数(在 4.1.2节 中介绍)。这可能是整个深度学习中最常见的损失函数。目前,可以轻松归为分类问题的深度学习应用数量远远超过那些更适合作为回归问题处理的应用。

回想一下,交叉熵取的是分配给真实标签的预测概率的负对数似然。为了提高效率,我们避免使用Python的for循环,而是使用索引。特别地,\(\mathbf{y}\) 中的独热编码使我们能够选择 \(\hat{\mathbf{y}}\) 中匹配的项。

为了看到它的实际效果,我们创建了样本数据 y_hat,其中包含2个样本在3个类别上的预测概率,以及它们对应的标签 y。正确的标签分别是 \(0\)\(2\)(即第一类和第三类)。使用 y 作为 y_hat 中概率的索引,我们可以高效地挑选出各项。

y = torch.tensor([0, 2])
y_hat = torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]
tensor([0.1000, 0.5000])

现在我们可以通过对所选概率的对数取平均值来实现交叉熵损失函数。

def cross_entropy(y_hat, y):
    return -torch.log(y_hat[list(range(len(y_hat))), y]).mean()

cross_entropy(y_hat, y)
tensor(1.4979)
@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)
y = np.array([0, 2])
y_hat = np.array([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]
array([0.1, 0.5])

现在我们可以通过对所选概率的对数取平均值来实现交叉熵损失函数。

def cross_entropy(y_hat, y):
    return -np.log(y_hat[list(range(len(y_hat))), y]).mean()

cross_entropy(y_hat, y)
array(1.4978662)
@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)
y = jnp.array([0, 2])
y_hat = jnp.array([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]
Array([0.1, 0.5], dtype=float32)

现在我们可以通过对所选概率的对数取平均值来实现交叉熵损失函数。

请注意,为了利用 jax.jit 来加速 JAX 的实现,并确保 loss 是一个纯函数,cross_entropy 函数在 loss 内部被重新定义,以避免使用任何可能使 loss 函数变得不纯的全局变量或函数。我们建议感兴趣的读者参考关于 jax.jit 和纯函数的 JAX 文档

def cross_entropy(y_hat, y):
    return -jnp.log(y_hat[list(range(len(y_hat))), y]).mean()

cross_entropy(y_hat, y)
Array(1.4978662, dtype=float32)
@d2l.add_to_class(SoftmaxRegressionScratch)
@partial(jax.jit, static_argnums=(0))
def loss(self, params, X, y, state):
    def cross_entropy(y_hat, y):
        return -jnp.log(y_hat[list(range(len(y_hat))), y]).mean()
    y_hat = state.apply_fn({'params': params}, *X)
    # The returned empty dictionary is a placeholder for auxiliary data,
    # which will be used later (e.g., for batch norm)
    return cross_entropy(y_hat, y), {}
y_hat = tf.constant([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y = tf.constant([0, 2])
tf.boolean_mask(y_hat, tf.one_hot(y, depth=y_hat.shape[-1]))
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.1, 0.5], dtype=float32)>

现在我们可以通过对所选概率的对数取平均值来实现交叉熵损失函数。

def cross_entropy(y_hat, y):
    return -tf.reduce_mean(tf.math.log(tf.boolean_mask(
        y_hat, tf.one_hot(y, depth=y_hat.shape[-1]))))

cross_entropy(y_hat, y)
<tf.Tensor: shape=(), dtype=float32, numpy=1.4978662>
@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)

4.4.4. 训练

我们重用在 3.4节 中定义的 fit 方法来训练模型10个周期。注意,周期数(max_epochs)、小批量大小(batch_size)和学习率(lr)都是可调的超参数。这意味着虽然这些值不是在我们的主训练循环中学习的,但它们仍然影响我们模型的性能,包括训练性能和泛化性能。在实践中,您将希望根据数据的*验证*集来选择这些值,然后,最终在*测试*集上评估您的最终模型。正如在 3.6.3节 中讨论的,我们将把Fashion-MNIST的测试数据视为验证集,因此报告此数据集上的验证损失和验证准确率。

data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegressionScratch(num_inputs=784, num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)
../_images/output_softmax-regression-scratch_aecde8_120_0.svg
data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegressionScratch(num_inputs=784, num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)
../_images/output_softmax-regression-scratch_aecde8_123_0.svg
data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegressionScratch(num_inputs=784, num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)
../_images/output_softmax-regression-scratch_aecde8_126_0.svg
data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegressionScratch(num_inputs=784, num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)
../_images/output_softmax-regression-scratch_aecde8_129_0.svg

4.4.5. 预测

现在训练已经完成,我们的模型准备好对一些图像进行分类了。

X, y = next(iter(data.val_dataloader()))
preds = model(X).argmax(axis=1)
preds.shape
torch.Size([256])
X, y = next(iter(data.val_dataloader()))
preds = model(X).argmax(axis=1)
preds.shape
(256,)
X, y = next(iter(data.val_dataloader()))
preds = model.apply({'params': trainer.state.params}, X).argmax(axis=1)
preds.shape
(256,)
X, y = next(iter(data.val_dataloader()))
preds = tf.argmax(model(X), axis=1)
preds.shape
TensorShape([256])

我们更感兴趣的是那些我们标记*错误*的图像。我们通过比较它们的实际标签(文本输出的第一行)和模型的预测(文本输出的第二行)来可视化它们。

wrong = preds.type(y.dtype) != y
X, y, preds = X[wrong], y[wrong], preds[wrong]
labels = [a+'\n'+b for a, b in zip(
    data.text_labels(y), data.text_labels(preds))]
data.visualize([X, y], labels=labels)
../_images/output_softmax-regression-scratch_aecde8_150_0.svg
wrong = preds.astype(y.dtype) != y
X, y, preds = X[wrong], y[wrong], preds[wrong]
labels = [a+'\n'+b for a, b in zip(
    data.text_labels(y), data.text_labels(preds))]
data.visualize([X, y], labels=labels)
../_images/output_softmax-regression-scratch_aecde8_153_0.svg
wrong = preds.astype(y.dtype) != y
X, y, preds = X[wrong], y[wrong], preds[wrong]
labels = [a+'\n'+b for a, b in zip(
    data.text_labels(y), data.text_labels(preds))]
data.visualize([X, y], labels=labels)
../_images/output_softmax-regression-scratch_aecde8_156_0.svg
wrong = tf.cast(preds, y.dtype) != y
X, y, preds = X[wrong], y[wrong], preds[wrong]
labels = [a+'\n'+b for a, b in zip(
    data.text_labels(y), data.text_labels(preds))]
data.visualize([X, y], labels=labels)
../_images/output_softmax-regression-scratch_aecde8_159_0.svg

4.4.6. 小结

到现在为止,我们已经开始积累了一些解决线性回归和分类问题的经验。借此,我们已经达到了可以说是1960-1970年代统计建模的最高水平。在下一节中,我们将向您展示如何利用深度学习框架来更有效地实现这个模型。

4.4.7. 练习

  1. 在本节中,我们直接根据softmax操作的数学定义实现了softmax函数。如 4.1节 所讨论的,这可能会导致数值不稳定性。

    1. 如果一个输入的值为 \(100\),测试 softmax 是否仍然能正常工作。

    2. 如果所有输入中的最大值小于 \(-100\),测试 softmax 是否仍然能正常工作?

    3. 通过查看相对于参数中最大项的值来实现一个修复方案。

  2. 实现一个遵循交叉熵损失函数定义 \(\sum_i y_i \log \hat{y}_i\)cross_entropy 函数。

    1. 在本节的代码示例中尝试一下。

    2. 你认为它为什么运行得更慢?

    3. 你应该使用它吗?什么时候使用它有意义?

    4. 你需要注意什么?提示:考虑对数的定义域。

  3. 返回最可能的标签总是好主意吗?例如,你会在医疗诊断中这样做吗?你会如何尝试解决这个问题?

  4. 假设我们想使用softmax回归根据一些特征来预测下一个词。一个大的词汇表可能会引起什么问题?

  5. 试验本节代码中的超参数。特别是:

    1. 绘制当你改变学习率时验证损失的变化情况。

    2. 当你改变小批量大小时,验证损失和训练损失会改变吗?你需要将其调到多大或多小才能看到效果?