Google Flax框架中的批归一化(BatchNorm)使用指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01065/article/details/148466197

Google Flax框架中的批归一化(BatchNorm)使用指南

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

批归一化概述

批归一化(Batch Normalization)是一种深度学习中常用的正则化技术，由Ioffe和Szegedy在2015年提出。在Google Flax框架中，批归一化通过flax.linen.BatchNorm模块实现，它能够显著加速神经网络训练过程并提高模型收敛性。

批归一化的核心思想是对每一层的输入进行标准化处理，使其均值为0，方差为1。这种操作有助于解决深度神经网络训练中的"内部协变量偏移"问题，使得各层输入分布更加稳定，从而允许使用更大的学习率。

Flax中BatchNorm的基本用法

在Flax中，BatchNorm是一个特殊的模块，它在训练和推理阶段有不同的行为。与PyTorch或TensorFlow不同，Flax通过显式的use_running_average参数来控制这种行为，而不是通过模块的eval()模式或training标志。

定义包含BatchNorm的模型

定义一个包含BatchNorm的多层感知机(MLP)示例：

import flax.linen as nn

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x, train: bool):
        x = nn.Dense(features=4)(x)
        x = nn.BatchNorm(use_running_average=not train)(x)
        x = nn.relu(x)
        x = nn.Dense(features=1)(x)
        return x

关键点：

模型接收一个train布尔参数
BatchNorm的use_running_average设置为not train
训练时(train=True)，使用当前批次的统计量
推理时(train=False)，使用运行平均值

模型初始化与变量结构

初始化包含BatchNorm的模型时，会生成额外的状态变量：

import jax
import jax.numpy as jnp

mlp = MLP()
x = jnp.ones((1, 3))
variables = mlp.init(jax.random.key(0), x, train=False)

此时variables包含两个集合：

params: 包含所有可训练参数(权重和偏置)
batch_stats: 包含BatchNorm的运行统计量(均值和方差)

具体变量结构如下：

{
    'batch_stats': {
        'BatchNorm_0': {
            'mean': (4,),  # 特征维度为4的运行均值
            'var': (4,),    # 特征维度为4的运行方差
        },
    },
    'params': {
        'BatchNorm_0': {
            'bias': (4,),   # 特征维度为4的偏置参数
            'scale': (4,),  # 特征维度为4的缩放参数
        },
        # 其他层参数...
    }
}

训练与推理的实现差异

训练阶段实现

训练时需要特别注意：

必须传入batch_stats集合
需要标记batch_stats为可变(mutable)
需要接收并更新返回的batch_stats

y, updates = mlp.apply(
    {'params': params, 'batch_stats': batch_stats},
    x,
    train=True,
    mutable=['batch_stats']
)
batch_stats = updates['batch_stats']

自定义TrainState

为了在训练循环中管理batch_stats，需要扩展基础的TrainState：

from flax.training import train_state
from typing import Any

class TrainState(train_state.TrainState):
    batch_stats: Any  # 存储BatchNorm的运行统计量

state = TrainState.create(
    apply_fn=mlp.apply,
    params=params,
    batch_stats=batch_stats,
    tx=optax.adam(1e-3)
)

训练步骤实现

训练步骤需要处理batch_stats的更新：

@jax.jit
def train_step(state: TrainState, batch):
    def loss_fn(params):
        logits, updates = state.apply_fn(
            {'params': params, 'batch_stats': state.batch_stats},
            x=batch['image'],
            train=True,
            mutable=['batch_stats']
        )
        loss = compute_loss(logits, batch['label'])
        return loss, (logits, updates)
    
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, (logits, updates)), grads = grad_fn(state.params)
    
    # 更新参数和batch_stats
    state = state.apply_gradients(grads=grads)
    state = state.replace(batch_stats=updates['batch_stats'])
    
    return state, compute_metrics(loss, logits, batch['label'])

评估步骤实现

评估阶段不需要更新batch_stats：

@jax.jit
def eval_step(state: TrainState, batch):
    logits = state.apply_fn(
        {'params': state.params, 'batch_stats': state.batch_stats},
        x=batch['image'],
        train=False
    )
    loss = compute_loss(logits, batch['label'])
    return state, compute_metrics(loss, logits, batch['label'])