从Haiku迁移到Flax NNX的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00170/article/details/148466206

从Haiku迁移到Flax NNX的完整指南

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

前言

在深度学习框架领域，Google的Flax NNX和Haiku都是基于JAX构建的优秀神经网络库。本文将深入探讨如何将现有的Haiku代码迁移到Flax NNX框架，帮助开发者理解两者之间的核心差异，并提供实用的迁移策略。

核心概念对比

模块(Module)系统

Flax NNX和Haiku都使用Module作为构建神经网络的基本单元，但它们在设计理念上有显著不同：

状态管理方式：
- Haiku采用**无状态(stateless)**设计，变量通过init()调用返回并由外部管理
- Flax NNX采用**有状态(stateful)**设计，变量作为Module对象的属性直接存储
初始化时机：
- Haiku是**惰性(lazy)**初始化，只有在实际看到输入数据时才创建变量
- Flax NNX是**急切(eager)**初始化，实例化时立即创建变量

代码结构示例

下面是一个简单的网络层定义对比：

# Haiku实现
class Block(hk.Module):
    def __init__(self, features: int):
        super().__init__()
        self.features = features
    
    def __call__(self, x, training: bool):
        x = hk.Linear(self.features)(x)
        x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)
        return jax.nn.relu(x)

# Flax NNX实现
class Block(nnx.Module):
    def __init__(self, in_features: int, out_features: int, 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)

变量创建与管理

Haiku的变量处理

Haiku需要将模型包装在转换函数中：

def forward(x, training: bool):
    return Model(256, 10)(x, training)

model = hk.transform(forward)
params = model.init(jax.random.key(0), sample_x, training=False)

Flax NNX的变量处理

Flax NNX在实例化时自动初始化变量：

model = Model(784, 256, 10, rngs=nnx.Rngs(0))
# 参数已初始化
print(model.linear.bias.value.shape)  # (10,)

训练流程对比

训练步骤实现

# Haiku训练步骤
@jax.jit
def train_step(key, params, inputs, labels):
    def loss_fn(params):
        logits = model.apply(params, key, inputs, training=True)
        return optax.softmax_cross_entropy(logits, labels).mean()
    
    grads = jax.grad(loss_fn)(params)
    return jax.tree_map(lambda p, g: p - 0.1 * g, params, grads)

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

@nnx.jit
def train_step(model, inputs, labels):
    def loss_fn(model):
        logits = model(inputs)
        return optax.softmax_cross_entropy(logits, labels).mean()
    
    grads = nnx.grad(loss_fn)(model)
    _, params, rest = nnx.split(model, nnx.Param, ...)
    params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
    nnx.update(model, nnx.merge_state(params, rest))

关键差异点

训练模式切换：
- Haiku需要显式传递training参数
- Flax NNX通过model.train()和model.eval()控制
梯度计算：
- Haiku使用jax.grad返回原始梯度字典
- Flax NNX使用nnx.grad返回状态字典
参数更新：
- Haiku需要返回更新后的参数树
- Flax NNX直接就地更新模型状态

处理非参数状态

对于包含BatchNorm等层的模型：

# Haiku实现
class Block(hk.Module):
    def __call__(self, x, training: bool):
        x = hk.BatchNorm(create_scale=True, create_offset=True)(x, is_training=training)
        return x

model = hk.transform_with_state(forward)
params, batch_stats = model.init(...)

# Flax NNX实现
class Block(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.batchnorm = nnx.BatchNorm(momentum=0.99, rngs=rngs)
    
    def __call__(self, x):
        return self.batchnorm(x)

model = Block(rngs=nnx.Rngs(0))
model.batchnorm.mean  # 直接访问批统计量

多方法模型实现

实现像自编码器这样的多方法模型：

# Haiku实现
class AutoEncoder(hk.Module):
    def encode(self, x): ...
    def decode(self, x): ...
    def __call__(self, x): ...

model = hk.multi_transform(lambda: (lambda x: module(x), (module.encode, module.decode))

# Flax NNX实现
class AutoEncoder(nnx.Module):
    def encode(self, x): ...
    def decode(self, x): ...

model = AutoEncoder(rngs=nnx.Rngs(0))
z = model.encode(inputs)  # 直接调用

转换(Transformations)处理

实现RNN单元的例子：

# Haiku实现
class RNN(hk.Module):
    def __call__(self, x):
        return hk.scan(self.cell)(self.initial_state(batch_size), x)

# Flax NNX实现
class RNN(nnx.Module):
    def __call__(self, x):
        return nnx.scan(self.cell, in_axes=0, out_axes=0)(self.initial_state(batch_size), x)