从Flax Linen迁移到NNX：深度学习框架的演进与对比

富茉钰Ida

于 2025-06-06 09:14:27 发布

阅读量336

点赞数 4

本文链接：https://blog.youkuaiyun.com/gitblog_01091/article/details/148466208

版权

从Flax Linen迁移到NNX：深度学习框架的演进与对比

flax Flax is a neural network library for JAX that is designed for flexibility. 项目地址: https://gitcode.com/gh_mirrors/fl/flax

引言

在深度学习框架的发展历程中，Google的Flax项目一直以其简洁的API设计和与JAX的无缝集成而著称。随着Flax生态系统的演进，NNX模块作为新一代的核心组件被引入，带来了更直观的编程模型和更强大的功能。本文将深入探讨从Flax Linen迁移到NNX的关键差异和实践指南。

核心概念对比

模块(Module)定义差异

在神经网络层的定义上，Linen和NNX都使用Module作为基本构建块，但存在三个根本性差异：

状态管理方式：
- Linen采用**无状态(stateless)**设计，变量通过Module.init()调用返回并单独管理
- NNX采用**有状态(stateful)**设计，变量作为Python对象的属性直接存储
初始化时机：
- Linen是**惰性(lazy)**初始化，只有在看到输入数据时才创建变量
- NNX是**急切(eager)**初始化，实例化时立即创建变量
API设计：
- Linen使用@nn.compact装饰器在单一方法中定义模型
- NNX在__init__中初始化参数，在__call__中定义计算逻辑

# Linen风格
class Block(nn.Module):
    features: int
    
    @nn.compact
    def __call__(self, x, training: bool):
        x = nn.Dense(self.features)(x)
        x = nn.Dropout(0.5)(x)
        return jax.nn.relu(x)

# NNX风格
class Block(nnx.Module):
    def __init__(self, in_features, out_features, rngs: nnx.Rngs):
        self.linear = nnx.Linear(in_features, out_features, rngs=rngs)
        self.dropout = nnx.Dropout(0.5, rngs=rngs)
    
    def __call__(self, x):
        x = self.linear(x)
        x = self.dropout(x)
        return jax.nn.relu(x)

变量创建与管理

变量初始化方式也有显著不同：

Linen：通过init方法显式初始化，返回参数字典
NNX：实例化时自动初始化，变量存储在模块属性中

# Linen初始化
model = Model(256, 10)
variables = model.init(jax.random.key(0), sample_input)

# NNX初始化
model = Model(784, 256, 10, rngs=nnx.Rngs(0))
# 参数已自动初始化
print(model.linear.bias.value.shape)  # 直接访问

训练流程对比

训练步骤实现

训练步骤的实现展示了两种框架的哲学差异：

# Linen训练步骤
@jax.jit
def train_step(params, inputs, labels):
    def loss_fn(params):
        logits = model.apply({'params': params}, inputs, training=True)
        return cross_entropy_loss(logits, labels)
    
    grads = jax.grad(loss_fn)(params)
    return update_params(params, grads)

# NNX训练步骤
model.train()  # 设置训练模式

@nnx.jit
def train_step(model, inputs, labels):
    def loss_fn(model):
        logits = model(inputs)
        return cross_entropy_loss(logits, labels)
    
    grads = nnx.grad(loss_fn)(model)
    update_model(model, grads)  # 就地更新

关键区别：

NNX使用nnx.jit而非jax.jit，支持有状态模块
NNX的nnx.grad返回模块状态而非原始梯度
NNX模型更新是就地(in-place)操作
训练模式通过model.train()/model.eval()控制

变量类型系统

Linen使用集合(collections)组织变量，而NNX使用变量类型：

| Linen集合 | NNX变量类型 | 示例层 | |---------------|---------------------|----------------| | params | nnx.Param | Dense/Linear | | batch_stats | nnx.BatchStat | BatchNorm | | intermediates | nnx.Intermediates | 中间值捕获 |

自定义变量示例：

class Counter(nnx.Variable): pass

class Block(nnx.Module):
    def __init__(self, ...):
        self.count = Counter(jnp.array(0))
    
    def __call__(self, x):
        self.count += 1  # 直接操作
        return x

高级模式对比

多方法模块

实现包含多个方法(如encode/decode)的模块：

# Linen实现
class AutoEncoder(nn.Module):
    def setup(self):
        self.encoder = nn.Dense(256)
        self.decoder = nn.Dense(784)
    
    def encode(self, x):
        return self.encoder(x)
    
    def decode(self, x):
        return self.decoder(x)

# 调用方式
z = model.apply(variables, x, method="encode")

# NNX实现
class AutoEncoder(nnx.Module):
    def __init__(self, in_dim, embed_dim, out_dim, rngs):
        self.encoder = nnx.Linear(in_dim, embed_dim, rngs=rngs)
        self.decoder = nnx.Linear(embed_dim, out_dim, rngs=rngs)
    
    def encode(self, x):
        return self.encoder(x)

# 调用方式
z = model.encode(x)  # 更直观

变换(Transformations)处理

两种框架都提供了对JAX变换的封装，但NNX的变换更自然地处理有状态模块：

# Linen的scan变换
scanned = nn.scan(
    nn.Dense,
    variable_axes={'params': 0},
    split_rngs={'params': True}
)

# NNX的scan变换
class ScannedLinear(nnx.Module):
    def __init__(self, dim, n_layers, rngs):
        self.layers = [
            nnx.Linear(dim, dim, rngs=rngs) 
            for _ in range(n_layers)
        ]
    
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x