我需要回顾一下

提取数据集

我们将训练 ResNet 模型，以从猫狗数据集中预测两个类别。请下载并解压猫和狗的图像。

解压数据集文件

1import zipfile
2
3with zipfile.ZipFile('../train.zip', 'r') as zip_ref:
4    zip_ref.extractall('.')

在上面的代码中：

• 第 1 行：我们导入 zipfile 库。

• 第 3–4 行：我们调用 zipfile 模块的 ZipFile() 方法，以只读模式打开 ZIP 文件，并将其赋值为 zip_ref。使用 with 语句可在代码执行后自动关闭文件。接着调用 extractall() 方法，将 ZIP 文件内容解压到当前目录。

在 Flax 中加载数据集

由于 JAX 和 Flax 不提供任何数据加载器，我们使用来自 PyTorch 或 TensorFlow 的数据加载工具。当使用 PyTorch 时，首先需要创建一个数据集类。

数据集类

1import torch
2import os
3from PIL import Image
4from torch.utils.data import Dataset
5import pandas as pd
6
7class CatsDogsDataset(Dataset):
8    def __init__(self, root_dir, annotation_file, transform=None):
9        self.root_dir = root_dir
10        self.annotations = pd.read_csv(annotation_file)
11        self.transform = transform
12
13    def __len__(self):
14        return len(self.annotations)
15
16    def __getitem__(self, index):
17        img_id = self.annotations.iloc[index, 0]
18        img = Image.open(os.path.join(self.root_dir, img_id)).convert("RGB")
19        y_label = torch.tensor(float(self.annotations.iloc[index, 1]))
20
21        if self.transform is not None:
22            img = self.transform(img)
23
24        return (img, y_label)

在上面的代码中：

• 第 1–5 行：我们导入所需的库：PyTorch 库、os、PIL 中的 Image 模块、torch.utils.data 中的 Dataset 模块，以及 pandas 库并简称为 pd。

• 第 7–24 行：我们定义了一个名为 CatsDogsDataset 的数据集类，用于加载图像和标签。在该类内部：

  • 第 8–11 行：我们定义构造函数 __init__()，接收根目录和标注文件作为参数，还有一个可选参数 transform，用于对图像进行变换。构造函数初始化类的实例变量。

  • 第 13–14 行：我们定义 __len__() 函数，返回数据集中样本的数量。

  • 第 16–24 行：我们定义 __getitem__() 函数，接收一个索引参数 index。从标注文件中获取 img_id 和 y_label，并将图像加载到 img 变量中。如果存在变换，则应用变换，最后返回图像及其标签。

接下来，我们创建一个包含图像路径和标签的 pandas DataFrame。

训练数据 DataFrame

1train_df = pd.DataFrame(columns=["img_path","label"])
2train_df["img_path"] = os.listdir("train/")
3
4for idx, i in enumerate(os.listdir("train/")):
5    if "cat" in i:
6        train_df["label"][idx] = 0
7    if "dog" in i:
8        train_df["label"][idx] = 1
9
10print(train_df)
11train_df.to_csv (r'train_csv.csv', index = False, header=True)

在上面的代码中：

• 第 1 行：我们创建一个名为 train_df 的 pandas DataFrame，包含两列：img_path 用于记录文件路径，label 用于存储标签。

• 第 2 行：我们调用 os 模块的 listdir() 方法，将 train 目录中的所有文件名填入 train_df 的 img_path 列。

• 第 4–8 行：我们使用 for 循环遍历 train 目录中的每个文件名。若文件名包含 "cat"，则赋予标签 0；若包含 "dog"，则赋予标签 1。

• 第 10–11 行：最后，打印 train_df 并调用其 to_csv() 方法，将其保存为 CSV 文件。

可视化数据集

我们将从猫狗数据集的每个类别中可视化一张图像。

可视化数据集

1from PIL import Image
2
3img1 = Image.open("train/cat.1.jpg")
4img1.save("output/img1.png")
5
6img2 = Image.open("train/dog.3.jpg")
7img2.save("output/img2.png")

在上面的代码中：

• 第 1 行：我们从 Python 图像库（PIL）中导入 Image 模块。

