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

本文链接：https://blog.youkuaiyun.com/gitblog_00277/article/details/148466183

从Haiku迁移到Flax的完整指南

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

前言

在深度学习框架领域，JAX生态系统提供了多个高级神经网络库，其中Haiku和Flax都是基于JAX构建的流行选择。本文将详细介绍如何将现有的Haiku模型迁移到Flax框架，并深入分析两者在设计和实现上的关键差异。

基础模块对比

模块定义方式

Haiku和Flax在模块定义上有显著差异：

类定义结构：
- Haiku使用传统的Python类继承hk.Module，需要在__init__中显式调用父类构造器
- Flax使用数据类(dataclass)风格，通过类属性定义配置参数
名称参数处理：
- Haiku要求显式声明name参数并传递给父类
- Flax自动处理name参数，无需显式定义
方法装饰器：
- Flax需要使用@nn.compact装饰器来启用内联子模块定义
- Haiku默认支持内联子模块定义

示例代码对比

# Haiku实现
class Block(hk.Module):
    def __init__(self, features: int, name=None):
        super().__init__(name=name)
        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)
        x = jax.nn.relu(x)
        return x

# Flax实现
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, deterministic=not training)(x)
        x = jax.nn.relu(x)
        return x

模型初始化与参数结构

模型构造方式

Haiku：使用hk.transform包装模型函数，返回包含init和apply方法的对象
Flax：直接实例化模块类

参数初始化

# Haiku参数初始化
params = model.init(random.key(0), sample_x, training=False)

# Flax参数初始化
variables = model.init(random.key(0), sample_x, training=False)
params = variables["params"]

参数结构差异

Haiku和Flax的参数组织结构有本质区别：

Haiku参数结构：
- 两级层次结构
- 使用"/"分隔的模块路径作为键
- 参数名称作为子键
Flax参数结构：
- 多级嵌套结构
- 使用模块类名和实例编号作为键
- 更自然的层次关系

训练过程对比

基本训练步骤

# Haiku训练步骤
def train_step(key, params, inputs, labels):
    def loss_fn(params):
        logits = model.apply(params, key, inputs, training=True)
        return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
    # ...梯度计算和参数更新...

# Flax训练步骤
def train_step(key, params, inputs, labels):
    def loss_fn(params):
        logits = model.apply(
            {'params': params},
            inputs, training=True,
            rngs={'dropout': key}
        )
        return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
    # ...梯度计算和参数更新...

关键区别：

Flax需要将参数包装在字典中
Flax的随机数生成器也需要特定格式
Haiku直接传递参数和随机键

状态处理机制

BatchNorm实现对比

# Haiku实现
x = hk.BatchNorm(create_scale=True, create_offset=True, decay_rate=0.99)(x, is_training=training)

# Flax实现
x = nn.BatchNorm(momentum=0.99)(x, use_running_average=not training)

状态初始化

# Haiku状态初始化
params, state = model.init(random.key(0), sample_x, training=True)

# Flax状态初始化
variables = model.init(random.key(0), sample_x, training=False)
params, batch_stats = variables["params"], variables["batch_stats"]

状态训练循环

# Haiku状态训练
def train_step(params, state, inputs, labels):
    def loss_fn(params):
        logits, new_state = model.apply(params, state, None, inputs, training=True)
        # ...损失计算...
    # ...梯度计算和状态更新...

# Flax状态训练
def train_step(params, batch_stats, inputs, labels):
    def loss_fn(params):
        logits, updates = model.apply(
            {'params': params, 'batch_stats': batch_stats},
            inputs, training=True,
            mutable='batch_stats'
        )
        # ...损失计算...
    # ...梯度计算和状态更新...

多方法模块实现

自动编码器示例

# Haiku实现
class AutoEncoder(hk.Module):
    def __init__(self, embed_dim: int, output_dim: int, name=None):
        super().__init__(name=name)
        self.encoder = hk.Linear(embed_dim, name="encoder")
        self.decoder = hk.Linear(output_dim, name="decoder")
    # ...方法实现...

# Flax实现
class AutoEncoder(nn.Module):
    embed_dim: int
    output_dim: int
    
    def setup(self):
        self.encoder = nn.Dense(self.embed_dim)
        self.decoder = nn.Dense(self.output_dim)
    # ...方法实现...

方法调用方式

# Haiku方法调用
encode, decode = model.apply
z = encode(params, None, x=jax.numpy.ones((1, 784)))

# Flax方法调用
z = model.apply({"params": params}, x=jax.numpy.ones((1, 784)), method="encode")

高阶变换应用

RNN单元实现

# Haiku实现
class RNNCell(hk.Module):
    def __init__(self, hidden_size: int, name=None):
        super().__init__(name=name)
        self.hidden_size = hidden_size
    # ...方法实现...

# Flax实现
class RNNCell(nn.Module):
    hidden_size: int
    
    @nn.compact
    def __call__(self, carry, x):
        # ...实现逻辑...