Flax 是一个用于 JAX 的神经网络库。JAX 是一个 Python 库，为机器学习研究提供高性能计算能力。JAX 提供了与 NumPy 类似的 API，易于上手，同时还包含多项有助于提升机器学习研究效率的功能，包括：

• 自动微分：JAX 支持通过 jacrev、grad 和 hessian 等函数对数值函数进行前向和反向自动微分。
• 向量化：JAX 通过 vmap 函数支持自动向量化，同时可通过 pmap 函数轻松实现大规模数据处理的并行化。
• 即时编译（JIT）：JAX 使用 XLA 实现代码在 GPU 和 TPU 上的即时（JIT）编译与执行。

在本节课中，我们将使用 JAX 和 Flax 构建一个简单的卷积神经网络。

加载数据集

我们将使用之前课程中用过的猫狗分类数据集。首先解压包含该数据集的压缩文件。

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

Flax 本身不提供数据加载工具，我们可以使用 PyTorch 或 TensorFlow 的数据加载器。这里我们选择使用 PyTorch 来加载数据。第一步是定义数据集类。

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

接下来，创建一个包含类别信息的 pandas DataFrame。

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

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

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

现在可以定义训练集和验证集，并结合 PyTorch 的 DataLoader 使用。同时定义一个图像缩放的变换操作。

1import torch
2from torch.utils.data import DataLoader
3from torchvision import transforms
4
5size_image = 64
6batch_size = 32
7
8transform = transforms.Compose([transforms.Resize((size_image, size_image)), np.array])
9dataset = CatsDogsDataset("train", "train_csv.csv", transform=transform)
10
11train_set, validation_set = torch.utils.data.random_split(dataset, [0.7,0.3])
12
13train_loader = DataLoader(dataset=train_set, collate_fn=custom_collate_fn, shuffle=True, batch_size=batch_size)
14validation_loader = DataLoader(dataset=validation_set, collate_fn=custom_collate_fn, shuffle=False, batch_size=batch_size)

可视化数据集

我们从猫狗数据集中各选取一张图像进行可视化展示。

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")

使用 Flax 定义卷积神经网络

在 Flax 中，通过继承 Module 类并使用 Linen API 来创建网络。所有 Flax 模块都是 Python 数据类，默认具有 __init__ 方法，因此应重写 setup() 方法来初始化网络。不过，我们可以使用 @nn.compact 装饰器使模型定义更加简洁。

1from flax import linen as nn  
2
3class CNN(nn.Module):
4
5    @nn.compact
6    def __call__(self, x):
7        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
8        x = nn.relu(x)
9        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
10        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
11        x = nn.relu(x)
12        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
13        x = x.reshape((x.shape[0], -1))  
14        x = nn.Dense(features=256)(x)
15        x = nn.relu(x)
16        x = nn.Dense(features=2)(x)
17        x = nn.log_softmax(x)
18        return x

定义损失函数

可以使用 optax 包计算损失。在传入 softmax 交叉熵函数前，需将整数标签进行 one-hot 编码。由于我们处理的是两类分类任务，num_classes 设为 2。

1import optax
2
3def cross_entropy_loss(*, logits, labels):
4    labels_onehot = jax.nn.one_hot(labels, num_classes=2)
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

创建训练状态

训练状态（TrainState）保存模型变量，如参数和优化器状态，并在每次迭代中由优化器更新。我们可以继承 flax.training.train_state 来追踪更多状态信息（例如 Dropout 或批归一化层的状态）。但对于当前简单模型，使用默认类即可满足需求。

1from flax.training import train_state
2
3def create_train_state(rng, learning_rate, momentum):
4    """创建初始 TrainState"""
5    cnn = CNN()
6    params = cnn.init(rng, jnp.ones([1, size_image, size_image, 3]))['params']
7    tx = optax.sgd(learning_rate, momentum)
8    return train_state.TrainState.create(apply_fn=cnn.apply, params=params, tx=tx)

定义训练步骤

在此函数中，我们使用 apply 方法对输入图像进行模型推理，得到 logits 后计算损失。接着使用 value_and_grad 同时计算损失函数及其梯度，并用梯度更新模型参数。最后调用前面定义的 compute_metrics 函数计算损失和准确率。该函数使用 @jit 装饰器进行 JIT 编译，以加快执行速度。