• 第 3–4 行：我们从 train 文件夹中打开一张猫类别的图像，并将其保存到 output 文件夹中，以便在运行环境中显示。

• 第 6–7 行：类似地，我们从 train 文件夹中打开一张狗类别的图像，并保存到 output 文件夹中。

注意：train 文件夹中每种类别各有 15 张图像。你可以修改 Image.open() 方法中的图像编号来查看其他图像。

Flax 中的数据变换

我们将定义一个函数，用于堆叠数据并以 NumPy 数组的形式返回。

拼合函数（Collation function）

1import numpy as np
2
3def custom_collate_fn(batch):
4    transposed_data = list(zip(*batch))
5    labels = np.array(transposed_data[1])
6    imgs = np.stack(transposed_data[0])
7    return imgs, labels

在上面的代码中：

• 第 1 行：我们导入 NumPy 库并简称为 np。

• 第 3–7 行：我们定义 custom_collate_fn() 函数，它接收一个包含图像和对应标签的批次元组。该函数从 batch 中提取数据并转置为 transposed_data。然后使用 np.array() 将标签转换为 NumPy 数组，使用 np.stack() 堆叠图像。最后返回图像和标签。

我们创建一个用于调整图像大小的变换，然后将其应用于之前创建的数据集。

图像变换

1from torchvision import transforms
2
3size_image = 224
4
5transform = transforms.Compose([
6    transforms.Resize((size_image,size_image)),
7    np.array])
8dataset = CatsDogsDataset("train", "train_csv.csv", transform=transform)

在上面的代码中：

• 第 1 行：我们从 torchvision 库中导入 transforms 模块。

• 第 3 行：我们定义变量 size_image，表示图像的尺寸。

• 第 5–8 行：我们定义图像缩放变换，并通过传入必要参数创建 CatsDogsDataset 类的实例。

我们将该数据集划分为训练集和测试集，并分别为它们创建数据加载器。

训练集与验证集

1from torch.utils.data import DataLoader
2
3batch_size = 32
4
5train_set, validation_set = torch.utils.data.random_split(dataset, [0.7,0.3])
6
7train_loader = DataLoader(dataset=train_set, collate_fn=custom_collate_fn,shuffle=True, batch_size=batch_size)
8validation_loader = DataLoader(dataset=validation_set, collate_fn=custom_collate_fn, shuffle=True, batch_size=batch_size)

在上面的代码中：

• 第 1 行：我们从 torch.utils.data 导入 DataLoader 模块。

• 第 3 行：我们定义变量 batch_size，表示每个批次中的图像数量。

• 第 5 行：我们将数据集按 7:3 的比例随机划分为训练集和验证集。

• 第 7–8 行：我们调用 DataLoader() 函数，使用 custom_collate_fn() 创建训练和验证数据加载器。这些加载器为各自的数据集提供了可迭代接口。

实例化 Flax ResNet 模型

准备好数据后，我们使用 flaxmodels 包实例化 Flax ResNet 模型。实例化需要以下参数：

• 输出类型
• pretrained 参数，用于加载预训练模型，此处设为 None
• 期望的类别数
• 数据类型

1import jax.numpy as jnp
2import flaxmodels as fm
3
4num_classes = 2
5dtype = jnp.float32
6model = fm.ResNet50(output='log_softmax', pretrained=None, num_classes=num_classes, dtype=dtype)

ResNet 模块有两种形式：原始 ResNet 模块和预激活 ResNet 模块。两者各有优缺点。下图展示了两种模块的结构：

计算指标

我们定义用于训练过程中评估模型的指标。首先从损失函数开始。

1import jax
2
3def cross_entropy_loss(*, logits, labels):
4    labels_onehot = jax.nn.one_hot(labels, num_classes=num_classes)
5    return optax.softmax_cross_entropy(logits=logits, labels=labels_onehot).mean()

接下来，我们定义一个计算并返回损失和准确率的函数。

1def compute_metrics(*, logits, labels):
2    loss = cross_entropy_loss(logits=logits, labels=labels)
3    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
4    metrics = {
5        'loss': loss,
6        'accuracy': accuracy,
7    }
8    return metrics

创建 Flax 模型训练状态

