Google Flax项目从flax.optim迁移到Optax的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01114/article/details/148466184

Google Flax项目从flax.optim迁移到Optax的完整指南

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

前言

在深度学习框架中，优化器(Optimizer)是训练神经网络模型的核心组件之一。Google Flax项目在2021年提出了用Optax替代原有的flax.optim模块的改进方案，并在v0.6.0版本中完全移除了flax.optim。本文将详细介绍如何将基于flax.optim的代码迁移到Optax，并解释Optax的设计理念和优势。

Optax简介

Optax是一个专注于梯度处理和优化的JAX库，它采用函数式编程范式，通过组合简单的梯度变换来构建复杂的优化器。与flax.optim相比，Optax具有以下优势：

模块化设计：通过组合基本变换构建复杂优化器
更清晰的参数管理：不保存参数副本，参数与优化器状态分离
更好的可组合性：支持灵活地组合各种梯度变换

基础迁移示例

让我们从一个最基本的SGD优化器迁移开始：

# 原flax.optim代码
optimizer_def = flax.optim.Momentum(learning_rate, momentum)
optimizer = optimizer_def.create(variables['params'])

# 迁移后的Optax代码
tx = optax.sgd(learning_rate, momentum)
params = variables['params']
opt_state = tx.init(params)

关键变化：

优化器定义更简洁
参数和优化器状态分离管理
需要显式初始化优化器状态

训练步骤的差异

训练循环的实现也有显著变化：

# flax.optim风格
@jax.jit
def train_step(optimizer, batch):
    grads = jax.grad(loss)(optimizer.target, batch)
    return optimizer.apply_gradient(grads)

# Optax风格
@jax.jit
def train_step(params, opt_state, batch):
    grads = jax.grad(loss)(params, batch)
    updates, opt_state = tx.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state

主要区别在于：

Optax需要显式管理参数和优化器状态
更新过程分为计算更新和应用更新两步
更符合函数式编程范式

理解Optax的模块化设计

Optax的核心思想是将优化过程分解为一系列可组合的梯度变换。例如，SGD优化器实际上是两个基本变换的组合：

tx = optax.chain(
    optax.trace(decay=momentum),  # 动量项
    optax.scale(-learning_rate),  # 学习率缩放
)

这种设计使得我们可以灵活地添加或修改优化器的各个组件。

常用功能迁移指南

1. 权重衰减

在flax.optim中，权重衰减通常作为优化器的参数：

optimizer_def = flax.optim.Adam(learning_rate, weight_decay=weight_decay)

在Optax中，权重衰减是一个独立的梯度变换：

tx = optax.chain(
    optax.scale_by_adam(),
    optax.add_decayed_weights(weight_decay),
    optax.scale(-learning_rate),
)

2. 梯度裁剪

梯度裁剪是训练稳定性的重要技术：

# flax.optim需要手动实现
grads_flat, _ = jax.tree_util.tree_flatten(grads)
global_l2 = jnp.sqrt(sum([jnp.vdot(p, p) for p in grads_flat]))
g_factor = jnp.minimum(1.0, grad_clip_norm / global_l2)
grads = jax.tree_util.tree_map(lambda g: g * g_factor, grads)

# Optax提供内置支持
tx = optax.chain(
    optax.clip_by_global_norm(grad_clip_norm),
    # 其他变换...
)

3. 学习率调度

学习率调度在训练中很常见：

# flax.optim方式
optimizer.apply_gradient(grads, learning_rate=schedule(step))

# Optax方式
tx = optax.chain(
    optax.scale_by_schedule(lambda step: -schedule(step)),
    # 其他变换...
)

高级用法：参数分组优化

有时我们需要对模型的不同部分使用不同的优化策略：

# 创建参数掩码
kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)
biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)

all_false = jax.tree_util.tree_map(lambda _: False, params)
kernels_mask = kernels.update(lambda _: True, all_false)
biases_mask = biases.update(lambda _: True, all_false)

# 创建分组优化器
tx = optax.chain(
    optax.trace(decay=momentum),
    optax.masked(optax.scale(-learning_rate), kernels_mask),
    optax.masked(optax.scale(-learning_rate * 0.1), biases_mask),
)