一、nn.Module 详解：一切模型的基石
说实话，我刚接触 Flax 时，第一反应是：这玩意儿跟 PyTorch 的 nn.Module 也太像了吧？
但用久了你会发现，Flax 的 Module 有它自己的脾气。它不是一个简单的「层容器」，而是一个 函数式 + 面向对象 的混合体。
1.1 核心概念：setup 与 __call__
每个 nn.Module 必须实现两个方法：
- setup()：定义子模块和参数的地方。只调用一次。
- __call__()：前向传播逻辑。每次调用都会执行。
我见过不少新手把参数定义写在 __call__ 里，结果每次前向都重新创建参数，模型直接炸了。嗯，这是个坑。
import flax.linen as nn
import jax.numpy as jnp
class SimpleMLP(nn.Module):
hidden_dim: int = 128
output_dim: int = 10
def setup(self):
# 这里定义子模块，只执行一次
self.dense1 = nn.Dense(self.hidden_dim)
self.dense2 = nn.Dense(self.output_dim)
def __call__(self, x):
x = self.dense1(x)
x = nn.relu(x)
x = self.dense2(x)
return x
model = SimpleMLP()
x = jnp.ones((32, 784))
params = model.init(jax.random.PRNGKey(0), x)
y = model.apply(params, x)
print(y.shape) # (32, 10)1.2 参数的生命周期
Flax 的参数是「显式管理」的。你调用 model.init() 时，它会返回一个参数字典。调用 model.apply() 时，你得把参数传进去。
这跟 PyTorch 那种「参数挂在模型上」的方式完全不同。我第一次迁移项目时，就因为忘了传参数，debug 了一整个下午。
二、参数初始化策略：别让模型输在起跑线上
初始化做不好，模型训练三天三夜都不收敛。我见过太多这样的案例了。
2.1 Flax 的默认初始化
Flax 的 Dense、Conv 等层都有默认初始化器。Dense 默认用 lecun_normal()，偏置默认用 zeros。
但说实话，默认值不一定适合你的场景。比如用 ReLU 激活时，我习惯用 kaiming_normal。
2.2 自定义初始化
from flax.linen import initializers
class MyDense(nn.Module):
features: int
kernel_init: callable = initializers.kaiming_normal()
bias_init: callable = initializers.zeros
def setup(self):
self.kernel = self.param(
'kernel',
self.kernel_init,
(self.features,), # 注意：这里只传 shape
jnp.float32
)
self.bias = self.param(
'bias',
self.bias_init,
(self.features,),
jnp.float32
)
def __call__(self, x):
return x @ self.kernel + self.bias关键点： self.param() 的第一个参数是名字，第二个是初始化函数，第三个是 shape 或 shape 元组。初始化函数签名是 init(key, shape, dtype)。
2.3 初始化策略选择表
| 激活函数 | 推荐初始化 | 适用场景 |
|---|---|---|
| ReLU / LeakyReLU | kaiming_normal | CNN、MLP |
| Tanh / Sigmoid | xavier_normal | RNN、全连接 |
| GELU / Swish | lecun_normal | Transformer |
| 无激活（回归） | orthogonal | 输出层 |
你想想看，如果初始化值太大，梯度爆炸；太小，梯度消失。这个平衡点，就是初始化策略要解决的问题。
三、模块组合与复用：搭积木的艺术
模块组合是 Flax 的核心能力。说白了，就是把小模块拼成大模型。
3.1 嵌套模块
class ResidualBlock(nn.Module):
features: int
def setup(self):
self.conv1 = nn.Conv(self.features, kernel_size=(3, 3))
self.bn1 = nn.BatchNorm()
self.conv2 = nn.Conv(self.features, kernel_size=(3, 3))
self.bn2 = nn.BatchNorm()
def __call__(self, x, training=True):
residual = x
x = self.conv1(x)
x = self.bn1(x, use_running_average=not training)
x = nn.relu(x)
x = self.conv2(x)
x = self.bn2(x, use_running_average=not training)
x = nn.relu(x + residual)
return x
class ResNet18(nn.Module):
num_classes: int = 10
def setup(self):
self.blocks = [
ResidualBlock(64),
ResidualBlock(64),
ResidualBlock(128),
ResidualBlock(128),
]
self.fc = nn.Dense(self.num_classes)
def __call__(self, x, training=True):
for block in self.blocks:
x = block(x, training=training)
x = x.mean(axis=(1, 2)) # global avg pooling
x = self.fc(x)
return x3.2 参数共享
有时候你想让两个子模块共享同一份参数。Flax 怎么做？
class SharedWeightModel(nn.Module):
def setup(self):
self.shared_dense = nn.Dense(64)
def __call__(self, x1, x2):
# 同一个 dense 被调用两次，参数共享
h1 = self.shared_dense(x1)
h2 = self.shared_dense(x2)
return h1 + h2嗯，就是这么简单。只要在 setup 里定义一次，多次调用就是共享参数。
四、自定义模块：从零开始造轮子
有时候内置模块不够用，你得自己写。比如实现一个自定义激活函数，或者一个特殊的归一化层。
4.1 自定义无参数模块
class Swish(nn.Module):
def __call__(self, x):
return x * nn.sigmoid(x)
# 使用
model = nn.Sequential([
nn.Dense(128),
Swish(), # 直接当层用
nn.Dense(10)
])4.2 自定义带参数模块
class LearnableSwish(nn.Module):
def setup(self):
# 可学习的缩放参数
self.beta = self.param('beta', initializers.ones, (1,), jnp.float32)
def __call__(self, x):
return x * nn.sigmoid(self.beta * x)
# 使用
model = LearnableSwish()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 10)))
print(params['params']['beta']) # 可以看到 beta 参数setup() 或 @nn.compact 装饰的方法里完成，不能依赖运行时条件。
4.3 自定义模块的最佳实践
- 命名规范： 模块名用大驼峰，参数名用小写下划线
- 类型注解： 给所有超参数加上类型注解，方便调试
- 文档字符串： 写清楚输入输出的 shape 和 dtype
- 测试先行： 写个简单的 init 和 apply 测试，确保模块能跑通
知识体系总览
下面这张图，是我自己梳理的本章知识结构。你看一眼，心里就有谱了。
这张图把本章四个核心知识点串起来了。你写代码时，脑子里可以过一遍这个结构：先定义模块，再初始化参数，然后组合复用，最后按需自定义。
总结一句话： nn.Module 是 Flax 的骨架，参数初始化是血液，模块组合是肌肉，自定义模块是灵活的手指。四者配合，才能构建出生产级的模型。