Flax 提供了训练状态（TrainState）用于存储训练信息，且可自定义扩展。本例中，我们需要添加批归一化统计量（batch_stats），因为 ResNet 模型会计算 batch_stats。

1import flax
2from flax.training import train_state
3
4class TrainState(train_state.TrainState):
5    batch_stats: flax.core.FrozenDict

为了创建训练状态，我们需要模型参数和批统计量。可以通过将 train 设为 False 来初始化模型获取这些值。

1key = jax.random.PRNGKey(0)
2variables = model.init(key, jnp.ones([1, size_image, size_image, 3]), train=False)

TrainState 的 create 方法需要以下参数：

• 模型的 apply_fn（应用函数）
• 模型参数 variables['params']
• 优化器（通常使用 Optax 定义）
• 批统计量 variables['batch_stats']

我们对该函数应用 pmap，以创建可在多个设备（如多 GPU 或 TPU）上运行的分布式版本。pmap 会将函数编译为多设备执行。

1import functools
2import optax
3
4@functools.partial(jax.pmap)
5def create_train_state(rng):
6    """创建初始的 `TrainState`。"""
7    return TrainState.create(apply_fn=model.apply, params=variables['params'], tx=optax.adam(0.01,0.9), batch_stats=variables['batch_stats'])

应用模型函数

接下来，我们定义一个并行的模型训练函数。传入 axis_name 参数，以便聚合所有设备上的指标。该函数：

• 计算损失
• 使用 jax.lax.pmean() 对所有设备的概率取平均，得到预测结果

在应用模型时，还需传入批统计量和用于 Dropout 的随机数。由于这是训练函数，train 参数设为 True。在计算梯度时也包含 batch_stats。update_model 函数则应用计算出的梯度并更新模型参数。

1@functools.partial(jax.pmap, axis_name='ensemble')
2def apply_model(state, images, labels):
3    def loss_fn(params, batch_stats):
4        logits, batch_stats = model.apply({'params': params,'batch_stats': batch_stats}, images, train=True, rngs={'dropout': jax.random.PRNGKey(0)}, mutable=['batch_stats'])
5        one_hot = jax.nn.one_hot(labels, num_classes)
6        loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()
7        return loss, (logits, batch_stats)
8
9    (loss, (logits, batch_stats)), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params, state.batch_stats)
10    probs = jax.lax.pmean(jax.nn.softmax(logits), axis_name='ensemble')
11
12    accuracy = jnp.mean(jnp.argmax(probs, -1) == labels)
13    return grads, loss, accuracy
14
15@jax.pmap
16def update_model(state, grads):
17    return state.apply_gradients(grads=grads)

Flax 中的 TensorBoard

下一步是训练 ResNet 模型。但可能希望使用 TensorBoard 跟踪训练过程。为此，我们需要配置 TensorBoard，并使用 PyTorch 的 SummaryWriter 将指标写入。

1!rm -rf ./flax_logs/
2from torch.utils.tensorboard import SummaryWriter
3
4logdir = "flax_logs"
5writer = SummaryWriter(logdir)

训练 Flax ResNet 模型

我们将在整个训练集上训练 ResNet 模型，并在部分测试集上进行评估（也可在整个测试集上评估）。先将测试集复制到所有可用设备。

1(test_images, test_labels) = next(iter(validation_loader))
2test_images = test_images / 255.0
3test_images = np.array(jax_utils.replicate(test_images))
4test_labels = np.array(jax_utils.replicate(test_labels))

我们创建一些列表来保存训练和评估指标。

1epoch_loss = []
2epoch_accuracy = []
3testing_accuracy = []
4testing_loss = []

接下来，我们定义 ResNet 模型的训练函数。该函数执行以下操作：

• 遍历训练数据集并进行归一化
• 将数据复制到所有可用设备
• 应用模型并计算指标
• 使用 jax_utils.unreplicate 获取各设备的指标
• 将指标追加到列表中
• 计算每轮次的平均损失和准确率
• 在测试集上应用模型并获取指标
• 将测试指标追加到列表
• 将训练和评估指标写入 TensorBoard
• 打印训练和评估指标

