Flax 与 TensorFlow

Flax 是基于 JAX 的神经网络库。TensorFlow 是一个拥有庞大生态系统和丰富工具资源的深度学习库。Flax 和 TensorFlow 在某些方面相似（例如，两者都可以在 XLA 上运行），但在其他方面存在差异。

下面我们来看一下 Flax 与 TensorFlow 之间的主要区别。

随机数生成

在 TensorFlow 中，我们可以设置全局或函数级别的随机种子。生成随机数非常直接。

1import tensorflow as tf
2
3tf.random.set_seed(6853)
4print(tf.random.uniform([1]))  
5print(tf.random.uniform([1]))  
6print(tf.random.uniform([1]))  

但在 Flax 中情况不同。Flax 建立在 JAX 之上，而 JAX 要求使用纯函数，即没有副作用的函数。为了实现这一点，JAX 引入了无状态的伪随机数生成器（PRNG）。例如，在 NumPy 中调用随机数生成函数每次都会产生不同的结果。

1import numpy as np
2
3print(np.random.random())
4print(np.random.random())
5print(np.random.random())

而在 JAX 和 Flax 中，相同的输入应始终产生相同的结果。因此，我们需要从一个随机状态中生成随机数，并且该状态不应被重复使用。可以通过分裂该状态来获得多个独立的伪随机数序列。

1import jax
2
3key = jax.random.PRNGKey(0)
4key1, key2, key3 = jax.random.split(key, num=3)
5
6print(key1)
7print(key2)
8print(key3)

模型定义

TensorFlow 通过 Keras API 简化了模型定义过程。我们可以使用 Keras 定义 Sequential（顺序）模型或 Functional（函数式）模型。Keras 提供了多种层类型，便于设计 CNN、LSTM 等各类网络结构。

1import tensorflow as tf
2
3class MyModel(tf.keras.Model):
4
5    def __init__(self):
6        super().__init__()
7        self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
8        self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax)
9
10    def call(self, inputs):
11        x = self.dense1(inputs)
12        return self.dense2(x)
13
14model = MyModel()

在 Flax 中，网络可以通过 setup 方法或 compact 方式定义。setup 方法更显式，而 compact 方法则更紧凑简洁。setup 方法的设计风格类似于 PyTorch。例如，以下是一个使用 setup 方法定义的网络。子模块名称由我们赋值的属性决定，此处为 dense1 和 dense2，这与 PyTorch 的逻辑一致。

1from flax import linen as nn
2
3class MLP(nn.Module):
4    def setup(self):
5        self.dense1 = nn.Dense(32)
6        self.dense2 = nn.Dense(32)
7
8    def __call__(self, x):
9        x = self.dense1(x)
10        x = nn.relu(x)
11        x = self.dense2(x)
12        return x

以下是采用 compact 方式定义的相同网络。这种方式更为简洁，减少了代码冗余。

1from flax import linen as nn
2
3class MLP(nn.Module):
4
5    @nn.compact
6    def __call__(self, x):
7        x = nn.Dense(32, name="dense1")(x)
8        x = nn.relu(x)
9        x = nn.Dense(32, name="dense2")(x)
10        return x

激活函数

TensorFlow 中的 tf.keras.activations 模块提供了构建神经网络所需的大多数激活函数。

1import tensorflow as tf
2import numpy as np
3
4layer = tf.keras.layers.Activation('relu')
5output = layer([-3.0, -1.0, 0.0, 2.0])
6print(list(output.numpy()))

在 Flax 中，激活函数通过 linen 模块提供。

1import flax.linen as nn
2import numpy as np
3
4print(nn.relu(np.array([-3.0, -1.0, 0.0, 2.0])))

优化器

TensorFlow 提供了常用的优化器函数，位于 tf.keras.optimizers 模块中。

1optimizer = tf.keras.optimizers.Adam(
2    learning_rate=0.001,
3    beta_1=0.9,
4    beta_2=0.999,
5    epsilon=1e-07,
6    amsgrad=False,
7    weight_decay=None,
8    clipnorm=None,
9    clipvalue=None,
10    global_clipnorm=None,
11    use_ema=False,
12    ema_momentum=0.99,
13    ema_overwrite_frequency=None,
14    jit_compile=True,
15    name='Adam',
16    **kwargs
17)

然而，Flax 本身并不包含优化器函数。Flax 中使用的优化器由另一个名为 Optax 的库提供。

1optimizers = optax.adamw(learning_rate=0.001, 
2                         b1=0.9, 
3                         b2=0.999, 
4                         eps=1e-08, 
5                         eps_root=0.0, 
6                         mu_dtype=None)

指标（Metrics）

在 TensorFlow 中，可通过 tf.keras.metrics 模块获取各种评估指标。

1import tensorflow as tf
2
3m = tf.keras.metrics.Accuracy()
4m.update_state([[1], [2], [3], [4]], [[0], [2], [3], [4]])
5print(m.result().numpy())

截至目前，Flax 尚未提供内置的 metrics 模块。我们需要自行定义评估指标函数，或借助第三方库实现。

1import optax
2import jax.numpy as jnp
3
4def compute_metrics(logits, labels):
5    loss = jnp.mean(optax.softmax_cross_entropy(logits, jax.nn.one_hot(labels, num_classes=2)))
6    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
7    metrics = {
8        'loss': loss,
9        'accuracy': accuracy
10    }
11    return metrics

梯度计算

在 Flax 中，使用 jax.grad 函数来计算梯度。它还支持返回辅助数据，例如可以同时返回损失值和梯度。

1@jax.jit
2def sum_logistic(x):
3    return jnp.sum(1.0 / (1.0 + jnp.exp(-x))),(x + 1)
4
5x_small = jnp.arange(6.)
6derivative_fn = jax.grad(sum_logistic, has_aux=True)
7print(derivative_fn(x_small))

还可以使用 jax.vjp() 和 jax.jvp() 实现更高级的自动微分功能。

在 TensorFlow 中，梯度是通过 tf.GradientTape 计算的。

1import tensorflow as tf
2
3def grad(model, inputs, targets):
4    with tf.GradientTape() as tape:
5        loss_value = loss(model, inputs, targets, training=True)
6    return loss_value, tape.gradient(loss_value, model.trainable_variables)

除非我们在 TensorFlow 中自定义训练循环，否则通常不需要手动定义梯度函数——在训练网络时系统会自动完成这一过程。

数据集加载

TensorFlow 提供了用于加载数据的工具函数。Flax 本身不提供任何数据加载器，我们必须使用其他库（如 TensorFlow）的数据加载功能。只要数据格式为 JAX NumPy 数组或普通数组，并具有正确的形状，就可以传递给 Flax 网络使用。

1import tensorflow as tf
2
3max_features = 5000  # 最大词汇表大小
4batch_size = 128
5max_len = 25 # 输出序列填充长度
6
7vectorize_layer = tf.keras.layers.TextVectorization(standardize='lower_and_strip_punctuation',max_tokens=max_features,output_mode='int',output_sequence_length=max_len)
8vectorize_layer.adapt(X_train, batch_size=None)
9
10X_train_padded =  vectorize_layer(X_train)
11X_test_padded =  vectorize_layer(X_test)
12
13training_data = tf.data.Dataset.from_tensor_slices((X_train_padded, y_train))
14validation_data = tf.data.Dataset.from_tensor_slices((X_test_padded, y_test))
15
16training_data = training_data.batch(batch_size)
17validation_data = validation_data.batch(batch_size)