1def compute_loss(params, images, labels):
2    logits = CNN().apply({'params': params}, images)
3    loss = cross_entropy_loss(logits=logits, labels=labels)
4    return loss, logits
5
6@jax.jit
7def train_step(state, images, labels):
8    """执行单步训练"""
9    (_, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(state.params, images, labels)
10    state = state.apply_gradients(grads=grads) 
11    metrics = compute_metrics(logits=logits, labels=labels) 
12    return state, metrics

定义评估步骤

评估函数使用 apply 方法在验证集上评估模型性能。

1@jax.jit
2def eval_step(state, images, labels):
3    logits = CNN().apply({'params': state.params}, images)
4    return compute_metrics(logits=logits, labels=labels)

训练函数

该函数应用上述定义的训练步骤，遍历数据加载器中的每个批次并进行优化。使用 jax.device_get 获取指标并在 CPU 上计算均值。

1def train_one_epoch(state, dataloader):
2    """在训练集上训练一个 epoch"""
3    batch_metrics = []
4    for cnt, (images, labels) in enumerate(dataloader):
5        images = images / 255.0
6        state, metrics = train_step(state, images, labels)
7        batch_metrics.append(metrics)
8
9    batch_metrics_np = jax.device_get(batch_metrics)  
10    epoch_metrics_np = {
11        k: np.mean([metrics[k] for metrics in batch_metrics_np])
12        for k in batch_metrics_np[0]
13    }
14
15    return state, epoch_metrics_np

评估模型

评估函数执行评估步骤并返回验证集上的指标。

1def evaluate_model(state, test_imgs, test_lbls):
2    """在验证集上评估模型"""
3    metrics = eval_step(state, test_imgs, test_lbls)
4    metrics = jax.device_get(metrics) 
5    metrics = jax.tree_map(lambda x: x.item(), metrics)  
6    return metrics

训练并评估模型

在训练前需要初始化训练状态。状态初始化函数需要一个伪随机数生成器（PRNG）密钥。我们使用 PRNGKey 生成密钥，并通过 split 得到另一个用于参数初始化的密钥。

将该密钥连同学习率和动量参数一起传入 create_train_state 函数。随后即可使用 train_one_epoch 进行训练，并使用 evaluate_model 进行评估。

1import jax 
2
3learning_rate = 0.001
4momentum = 0.9
5seed = 0 
6num_epochs = 10
7
8(test_images, test_labels) = next(iter(validation_loader))
9test_images = test_images / 255.0
10
11state = create_train_state(jax.random.PRNGKey(seed), learning_rate, momentum)
12
13training_loss = []
14training_accuracy = []
15testing_loss = []
16testing_accuracy = []
17
18for epoch in range(1, num_epochs + 1):
19    state, train_metrics = train_one_epoch(state, train_loader)
20    training_loss.append(train_metrics['loss'])
21    training_accuracy.append(train_metrics['accuracy'])
22    print(f"Train epoch: {epoch}, loss: {train_metrics['loss']}, accuracy: {train_metrics['accuracy'] * 100}")
23
24    test_metrics = evaluate_model(state, test_images, test_labels)
25    testing_loss.append(test_metrics['loss'])
26    testing_accuracy.append(test_metrics['accuracy'])
27    print(f"Test epoch: {epoch}, loss: {test_metrics['loss']}, accuracy: {test_metrics['accuracy'] * 100}")

模型性能

训练过程中会打印训练和验证的指标。我们还可以利用这些结果绘制训练和验证的准确率与损失曲线。

1import matplotlib.pyplot as plt
2
3plt.figure(0, figsize=(8,6))
4plt.plot(training_accuracy, label="Training")
5plt.plot(testing_accuracy, label="Test")
6plt.xlabel("Epoch")
7plt.ylabel("Accuracy")
8plt.legend()
9plt.savefig('output/accuracy.png')
10
11plt.figure(1, figsize=(8,6))
12plt.plot(training_loss, label="Training")
13plt.plot(testing_loss, label="Test")
14plt.xlabel("Epoch")
15plt.ylabel("Loss")
16plt.legend()
17plt.savefig('output/loss.png')