1from tqdm.notebook import tqdm
2import math 
3from flax import jax_utils
4
5def train_model(state, dataloader, num_epochs):
6    for epoch in range(num_epochs):
7        for cnt, (images, labels) in tqdm(enumerate(dataloader), total=(math.ceil(len(train_set)/batch_size))):
8            images = images / 255.0
9            images = jax_utils.replicate(images)
10            labels = jax_utils.replicate(labels)
11            grads, loss, accuracy = apply_model(state, images, labels)
12            state = update_model(state, grads)
13
14        epoch_loss.append(jax_utils.unreplicate(loss))
15        epoch_accuracy.append(jax_utils.unreplicate(accuracy))
16        train_loss = np.mean(epoch_loss)
17        train_accuracy = np.mean(epoch_accuracy)
18
19        _, test_loss, test_accuracy = jax_utils.unreplicate(apply_model(state, test_images, test_labels))
20        testing_accuracy.append(test_accuracy)
21        testing_loss.append(test_loss)
22
23        writer.add_scalar('Loss/train', np.array(train_loss), epoch)
24        writer.add_scalar('Loss/test', np.array(test_loss), epoch)
25
26        writer.add_scalar('Accuracy/train', np.array(train_accuracy), epoch)
27        writer.add_scalar('Accuracy/test', np.array(test_accuracy), epoch)
28
29        print(f"Epoch: {epoch + 1}, train loss: {train_loss:.4f}, train accuracy: {train_accuracy * 100:.4f}, test loss: {test_loss:.4f}, test accuracy: {test_accuracy* 100:.4f}", flush=True)
30    return state, epoch_loss, epoch_accuracy, testing_accuracy, testing_loss

我们生成与设备数量相等的随机数来创建训练状态。

1rng = jax.random.PRNGKey(0)
2rng, init_rng = jax.random.split(rng)
3
4state = create_train_state(jax.random.split(init_rng, jax.device_count()))

我们通过传入训练数据和训练轮数来训练 ResNet 模型。

1import time
2
3start = time.time()
4num_epochs = 30
5state, epoch_loss, epoch_accuracy, testing_accuracy, testing_loss = train_model(state, train_loader, num_epochs)
6print("Total time: ", time.time() - start, "seconds")

使用 TensorBoard 评估模型

我们将启动 TensorBoard 查看记录的标量数据。

1%load_ext tensorboard
2%tensorboard --logdir={logdir}

可视化 Flax 模型性能

存储在列表中的指标可以使用 Matplotlib 绘制图表。

1import matplotlib.pyplot as plt
2%matplotlib inline
3
4plt.plot(epoch_accuracy, label="Training")
5plt.plot(testing_accuracy, label="Test")
6plt.xlabel("Epoch")
7plt.ylabel("Accuracy")
8plt.legend()
9plt.show()
10
11plt.plot(epoch_loss, label="Training")
12plt.plot(testing_loss, label="Test")
13plt.xlabel("Epoch")
14plt.ylabel("Loss")
15plt.legend()
16plt.show()

保存 Flax ResNet 模型

要保存训练好的 Flax ResNet 模型，我们使用 save_checkpoint 函数。该函数需要以下参数：

• 保存模型的目录
• 要保存的模型对象（target）
• 训练步数（step）
• 模型文件前缀
• overwrite 参数，决定是否覆盖已有模型

1!pip install tensorstore
2from flax.training import checkpoints
3
4ckpt_dir = 'model_checkpoint/'
5checkpoints.save_checkpoint(ckpt_dir=ckpt_dir,
6                            target=state,
7                            step=100,
8                            prefix='flax_model',
9                            overwrite=True
10                           )

加载 Flax ResNet 模型

已保存的 Flax ResNet 模型也可以被加载以进行预测。使用 restore_checkpoint 函数加载模型，该函数需要：

• 目标状态（target）
• 包含已保存模型的目录
• 模型文件前缀

1loaded_model = checkpoints.restore_checkpoint(ckpt_dir=ckpt_dir,
2                                              target=state,
3                                              prefix='flax_model'
4                                              )

注意： 由于平台限制，ResNet 模型无法在我们的平台上完成训练。但你可以在本地设备上运行该代码。

现在，让我们来看完整的代码：