使用 tensorflow/datasets 加载数据训练简单的神经网络

内容

根据 Apache 许可证 2.0 版（“许可证”）授权；

根据 Apache 许可证 2.0 版（“许可证”）授权；除非符合许可证，否则您不得使用此文件。您可以在以下网址获取许可证副本：

https://www.apache.org/licenses/LICENSE-2.0

除非适用法律要求或以书面形式达成协议，否则根据许可证分发的软件按“原样”分发，不附带任何明示或暗示的保证和条件。请参阅许可证以了解管理权限和限制的具体语言。

使用 TensorFlow/Datasets 加载数据训练简单的神经网络#

源自 neural_network_and_data_loading.ipynb

让我们结合我们在快速入门中展示的所有内容来训练一个简单的神经网络。我们首先将使用 JAX 进行计算，并在 MNIST 上指定并训练一个简单的 MLP。我们将使用 tensorflow/datasets 数据加载 API 加载图像和标签（因为它非常棒，而且世界不需要另一个数据加载库 :P）。

当然，您可以将 JAX 与任何与 NumPy 兼容的 API 配合使用，以使模型的指定更加即插即用。这里，仅出于说明目的，我们不会使用任何神经网络库或特殊 API 来构建我们的模型。

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

超参数#

让我们先处理一些簿记事项。

# A helper function to randomly initialize weights and biases
# for a dense neural network layer
def random_layer_params(m, n, key, scale=1e-2):
  w_key, b_key = random.split(key)
  return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))

# Initialize all layers for a fully-connected neural network with sizes "sizes"
def init_network_params(sizes, key):
  keys = random.split(key, len(sizes))
  return [random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]

layer_sizes = [784, 512, 512, 10]
step_size = 0.01
num_epochs = 10
batch_size = 128
n_targets = 10
params = init_network_params(layer_sizes, random.key(0))

自动批处理预测#

首先，让我们定义我们的预测函数。请注意，我们是在为单个图像示例定义此函数。我们将使用 JAX 的vmap函数自动处理小批量，而不会造成性能损失。

from jax.scipy.special import logsumexp

def relu(x):
  return jnp.maximum(0, x)

def predict(params, image):
  # per-example predictions
  activations = image
  for w, b in params[:-1]:
    outputs = jnp.dot(w, activations) + b
    activations = relu(outputs)

  final_w, final_b = params[-1]
  logits = jnp.dot(final_w, activations) + final_b
  return logits - logsumexp(logits)

让我们检查一下我们的预测函数是否只对单个图像起作用。

# This works on single examples
random_flattened_image = random.normal(random.key(1), (28 * 28,))
preds = predict(params, random_flattened_image)
print(preds.shape)

(10,)

# Doesn't work with a batch
random_flattened_images = random.normal(random.key(1), (10, 28 * 28))
try:
  preds = predict(params, random_flattened_images)
except TypeError:
  print('Invalid shapes!')

Invalid shapes!

# Let's upgrade it to handle batches using `vmap`

# Make a batched version of the `predict` function
batched_predict = vmap(predict, in_axes=(None, 0))

# `batched_predict` has the same call signature as `predict`
batched_preds = batched_predict(params, random_flattened_images)
print(batched_preds.shape)

(10, 10)

至此，我们已经具备了定义神经网络并对其进行训练的所有要素。我们构建了一个predict的自动批处理版本，我们应该能够在损失函数中使用它。我们应该能够使用grad来计算损失相对于神经网络参数的导数。最后，我们应该能够使用jit来加速所有操作。

实用程序和损失函数#

def one_hot(x, k, dtype=jnp.float32):
  """Create a one-hot encoding of x of size k."""
  return jnp.array(x[:, None] == jnp.arange(k), dtype)

def accuracy(params, images, targets):
  target_class = jnp.argmax(targets, axis=1)
  predicted_class = jnp.argmax(batched_predict(params, images), axis=1)
  return jnp.mean(predicted_class == target_class)

def loss(params, images, targets):
  preds = batched_predict(params, images)
  return -jnp.mean(preds * targets)

@jit
def update(params, x, y):
  grads = grad(loss)(params, x, y)
  return [(w - step_size * dw, b - step_size * db)
          for (w, b), (dw, db) in zip(params, grads)]